kmir prove takes multiple start-symbols (#1071)

This PR adds the ability to provide multiple start symbols to `kmir
prove`, and importantly, a single call to `llvm-kompile` for the proofs.

In particular:
- Names can be provided to `--start-symbol(s)` either comma separated or
repeated `kmir prove --start-symbols name1, name2 --start-symbol name3`;
- Function `_prove_multi` will reduce the items to the provided names
and do a single `llvm-kompile` to be shared;
- Each proof is run _sequentially_ and the definition and smir.json is
stored in a directory `filename_stem.kompiled`, and the proof dirs just
store the kcfg (this is only if multiple start symbols are provided);
- `prove_with_kmir` is not removed. It was considered but the way it
interacts with the test runner is nicer so I left it;
- Test is added to `test_cli`;
- Observed to correctly reduce the calls to `llvm-kompile` and has
greatly increased performance for mutliple start symbol proofs;

Author: dkcumming

kmir/src/kmir/__main__.py

@@ -43,6 +43,13 @@
 _LOG_FORMAT: Final = '%(levelname)s %(asctime)s %(name)s - %(message)s'
 
 
+def _flatten_comma_list(values: list[str] | None) -> list[str] | None:
+    """Flatten a list that may contain comma-separated entries, e.g. ['a,b', 'c'] -> ['a', 'b', 'c']."""
+    if values is None:
+        return None
+    return [item for v in values for item in v.split(',') if item.strip()]
+
+
 def _flush_lines(lines: Iterable[str]) -> None:
     """Print lines to stdout one at a time, then release the list if possible."""
     for line in lines:
@@ -82,9 +89,13 @@ def run(target_dir: Path):
 
 
 def _kmir_prove(opts: ProveOpts) -> None:
-    proof = KMIR.prove_program(opts)
-    print(str(proof.summary))
-    if not proof.passed:
+    proofs = KMIR.prove_programs(opts)
+    any_failed = False
+    for proof in proofs:
+        print(str(proof.summary))
+        if not proof.passed:
+            any_failed = True
+    if any_failed:
         sys.exit(1)
 
 
@@ -587,7 +598,14 @@ def _arg_parser() -> ArgumentParser:
     )
     prove_parser.add_argument('--smir', action='store_true', help='Treat the input file as a smir json.')
     prove_parser.add_argument(
-        '--start-symbol', type=str, metavar='SYMBOL', default='main', help='Symbol name to begin execution from'
+        '--start-symbol',
+        '--start-symbols',
+        dest='start_symbols',
+        type=str,
+        metavar='SYMBOL',
+        action='append',
+        default=None,
+        help='Symbol name(s) to prove (repeatable, comma-separated allowed)',
     )
     prove_parser.add_argument(
         '--add-module',
@@ -710,7 +728,7 @@ def _parse_args(ns: Namespace) -> KMirOpts:
                 maintenance_rate=ns.maintenance_rate,
                 save_smir=ns.save_smir,
                 smir=ns.smir,
-                start_symbol=ns.start_symbol,
+                start_symbols=_flatten_comma_list(ns.start_symbols),
                 break_on_calls=ns.break_on_calls,
                 break_on_function_calls=ns.break_on_function_calls,
                 break_on_intrinsic_calls=ns.break_on_intrinsic_calls,

kmir/src/kmir/_prove.py

@@ -31,78 +31,96 @@
 _LOGGER: Final = logging.getLogger(__name__)
 
 
-def prove(opts: ProveOpts) -> APRProof:
-    """Run a proof creating a new KMIR instance."""
+def prove(opts: ProveOpts) -> list[APRProof]:
+    """Prove one or more start symbols from the same file, kompiling only once."""
     if not opts.rs_file.is_file():
         raise ValueError(f'Input file does not exist: {opts.rs_file}')
 
     if opts.max_workers is not None and opts.max_workers < 1:
         raise ValueError(f'Expected positive integer for `max_workers, got: {opts.max_workers}')
 
-    label = f'{opts.rs_file.stem}.{opts.start_symbol}'
-
     if opts.proof_dir is not None:
-        target_path = opts.proof_dir / label
-        return _prove(opts, target_path, label)
+        if len(opts.start_symbols) == 1:
+            label = f'{opts.rs_file.stem}.{opts.start_symbols[0]}'
+            target_path = opts.proof_dir / label
+        else:
+            # Multiple start symbols share a separate target dir
+            target_path = opts.proof_dir / f'{opts.rs_file.stem}.kompiled'
+        return _prove_multi(opts, target_path)
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         target_path = Path(tmp_dir)
-        return _prove(opts, target_path, label)
+        return _prove_multi(opts, target_path)
 
 
 def prove_with_kmir(
     kmir: KMIR,
     smir_info: SMIRInfo,
     opts: ProveOpts,
 ) -> APRProof:
-    """Run a proof using a pre-built KMIR instance, avoiding redundant kompilation."""
+    """Prove a single symbol using a pre-built KMIR instance.
+
+    The intended use case for this function is internal with the test runner.
+    Use this instead of `prove()` when the caller manages kompilation externally
+    (e.g. the test harness kompiles once and loops over symbols). Only the first
+    entry in `opts.start_symbols` is used; for multi-symbol proving use `prove()`
+    which kompiles once internally but is for external use via `kmir prove`.
+    """
+    assert len(opts.start_symbols) == 1, f'prove_with_kmir handles a single symbol; got {len(opts.start_symbols)}.'
 
     if opts.max_workers is not None and opts.max_workers < 1:
         raise ValueError(f'Expected positive integer for `max_workers, got: {opts.max_workers}')
 
     # No check for rs_file as smir_info already exists
-    label = f'{opts.rs_file.stem}.{opts.start_symbol}'
+    start_symbol = opts.start_symbols[0]
+    label = f'{opts.rs_file.stem}.{start_symbol}'
 
     _LOGGER.info(f'Using pre-built KMIR for proof: {label}')
     proof = apr_proof_from_smir(
         kmir,
         label,
         smir_info,
-        start_symbol=opts.start_symbol,
+        start_symbol=start_symbol,
         proof_dir=opts.proof_dir,
     )
-    if proof.proof_dir is not None and (proof.proof_dir / label).is_dir():
-        smir_info.dump(proof.proof_dir / proof.id / 'smir.json')
+    if opts.proof_dir is not None:
+        kompiled_smir_path = opts.proof_dir / f'{opts.rs_file.stem}.kompiled' / 'smir.json'
+        kompiled_smir_path.parent.mkdir(parents=True, exist_ok=True)
+        smir_info.dump(kompiled_smir_path)
 
     return _advance_proof(kmir, proof, opts, label)
 
 
-def _prove(opts: ProveOpts, target_path: Path, label: str) -> APRProof:
-    if not opts.reload and opts.proof_dir is not None and APRProof.proof_data_exists(label, opts.proof_dir):
-        _LOGGER.info(f'Reading proof from disc: {opts.proof_dir}, {label}')
-        proof = APRProof.read_proof_data(opts.proof_dir, label)
-
-        smir_info = SMIRInfo.from_file(target_path / 'smir.json')
-        kmir = KMIR.from_kompiled_kore(
-            smir_info,
-            target_dir=target_path,
-            extra_module=opts.add_module,
-            bug_report=opts.bug_report,
-            symbolic=True,
-            haskell_target=opts.haskell_target,
-            llvm_lib_target=opts.llvm_lib_target,
-            break_on_function=opts.break_on_function or None,
-        )
+def _prove_multi(opts: ProveOpts, target_path: Path) -> list[APRProof]:
+    """Prove single or multiple symbols with a single kompilation."""
+    labels = [f'{opts.rs_file.stem}.{sym}' for sym in opts.start_symbols]
+
+    if not labels:
+        raise ValueError('No label to prove')
+
+    # Check which proofs can be resumed
+    resumable: dict[str, APRProof] = {}
+    if not opts.reload and opts.proof_dir is not None:
+        for label in labels:
+            if APRProof.proof_data_exists(label, opts.proof_dir):
+                _LOGGER.info(f'Reading proof from disc: {opts.proof_dir}, {label}')
+                resumable[label] = APRProof.read_proof_data(opts.proof_dir, label)
+
+    # Load SMIR info (once)
+    kompiled_smir_path = target_path / 'smir.json'
+    if len(resumable) == len(labels):
+        # All proofs are resumed, load SMIR from saved data
+        smir_info = SMIRInfo.from_file(kompiled_smir_path)
     else:
-        _LOGGER.info(f'Constructing initial proof: {label}')
+        # Need fresh SMIR for at least one proof
         if opts.parsed_smir is not None:
             smir_info = SMIRInfo(opts.parsed_smir)
         elif opts.smir:
             smir_info = SMIRInfo.from_file(opts.rs_file)
         else:
             smir_info = SMIRInfo(cargo_get_smir_json(opts.rs_file, save_smir=opts.save_smir))
 
-        smir_info = smir_info.reduce_to(opts.start_symbol)
+        smir_info = smir_info.reduce_to(opts.start_symbols)
         # Report whether the reduced call graph includes any functions without MIR bodies
         missing_body_syms = [
             sym
@@ -114,28 +132,38 @@ def _prove(opts: ProveOpts, target_path: Path, label: str) -> APRProof:
             _LOGGER.info(f'missing-bodies-present={has_missing} count={len(missing_body_syms)}')
             _LOGGER.debug(f'Missing-body function symbols (first 5): {missing_body_syms[:5]}')
 
-        kmir = KMIR.from_kompiled_kore(
-            smir_info,
-            target_dir=target_path,
-            extra_module=opts.add_module,
-            bug_report=opts.bug_report,
-            symbolic=True,
-            haskell_target=opts.haskell_target,
-            llvm_lib_target=opts.llvm_lib_target,
-            break_on_function=opts.break_on_function or None,
-        )
+    kmir = KMIR.from_kompiled_kore(
+        smir_info,
+        target_dir=target_path,
+        extra_module=opts.add_module,
+        bug_report=opts.bug_report,
+        symbolic=True,
+        haskell_target=opts.haskell_target,
+        llvm_lib_target=opts.llvm_lib_target,
+        break_on_function=opts.break_on_function or None,
+    )
 
-        proof = apr_proof_from_smir(
-            kmir,
-            label,
-            smir_info,
-            start_symbol=opts.start_symbol,
-            proof_dir=opts.proof_dir,
-        )
-        if proof.proof_dir is not None and (proof.proof_dir / label).is_dir():
-            smir_info.dump(proof.proof_dir / proof.id / 'smir.json')
+    smir_info.dump(kompiled_smir_path)
 
-    return _advance_proof(kmir, proof, opts, label)
+    # Prove each symbol sequentially using shared definition
+    results: list[APRProof] = []
+    for label, start_symbol in zip(labels, opts.start_symbols, strict=True):
+        if label in resumable:
+            proof = resumable[label]
+        else:
+            _LOGGER.info(f'Constructing initial proof: {label}')
+            proof = apr_proof_from_smir(
+                kmir,
+                label,
+                smir_info,
+                start_symbol=start_symbol,
+                proof_dir=opts.proof_dir,
+            )
+
+        proof = _advance_proof(kmir, proof, opts, label)
+        results.append(proof)
+
+    return results
 
 
 def _advance_proof(kmir: KMIR, proof: APRProof, opts: ProveOpts, label: str) -> APRProof:

kmir/src/kmir/kmir.py

@@ -121,11 +121,17 @@ def run_smir(
         return result
 
     @staticmethod
-    def prove_program(opts: ProveOpts) -> APRProof:
+    def prove_programs(opts: ProveOpts) -> list[APRProof]:
         from ._prove import prove
 
         return prove(opts)
 
+    @staticmethod
+    def prove_program(opts: ProveOpts) -> APRProof:
+        proofs = KMIR.prove_programs(opts)
+        assert len(proofs) == 1, f'Expected single proof, got {len(proofs)}'
+        return proofs[0]
+
     @staticmethod
     def prove_program_with_kmir(kmir: KMIR, smir_info: SMIRInfo, opts: ProveOpts) -> APRProof:
         from ._prove import prove_with_kmir
@@ -171,12 +177,20 @@ def __init__(self, cterm_show: CTermShow, proof: APRProof, opts: DisplayOpts) ->
         self.smir_info = None
         if opts.smir_info:
             self.smir_info = SMIRInfo.from_file(opts.smir_info)
-        elif (
-            proof.proof_dir is not None
-            and (proof.proof_dir / proof.id).is_dir()
-            and (proof.proof_dir / proof.id / 'smir.json').is_file()
-        ):
-            self.smir_info = SMIRInfo.from_file(proof.proof_dir / proof.id / 'smir.json')
+        elif proof.proof_dir is not None:
+            file_stem = proof.id.rsplit('.', 1)[0]
+            # Single-symbol: smir.json in the proof's own directory
+            label_smir = proof.proof_dir / proof.id / 'smir.json'
+            # Multi-symbol: smir.json in the shared kompiled directory
+            kompiled_smir = proof.proof_dir / f'{file_stem}.kompiled' / 'smir.json'
+            if label_smir.is_file():
+                self.smir_info = SMIRInfo.from_file(label_smir)
+            elif kompiled_smir.is_file():
+                self.smir_info = SMIRInfo.from_file(kompiled_smir)
+            else:
+                _LOGGER.warning('SMIR info not found, span/function annotations unavailable')
+        else:
+            _LOGGER.warning('No SMIR Info or proof dir found, span/function annotations unavailable')
 
     def _span(self, node: KCFG.Node) -> str | None:
         curr_span: int | None = None

kmir/src/kmir/options.py

@@ -75,7 +75,7 @@ class ProveOpts(KMirOpts):
     save_smir: bool
     smir: bool
     parsed_smir: dict | None
-    start_symbol: str
+    start_symbols: list[str]
     add_module: Path | str | None
     break_on_calls: bool
     break_on_function_calls: bool
@@ -111,7 +111,7 @@ def __init__(
         save_smir: bool = False,
         smir: bool = False,
         parsed_smir: dict | None = None,
-        start_symbol: str = 'main',
+        start_symbols: list[str] | None = None,
         break_on_calls: bool = False,
         break_on_function_calls: bool = False,
         break_on_intrinsic_calls: bool = False,
@@ -144,7 +144,12 @@ def __init__(
         self.save_smir = save_smir
         self.smir = smir
         self.parsed_smir = parsed_smir
-        self.start_symbol = start_symbol
+
+        if start_symbols is not None:
+            self.start_symbols = start_symbols
+        else:
+            self.start_symbols = ['main']
+
         self.break_on_calls = break_on_calls
         self.break_on_function_calls = break_on_function_calls
         self.break_on_intrinsic_calls = break_on_intrinsic_calls