From bf0347188769cec3683b181eb4f4632ad67cc578 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20T=C3=B3th?=
 <tothtamas28@users.noreply.github.com>
Date: Tue, 14 May 2024 12:59:42 +0000
Subject: [PATCH] Implement kompilation for `PykBackend.BOOSTER`

---
 pyk/src/pyk/ktool/kompile.py              | 105 +++++++++++++++++++++-
 pyk/src/tests/integration/test_kompile.py |  19 +++-
 2 files changed, 121 insertions(+), 3 deletions(-)

diff --git a/pyk/src/pyk/ktool/kompile.py b/pyk/src/pyk/ktool/kompile.py
index 58d2c40fc62..ff71cda144e 100644
--- a/pyk/src/pyk/ktool/kompile.py
+++ b/pyk/src/pyk/ktool/kompile.py
@@ -2,6 +2,7 @@
 
 __all__ = ['PykBackend', 'kompile']
 
+import concurrent.futures
 import dataclasses
 import logging
 import shlex
@@ -69,7 +70,20 @@ def kompile(
 
     pyk_backend = PykBackend(backend) if backend else None
     if pyk_backend is PykBackend.BOOSTER:
-        raise ValueError('Backend not supported')
+        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
 
@@ -89,6 +103,95 @@ 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
 # -----------
diff --git a/pyk/src/tests/integration/test_kompile.py b/pyk/src/tests/integration/test_kompile.py
index afcdb4e3293..3491531f22d 100644
--- a/pyk/src/tests/integration/test_kompile.py
+++ b/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