Implement caching for SMIR kompilation (#777)

* Upon SMIR kompilation, save the digest of the `smir.json` file into a
file `smir-digest.json` in the target directory. This is then used on
subsequent calls to check if re-kompilation is necessary.
* Add option `--target-dir` for `kmir run`.

Author: tothtamas28

kmir/src/kmir/__main__.py

@@ -14,7 +14,7 @@
 from pyk.proof.show import APRProofShow
 from pyk.proof.tui import APRProofViewer
 
-from .build import HASKELL_DEF_DIR, LLVM_DEF_DIR, LLVM_LIB_DIR
+from .build import HASKELL_DEF_DIR, LLVM_LIB_DIR
 from .cargo import CargoProject
 from .kmir import KMIR, KMIRAPRNodePrinter
 from .linker import link
@@ -42,8 +42,6 @@
 
 
 def _kmir_run(opts: RunOpts) -> None:
-    kmir = KMIR(HASKELL_DEF_DIR) if opts.haskell_backend else KMIR(LLVM_DEF_DIR)
-
     if opts.file:
         smir_info = SMIRInfo.from_file(Path(opts.file))
     else:
@@ -54,11 +52,17 @@ def _kmir_run(opts: RunOpts) -> None:
         # target = opts.bin if opts.bin else cargo.default_target
         smir_info = cargo.smir_for_project(clean=False)
 
-    with tempfile.TemporaryDirectory() as work_dir:
-        kmir = KMIR.from_kompiled_kore(smir_info, symbolic=opts.haskell_backend, target_dir=work_dir)
+    def run(target_dir: Path):
+        kmir = KMIR.from_kompiled_kore(smir_info, symbolic=opts.haskell_backend, target_dir=target_dir)
         result = kmir.run_smir(smir_info, start_symbol=opts.start_symbol, depth=opts.depth)
         print(kmir.kore_to_pretty(result))
 
+    if opts.target_dir:
+        run(target_dir=opts.target_dir)
+    else:
+        with tempfile.TemporaryDirectory() as target_dir:
+            run(target_dir=Path(target_dir))
+
 
 def _kmir_prove_rs(opts: ProveRSOpts) -> None:
     proof = KMIR.prove_rs(opts)
@@ -185,6 +189,7 @@ def _arg_parser() -> ArgumentParser:
     run_target_selection = run_parser.add_mutually_exclusive_group()
     run_target_selection.add_argument('--bin', metavar='TARGET', help='Target to run')
     run_target_selection.add_argument('--file', metavar='SMIR', help='SMIR json file to execute')
+    run_parser.add_argument('--target-dir', type=Path, metavar='TARGET_DIR', help='SMIR kompilation target directory')
     run_parser.add_argument('--depth', type=int, metavar='DEPTH', help='Depth to execute')
     run_parser.add_argument(
         '--start-symbol', type=str, metavar='SYMBOL', default='main', help='Symbol name to begin execution from'
@@ -313,6 +318,7 @@ def _parse_args(ns: Namespace) -> KMirOpts:
             return RunOpts(
                 bin=ns.bin,
                 file=ns.file,
+                target_dir=ns.target_dir,
                 depth=ns.depth,
                 start_symbol=ns.start_symbol,
                 haskell_backend=ns.haskell_backend,

kmir/src/kmir/kmir.py

@@ -56,7 +56,7 @@ def __init__(
 
     @staticmethod
     def from_kompiled_kore(
-        smir_info: SMIRInfo, target_dir: str, bug_report: Path | None = None, symbolic: bool = True
+        smir_info: SMIRInfo, target_dir: Path, bug_report: Path | None = None, symbolic: bool = True
     ) -> KMIR:
         from .kompile import kompile_smir
 
@@ -147,12 +147,9 @@ def prove_rs(opts: ProveRSOpts) -> APRProof:
                 _LOGGER.info(f'Reading proof from disc: {opts.proof_dir}, {label}')
                 apr_proof = APRProof.read_proof_data(opts.proof_dir, label)
 
-                # TODO avoid compilation, use compilation output from the proof directory
-                # kmir = KMIR(opts.proof_dir / label / haskell, opts.proof_dir / label / llvm-library) if they exist
-                # or else implement this in the `from_kompiled_kore` constructor
                 smir_info = SMIRInfo.from_file(target_path / 'smir.json')
                 kmir = KMIR.from_kompiled_kore(
-                    smir_info, symbolic=True, bug_report=opts.bug_report, target_dir=str(target_path)
+                    smir_info, symbolic=True, bug_report=opts.bug_report, target_dir=target_path
                 )
             else:
                 _LOGGER.info(f'Constructing initial proof: {label}')
@@ -176,7 +173,7 @@ def prove_rs(opts: ProveRSOpts) -> APRProof:
                     _LOGGER.debug(f'Missing-body function symbols (first 5): {missing_body_syms[:5]}')
 
                 kmir = KMIR.from_kompiled_kore(
-                    smir_info, symbolic=True, bug_report=opts.bug_report, target_dir=str(target_path)
+                    smir_info, symbolic=True, bug_report=opts.bug_report, target_dir=target_path
                 )
 
                 apr_proof = kmir.apr_proof_from_smir(

kmir/src/kmir/kompile.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+import json
 import logging
 import shutil
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from pyk.kast.inner import KApply, KSort
@@ -15,6 +15,7 @@
 from .kmir import KMIR
 
 if TYPE_CHECKING:
+    from pathlib import Path
     from typing import Any, Final
 
     from pyk.kast import KInner
@@ -54,45 +55,74 @@ def create_kmir(self, *, bug_report_file: Path | None = None) -> KMIR:
         )
 
 
+@dataclass
+class KompileDigest:
+    digest: str
+    symbolic: bool
+
+    @staticmethod
+    def load(target_dir: Path) -> KompileDigest:
+        digest_file = KompileDigest._digest_file(target_dir)
+
+        if not digest_file.exists():
+            raise ValueError(f'Digest file does not exist: {digest_file}')
+
+        data = json.loads(digest_file.read_text())
+        return KompileDigest(
+            digest=data['digest'],
+            symbolic=data['symbolic'],
+        )
+
+    def write(self, target_dir: Path) -> None:
+        self._digest_file(target_dir).write_text(
+            json.dumps(
+                {
+                    'digest': self.digest,
+                    'symbolic': self.symbolic,
+                },
+            ),
+        )
+
+    @staticmethod
+    def _digest_file(target_dir: Path) -> Path:
+        return target_dir / 'smir-digest.json'
+
+
 def kompile_smir(
     smir_info: SMIRInfo,
-    target_dir: str,
+    target_dir: Path,
     bug_report: Path | None = None,
     symbolic: bool = True,
 ) -> KompiledSMIR:
-    kmir = KMIR(HASKELL_DEF_DIR)
-
-    def _insert_rules_and_write(input_file: Path, rules: list[str], output_file: Path) -> None:
-        with open(input_file, 'r') as f:
-            lines = f.readlines()
+    kompile_digest: KompileDigest | None = None
+    try:
+        kompile_digest = KompileDigest.load(target_dir)
+    except Exception:
+        pass
 
-        # last line must start with 'endmodule'
-        last_line = lines[-1]
-        assert last_line.startswith('endmodule')
-        new_lines = lines[:-1]
+    target_hs_path = target_dir / 'haskell'
+    target_llvm_lib_path = target_dir / 'llvm-library'
+    target_llvm_path = target_dir / 'llvm'
 
-        # Insert rules before the endmodule line
-        new_lines.append(f'\n// Generated from file {input_file}\n\n')
-        new_lines.extend(rules)
-        new_lines.append('\n' + last_line)
+    if kompile_digest is not None and kompile_digest.digest == smir_info.digest and kompile_digest.symbolic == symbolic:
+        _LOGGER.info(f'Kompiled SMIR up-to-date, no kompilation necessary: {target_dir}')
+        if symbolic:
+            return KompiledSymbolic(haskell_dir=target_hs_path, llvm_lib_dir=target_llvm_lib_path)
+        else:
+            return KompiledConcrete(llvm_dir=target_llvm_path)
 
-        # Write to output file
-        with open(output_file, 'w') as f:
-            f.writelines(new_lines)
+    _LOGGER.info(f'Kompiling SMIR program: {target_dir}')
 
-    target_path = Path(target_dir)
-    # TODO if target dir exists and contains files, check file dates (definition files and interpreter)
-    # to decide whether or not to recompile. For now we always recompile.
-    target_path.mkdir(parents=True, exist_ok=True)
+    kompile_digest = KompileDigest(digest=smir_info.digest, symbolic=symbolic)
+    target_dir.mkdir(parents=True, exist_ok=True)
 
+    kmir = KMIR(HASKELL_DEF_DIR)
     rules = _make_kore_rules(kmir, smir_info)
     _LOGGER.info(f'Generated {len(rules)} function equations to add to `definition.kore')
 
     if symbolic:
         # Create output directories
-        target_llvm_path = target_path / 'llvm-library'
-        target_llvmdt_path = target_llvm_path / 'dt'
-        target_hs_path = target_path / 'haskell'
+        target_llvmdt_path = target_llvm_lib_path / 'dt'
 
         _LOGGER.info(f'Creating directories {target_llvmdt_path} and {target_hs_path}')
         target_llvmdt_path.mkdir(parents=True, exist_ok=True)
@@ -101,7 +131,7 @@ def _insert_rules_and_write(input_file: Path, rules: list[str], output_file: Pat
         # Process LLVM definition
         _LOGGER.info('Writing LLVM definition file')
         llvm_def_file = LLVM_LIB_DIR / 'definition.kore'
-        llvm_def_output = target_llvm_path / 'definition.kore'
+        llvm_def_output = target_llvm_lib_path / 'definition.kore'
         _insert_rules_and_write(llvm_def_file, rules, llvm_def_output)
 
         # Run llvm-kompile-matching and llvm-kompile for LLVM
@@ -123,7 +153,7 @@ def _insert_rules_and_write(input_file: Path, rules: list[str], output_file: Pat
                 '-O2',
                 '--',
                 '-o',
-                target_llvm_path / 'interpreter',
+                target_llvm_lib_path / 'interpreter',
             ],
             check=True,
         )
@@ -141,13 +171,11 @@ def _insert_rules_and_write(input_file: Path, rules: list[str], output_file: Pat
                     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)
-        return KompiledSymbolic(
-            haskell_dir=target_hs_path,
-            llvm_lib_dir=target_llvm_path,
-        )
-    else:
 
-        target_llvm_path = target_path / 'llvm'
+        kompile_digest.write(target_dir)
+        return KompiledSymbolic(haskell_dir=target_hs_path, llvm_lib_dir=target_llvm_lib_path)
+
+    else:
         target_llvmdt_path = target_llvm_path / 'dt'
         _LOGGER.info(f'Creating directory {target_llvmdt_path}')
         target_llvmdt_path.mkdir(parents=True, exist_ok=True)
@@ -183,6 +211,8 @@ def _insert_rules_and_write(input_file: Path, rules: list[str], output_file: Pat
         for file in to_copy:
             _LOGGER.info(f'Copying file {file}')
             shutil.copy2(LLVM_DEF_DIR / file, target_llvm_path / file)
+
+        kompile_digest.write(target_dir)
         return KompiledConcrete(llvm_dir=target_llvm_path)
 
 
@@ -294,3 +324,22 @@ def _decode_alloc(smir_info: SMIRInfo, raw_alloc: Any) -> tuple[KInner, KInner]:
 
     alloc_id_term = KApply('allocId', intToken(alloc_id))
     return alloc_id_term, value.to_kast()
+
+
+def _insert_rules_and_write(input_file: Path, rules: list[str], output_file: Path) -> None:
+    with open(input_file, 'r') as f:
+        lines = f.readlines()
+
+    # last line must start with 'endmodule'
+    last_line = lines[-1]
+    assert last_line.startswith('endmodule')
+    new_lines = lines[:-1]
+
+    # Insert rules before the endmodule line
+    new_lines.append(f'\n// Generated from file {input_file}\n\n')
+    new_lines.extend(rules)
+    new_lines.append('\n' + last_line)
+
+    # Write to output file
+    with open(output_file, 'w') as f:
+        f.writelines(new_lines)

kmir/src/kmir/options.py

@@ -22,6 +22,7 @@ class KMirOpts(LoggingOptions): ...
 class RunOpts(KMirOpts):
     bin: str | None
     file: str | None
+    target_dir: Path | None
     depth: int
     start_symbol: str
     haskell_backend: bool

kmir/src/kmir/smir.py

@@ -40,6 +40,13 @@ def dump(self, smir_json_file: Path) -> None:
     def name(self) -> str:
         return self._smir['name']
 
+    @cached_property
+    def digest(self) -> str:
+        import hashlib
+
+        hash_object = hashlib.sha256(str(self._smir).encode('UTF-8'))
+        return hash_object.hexdigest()
+
     @cached_property
     def allocs(self) -> dict[AllocId, AllocInfo]:
         return {

kmir/src/tests/integration/test_decode_value.py

@@ -31,7 +31,7 @@ def definition_dir():  # -> Path:
     try:
         with open(lock_file, 'x') as _:
             # generate and compile an LLVM interpreter with the type-table
-            _ = KMIR.from_kompiled_kore(TEST_SMIR, target_dir=str(target_dir), symbolic=False)
+            _ = KMIR.from_kompiled_kore(TEST_SMIR, target_dir=target_dir, symbolic=False)
         lock_file.unlink()
     except FileExistsError:
         # wait loop until interpreter exists, max 1min

kmir/src/tests/integration/test_integration.py

@@ -329,16 +329,14 @@ def test_exec_smir(
     test_case: tuple[str, Path, Path, int],
     symbolic: bool,
     update_expected_output: bool,
+    tmp_path: Path,
 ) -> None:
-
-    (_, input_json, output_kast, depth) = test_case
+    _, input_json, output_kast, depth = test_case
     smir_info = SMIRInfo.from_file(input_json)
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        kmir_backend = KMIR.from_kompiled_kore(smir_info, target_dir=temp_dir, symbolic=symbolic)
-        result = kmir_backend.run_smir(smir_info, depth=depth)
-        result_pretty = kmir_backend.kore_to_pretty(result).rstrip()
-        assert_or_update_show_output(result_pretty, output_kast, update=update_expected_output)
+    kmir_backend = KMIR.from_kompiled_kore(smir_info, target_dir=tmp_path, symbolic=symbolic)
+    result = kmir_backend.run_smir(smir_info, depth=depth)
+    result_pretty = kmir_backend.kore_to_pretty(result).rstrip()
+    assert_or_update_show_output(result_pretty, output_kast, update=update_expected_output)
 
 
 @pytest.mark.parametrize(