From 5883587d2868a419ab96e860ff6db16f828149da Mon Sep 17 00:00:00 2001
From: LittleStar <59785146+LinZhihao-723@users.noreply.github.com>
Date: Sat, 16 Mar 2024 23:04:00 +0800
Subject: [PATCH] Add parallel compilation support

---
 setup.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 57edf165..772c762a 100644
--- a/setup.py
+++ b/setup.py
@@ -1,8 +1,10 @@
+import distutils.ccompiler
+import multiprocessing.pool
 import os
 import platform
 import sys
 from setuptools import setup, Extension
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 ir_native: Extension = Extension(
     name="clp_ffi_py.ir.native",
@@ -41,6 +43,33 @@
     define_macros=[("SOURCE_PATH_SIZE", str(len(os.path.abspath("./src/clp/components/core"))))],
 )
 
+def _parallel_compile(
+    self: distutils.ccompiler.CCompiler,
+    sources: List[str],
+    output_dir: Optional[str] = None,
+    macros: Optional[List[Tuple[str, Optional[str]]]] = None,
+    include_dirs: Optional[List[str]] = None,
+    debug: int = 0,
+    extra_preargs: Optional[List[str]] = None,
+    extra_postargs: Optional[List[str]] = None,
+    depends: Optional[List[str]] = None,
+) -> List[str]:
+    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
+        output_dir, macros, include_dirs, sources, depends, extra_postargs
+    )
+    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
+
+    def _compile_single_file(obj: str) -> None:
+        try:
+            src, ext = build[obj]
+        except KeyError:
+            return
+        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
+    num_cores: int = multiprocessing.cpu_count()
+    multiprocessing.pool.ThreadPool(num_cores).map(_compile_single_file, objects)
+    return objects
+
 if "__main__" == __name__:
     try:
         if "Darwin" == platform.system():
@@ -48,6 +77,7 @@
             if None is target or float(target) < 10.15:
                 os.environ["MACOSX_DEPLOYMENT_TARGET"] = "10.15"
 
+        distutils.ccompiler.CCompiler.compile = _parallel_compile
         project_name: str = "clp_ffi_py"
         description: str = "CLP FFI Python Interface"
         extension_modules: List[Extension] = [ir_native]