fix(kmir): preserve symbolic parallel proving server selection

kmir/src/kmir/_prove.py

@@ -3,15 +3,15 @@
 import logging
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, cast
 
 from pyk.cterm import CTerm
 from pyk.cterm.symbolic import CTermSymbolic
 from pyk.kast.inner import KSequence, KVariable, Subst
 from pyk.kast.manip import abstract_term_safely, split_config_from
 from pyk.kcfg import KCFG
 from pyk.kcfg.explore import KCFGExplore
-from pyk.kore.rpc import BoosterServer, KoreClient
+from pyk.kore.rpc import BoosterServer, KoreClient, KoreServer
 from pyk.proof.proof import parallel_advance_proof
 from pyk.proof.reachability import APRProof, APRProver
 
@@ -24,6 +24,7 @@
     from typing import Final
 
     from pyk.kast.inner import KInner
+    from pyk.kore.rpc import BoosterServerArgs, KoreServerArgs
 
     from .options import ProveRSOpts
 
@@ -143,18 +144,30 @@ def _prove_parallel(
     cut_point_rules: list[str],
 ) -> None:
     assert opts.max_workers
-    assert kmir.llvm_library_dir
+    server_args: dict[str, Any] = {
+        'kompiled_dir': kmir.definition_dir,
+        'module_name': kmir.definition.main_module_name,
+        'bug_report': kmir.bug_report,
+        'haskell_threads': opts.max_workers,
+    }
+    server_ctx: BoosterServer | KoreServer
 
-    with BoosterServer(
-        {
-            'kompiled_dir': kmir.definition_dir,
+    if kmir.llvm_library_dir is not None:
+        booster_server_args = {
+            **server_args,
             'llvm_kompiled_dir': kmir.llvm_library_dir,
-            'module_name': kmir.definition.main_module_name,
-            'bug_report': kmir.bug_report,
             'simplify_each': 30,
-            'haskell_threads': opts.max_workers,
         }
-    ) as server:
+        command = KMIR.kore_rpc_booster_command_from_env()
+        if command is not None:
+            booster_server_args['command'] = command
+            _LOGGER.info(f'Passing HS-only symbols to kore-rpc-booster: {KMIR.hs_only_symbols_from_env()}')
+        server_ctx = BoosterServer(cast('BoosterServerArgs', booster_server_args))
+    else:
+        _LOGGER.info('No proof-specific LLVM library available; using KoreServer for parallel proof execution')
+        server_ctx = KoreServer(cast('KoreServerArgs', server_args))
+
+    with server_ctx as server:
 
         def create_prover() -> APRProver:
             client = KoreClient(

kmir/src/kmir/kmir.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from contextlib import contextmanager
 from functools import cached_property
 from typing import TYPE_CHECKING
@@ -84,15 +85,54 @@ class Symbols:
         END_PROGRAM: Final = KApply('#EndProgram_KMIR-CONTROL-FLOW_KItem')
         THUNK: Final = KLabel('thunk(_)_RT-DATA_Value_Evaluation')
 
+    @staticmethod
+    def hs_only_symbols_from_env() -> list[str]:
+        raw_value = os.environ.get('KMIR_HS_ONLY_SYMBOLS', '')
+        if not raw_value:
+            return []
+        symbols: list[str] = []
+        seen: set[str] = set()
+        for token in raw_value.split(','):
+            symbol = token.strip()
+            if symbol and symbol not in seen:
+                seen.add(symbol)
+                symbols.append(symbol)
+        return symbols
+
+    @staticmethod
+    def kore_rpc_booster_command_from_env() -> list[str] | None:
+        hs_only_symbols = KMIR.hs_only_symbols_from_env()
+        if not hs_only_symbols:
+            return None
+        command = ['kore-rpc-booster']
+        for symbol in hs_only_symbols:
+            command += ['--hs-only-symbol', symbol]
+        return command
+
     @cached_property
     def parser(self) -> Parser:
         return Parser(self.definition)
 
     @contextmanager
     def kcfg_explore(self, label: str | None = None, terminate_on_thunk: bool = False) -> Iterator[KCFGExplore]:
+        kore_rpc_command = KMIR.kore_rpc_booster_command_from_env()
+        if kore_rpc_command is None:
+            with cterm_symbolic(
+                self.definition,
+                self.definition_dir,
+                llvm_definition_dir=self.llvm_library_dir,
+                bug_report=self.bug_report,
+                id=label if self.bug_report is not None else None,  # NB bug report arg.s must be coherent
+                simplify_each=30,
+            ) as cts:
+                yield KCFGExplore(cts, kcfg_semantics=KMIRSemantics(terminate_on_thunk=terminate_on_thunk))
+            return
+
+        _LOGGER.info(f'Passing HS-only symbols to kore-rpc-booster: {KMIR.hs_only_symbols_from_env()}')
         with cterm_symbolic(
             self.definition,
             self.definition_dir,
+            kore_rpc_command=kore_rpc_command,
             llvm_definition_dir=self.llvm_library_dir,
             bug_report=self.bug_report,
             id=label if self.bug_report is not None else None,  # NB bug report arg.s must be coherent

kmir/src/kmir/kompile.py

@@ -219,6 +219,8 @@ def kompile_smir(
     llvm_target = llvm_target or 'mir-semantics.llvm'
     llvm_lib_target = llvm_lib_target or 'mir-semantics.llvm-library'
     haskell_target = haskell_target or 'mir-semantics.haskell'
+    base_llvm_lib_path = kdist.which(llvm_lib_target)
+    use_hs_only_symbols = symbolic and bool(KMIR.hs_only_symbols_from_env())
 
     expected_digest = KompileDigest(
         digest=smir_info.digest,
@@ -236,6 +238,8 @@ def kompile_smir(
     if kompile_digest == expected_digest:
         _LOGGER.info(f'Kompiled SMIR up-to-date, no kompilation necessary: {target_dir}')
         if symbolic:
+            if use_hs_only_symbols:
+                return KompiledSymbolic(haskell_dir=target_hs_path, llvm_lib_dir=base_llvm_lib_path)
             return KompiledSymbolic(haskell_dir=target_hs_path, llvm_lib_dir=target_llvm_lib_path)
         else:
             return KompiledConcrete(llvm_dir=target_llvm_path)
@@ -261,6 +265,26 @@ def kompile_smir(
     all_rules = smir_rules + extra_rules
 
     if symbolic:
+        if use_hs_only_symbols:
+            _LOGGER.info(f'Creating directory {target_hs_path}')
+            target_hs_path.mkdir(parents=True, exist_ok=True)
+
+            _LOGGER.info('Writing Haskell definition file')
+            hs_def_file = haskell_def_dir / 'definition.kore'
+            _insert_rules_and_write(hs_def_file, all_rules, target_hs_path / 'definition.kore')
+
+            _LOGGER.info('Copying other artefacts into HS output directory')
+            for file_path in haskell_def_dir.iterdir():
+                if file_path.name != 'definition.kore' and file_path.name != 'haskellDefinition.bin':
+                    if file_path.is_file():
+                        shutil.copy2(file_path, target_hs_path / file_path.name)
+                    elif file_path.is_dir():
+                        shutil.copytree(file_path, target_hs_path / file_path.name, dirs_exist_ok=True)
+
+            _LOGGER.info('Using HS-only symbol mode: skip proof-specific llvm-kompile and reuse base llvm-library')
+            kompile_digest.write(target_dir)
+            return KompiledSymbolic(haskell_dir=target_hs_path, llvm_lib_dir=base_llvm_lib_path)
+
         # Create output directories
         target_llvmdt_path = target_llvm_lib_path / 'dt'
 

kmir/src/tests/unit/test_kmir_hs_only_symbols.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import os
+
+from kmir.kmir import KMIR
+
+
+def test_hs_only_symbols_from_env_parses_and_deduplicates(monkeypatch) -> None:
+    monkeypatch.setenv('KMIR_HS_ONLY_SYMBOLS', ' lookupTy , #getBlocks,lookupTy ,, ')
+    assert KMIR.hs_only_symbols_from_env() == ['lookupTy', '#getBlocks']
+
+
+def test_kore_rpc_booster_command_from_env(monkeypatch) -> None:
+    monkeypatch.setenv('KMIR_HS_ONLY_SYMBOLS', 'lookupFunction,#metadataSize')
+    assert KMIR.kore_rpc_booster_command_from_env() == [
+        'kore-rpc-booster',
+        '--hs-only-symbol',
+        'lookupFunction',
+        '--hs-only-symbol',
+        '#metadataSize',
+    ]
+
+
+def test_kore_rpc_booster_command_from_env_empty(monkeypatch) -> None:
+    monkeypatch.setitem(os.environ, 'KMIR_HS_ONLY_SYMBOLS', '')
+    assert KMIR.kore_rpc_booster_command_from_env() is None

kmir/src/tests/unit/test_prove.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, cast
+
+from kmir import _prove
+
+if TYPE_CHECKING:
+    from pyk.proof.reachability import APRProof
+
+    from kmir.kmir import KMIR
+    from kmir.options import ProveRSOpts
+
+
+def test_prove_parallel_uses_kore_server_without_llvm_library(monkeypatch) -> None:
+    opts = cast(
+        'ProveRSOpts',
+        SimpleNamespace(
+            max_workers=2,
+            terminate_on_thunk=False,
+            max_depth=None,
+            max_iterations=None,
+            fail_fast=False,
+            maintenance_rate=1,
+        ),
+    )
+    kmir = cast(
+        'KMIR',
+        SimpleNamespace(
+            llvm_library_dir=None,
+            definition_dir=Path('/tmp/definition'),
+            definition=SimpleNamespace(main_module_name='KMIR-TEST'),
+            bug_report=None,
+        ),
+    )
+    proof = cast('APRProof', object())
+    calls: list[tuple[str, dict[str, object]]] = []
+
+    class FakeKoreServer:
+        def __init__(self, args: dict[str, object]) -> None:
+            calls.append(('kore', args))
+            self.port = 31337
+
+        def __enter__(self) -> FakeKoreServer:
+            return self
+
+        def __exit__(self, _exc_type, _exc_val, _exc_tb) -> None:
+            return None
+
+    def fail_booster(*_args, **_kwargs) -> None:
+        raise AssertionError('BoosterServer should not be used without llvm_library_dir')
+
+    monkeypatch.setattr(_prove, 'KoreServer', FakeKoreServer)
+    monkeypatch.setattr(_prove, 'BoosterServer', fail_booster)
+    monkeypatch.setattr(_prove, 'parallel_advance_proof', lambda *_args, **_kwargs: None)
+
+    _prove._prove_parallel(kmir, proof, opts=opts, label='demo', cut_point_rules=['rule'])
+
+    assert calls == [
+        (
+            'kore',
+            {
+                'kompiled_dir': Path('/tmp/definition'),
+                'module_name': 'KMIR-TEST',
+                'bug_report': None,
+                'haskell_threads': 2,
+            },
+        )
+    ]