Add support for parallel Booster kompilation

pyk/src/pyk/ktool/kompile.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-__all__ = ['KompileBackend', 'kompile']
+__all__ = ['PykBackend', 'kompile']
 
+import concurrent.futures
 import dataclasses
 import logging
 import shlex
@@ -31,6 +32,14 @@ def __init__(self, kompile_command: str):
         super().__init__(f'Kompile command not found: {str}')
 
 
+class PykBackend(Enum):
+    LLVM = 'llvm'
+    HASKELL = 'haskell'
+    KORE = 'kore'
+    MAUDE = 'maude'
+    BOOSTER = 'booster'
+
+
 class Warnings(Enum):
     ALL = 'all'
     NORMAL = 'normal'
@@ -40,20 +49,44 @@ class Warnings(Enum):
 def kompile(
     main_file: str | Path,
     *,
+    backend: str | PykBackend | None = None,
+    # ---
     command: Iterable[str] = ('kompile',),
     output_dir: str | Path | None = None,
     temp_dir: str | Path | None = None,
     type_inference_mode: str | TypeInferenceMode | None = None,
     warnings: str | Warnings | None = None,
     warnings_to_errors: bool = False,
     no_exc_wrap: bool = False,
+    # ---
     debug: bool = False,
     verbose: bool = False,
     cwd: Path | None = None,
     check: bool = True,
+    # ---
     **kwargs: Any,
 ) -> Path:
     kwargs['main_file'] = main_file
+
+    pyk_backend = PykBackend(backend) if backend else None
+    if pyk_backend is PykBackend.BOOSTER:
+        return _booster_kompile(
+            command=command,
+            output_dir=output_dir,
+            temp_dir=temp_dir,
+            type_inference_mode=type_inference_mode,
+            warnings=warnings,
+            warnings_to_errors=warnings_to_errors,
+            no_exc_wrap=no_exc_wrap,
+            debug=debug,
+            verbose=verbose,
+            cwd=cwd,
+            check=check,
+            kwargs=kwargs,
+        )
+
+    kwargs['backend'] = KompileBackend(pyk_backend.value) if pyk_backend else None
+
     kompiler = Kompile.from_dict(kwargs)
     return kompiler(
         command=command,
@@ -70,6 +103,100 @@ def kompile(
     )
 
 
+def _booster_kompile(
+    command: Iterable[str],
+    output_dir: str | Path | None,
+    temp_dir: str | Path | None,
+    type_inference_mode: str | TypeInferenceMode | None,
+    warnings: str | Warnings | None,
+    warnings_to_errors: bool,
+    no_exc_wrap: bool,
+    # ---
+    debug: bool,
+    verbose: bool,
+    cwd: Path | None,
+    check: bool,
+    # ---
+    kwargs: Mapping[str, Any],
+) -> Path:
+    llvm_kt = kwargs.get('llvm_kompile_type')
+    llvm_kt = LLVMKompileType(llvm_kt) if llvm_kt else None
+    if llvm_kt and llvm_kt is not LLVMKompileType.C:
+        raise ValueError(f'Unsupported argument value for Booster kompilation: llvm_kompile_type: {llvm_kt.value}')
+
+    llvm_args, haskell_args = _group_args(kwargs)
+
+    llvm_args['backend'] = KompileBackend.LLVM
+    llvm_args['llvm_kompile_type'] = LLVMKompileType.C
+    llvm_kompile = LLVMKompile.from_dict(llvm_args)
+
+    haskell_args['backend'] = KompileBackend.HASKELL
+    haskell_kompile = HaskellKompile.from_dict(haskell_args)
+
+    main_file = Path(kwargs['main_file'])
+    output_dir = Path(output_dir) if output_dir else _default_output_dir(main_file)
+    temp_dir = Path(temp_dir) if temp_dir else None
+
+    def kompile_llvm() -> None:
+        llvm_kompile(
+            command=command,
+            output_dir=output_dir / 'llvm-library',
+            temp_dir=temp_dir / 'llvm-library' if temp_dir else None,
+            type_inference_mode=type_inference_mode,
+            warnings=warnings,
+            warnings_to_errors=warnings_to_errors,
+            no_exc_wrap=no_exc_wrap,
+            debug=debug,
+            verbose=verbose,
+            cwd=cwd,
+            check=check,
+        )
+
+    def kompile_haskell() -> None:
+        haskell_kompile(
+            command=command,
+            output_dir=output_dir,
+            temp_dir=temp_dir,
+            type_inference_mode=type_inference_mode,
+            warnings=warnings,
+            warnings_to_errors=warnings_to_errors,
+            no_exc_wrap=no_exc_wrap,
+            debug=debug,
+            verbose=verbose,
+            cwd=cwd,
+            check=check,
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [executor.submit(f) for f in [kompile_llvm, kompile_haskell]]
+        for future in concurrent.futures.as_completed(futures):
+            future.result()
+
+    assert output_dir.is_dir()
+    return output_dir
+
+
+def _group_args(args: Mapping[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
+    llvm_args = {}
+    haskell_args = {}
+
+    for arg, value in args.items():
+        if arg in COMMON_ARGS:
+            llvm_args[arg] = value
+            haskell_args[arg] = value
+        elif arg in KompileBackend.LLVM.args:
+            llvm_args[arg] = value
+        elif arg in KompileBackend.HASKELL.args:
+            haskell_args[arg] = value
+
+    return llvm_args, haskell_args
+
+
+# -----------
+# kompile CLI
+# -----------
+
+
 class KompileBackend(Enum):
     LLVM = 'llvm'
     HASKELL = 'haskell'
@@ -210,7 +337,7 @@ def __call__(
             if bug_report:
                 bug_report.add_file_contents(out, Path('kompile.log'))
 
-        definition_dir = output_dir if output_dir else Path(self.base_args.main_file.stem + '-kompiled')
+        definition_dir = output_dir if output_dir else _default_output_dir(self.base_args.main_file)
         assert definition_dir.is_dir()
 
         return definition_dir
@@ -219,6 +346,10 @@ def __call__(
     def args(self) -> list[str]: ...
 
 
+def _default_output_dir(main_file: Path) -> Path:
+    return Path(main_file.stem + '-kompiled')
+
+
 @final
 @dataclass(frozen=True)
 class HaskellKompile(Kompile):

pyk/src/tests/integration/test_kompile.py

@@ -7,16 +7,17 @@
 from pyk.kast import Atts
 from pyk.kast.inner import KSort
 from pyk.kore.syntax import SortApp
-from pyk.ktool.kompile import KompileBackend, KompileNotFoundError, kompile
+from pyk.ktool.kompile import DefinitionInfo, KompileBackend, KompileNotFoundError, PykBackend, kompile
 from pyk.testing import KompiledTest
 
 from .utils import K_FILES
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from pyk.kast.inner import KLabel
     from pyk.kast.outer import KDefinition
     from pyk.kore.kompiled import KompiledKore
-    from pyk.ktool.kompile import DefinitionInfo
 
 
 class TestHaskellKompile(KompiledTest):
@@ -50,6 +51,20 @@ def test_kompile_not_found(monkeypatch: pytest.MonkeyPatch) -> None:
         kompile(k_file, command=[bad_kompile])
 
 
+def test_booster_kompile(tmp_path: Path) -> None:
+    # Given
+    output_dir = tmp_path / 'kompiled'
+    main_file = tmp_path / 'test.k'
+    main_file.write_text('module TEST endmodule')
+
+    # When
+    kompile(main_file, backend=PykBackend.BOOSTER, output_dir=output_dir)
+
+    # Then
+    assert DefinitionInfo(output_dir).backend == KompileBackend.HASKELL
+    assert DefinitionInfo(output_dir / 'llvm-library').backend == KompileBackend.LLVM
+
+
 class TestKLabel(KompiledTest):
     KOMPILE_DEFINITION = """
         module KLABEL