komet prove options for advanced users (#64)

* implement `--extra-module`

* implement `--view-node` and `--remove-node`

* add progress display to `komet prove run`

* Set Version: 0.1.57

---------

Co-authored-by: devops <[email protected]>

Author: bbyalcinkaya

package/version

@@ -1 +1 @@
-0.1.56
+0.1.57

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "komet"
-version = "0.1.56"
+version = "0.1.57"
 description = "K tooling for the Soroban platform"
 authors = [
     "Runtime Verification, Inc. <[email protected]>",

src/komet/kasmer.py

@@ -50,6 +50,7 @@
 
     from hypothesis.strategies import SearchStrategy
     from pyk.kast.inner import KInner
+    from pyk.kast.outer import KFlatModule
     from pyk.kore.syntax import Pattern
     from pyk.proof import APRProof, EqualityProof
     from pyk.utils import BugReport
@@ -63,9 +64,11 @@ class Kasmer:
     """Reads soroban contracts, and runs tests for them."""
 
     definition: SorobanDefinition
+    extra_module: KFlatModule | None
 
-    def __init__(self, definition: SorobanDefinition) -> None:
+    def __init__(self, definition: SorobanDefinition, extra_module: KFlatModule | None = None) -> None:
         self.definition = definition
+        self.extra_module = extra_module
 
     def _which(self, cmd: str) -> Path:
         path_str = shutil.which(cmd)
@@ -290,7 +293,7 @@ def make_steps(*args: KInner) -> KInner:
 
         claim, _ = cterm_build_claim(name, lhs, rhs)
 
-        return run_claim(name, claim, proof_dir, bug_report)
+        return run_claim(name, claim, self.extra_module, proof_dir, bug_report)
 
     def deploy_and_run(
         self, contract_wasm: Path, child_wasms: tuple[Path, ...], max_examples: int = 100, id: str | None = None
@@ -306,26 +309,13 @@ def deploy_and_run(
         Raises:
             AssertionError if any of the tests fail
         """
-        print(f'Processing contract: {contract_wasm.stem}')
-
-        bindings = self.contract_bindings(contract_wasm)
-        has_init = 'init' in (b.name for b in bindings)
+        test_bindings, has_init = self.read_bindings(contract_wasm, id)
 
         contract_kast = self.kast_from_wasm(contract_wasm)
         child_kasts = tuple(self.kast_from_wasm(c) for c in child_wasms)
 
         conf, subst = self.deploy_test(contract_kast, child_kasts, has_init)
 
-        test_bindings = [b for b in bindings if b.name.startswith('test_') and (id is None or b.name == id)]
-
-        if id is None:
-            print(f'Discovered {len(test_bindings)} test functions:')
-        elif not test_bindings:
-            raise KeyError(f'Test function {id!r} not found.')
-        else:
-            print('Selected a single test function:')
-        print()
-
         failed: list[FuzzError] = []
         with FuzzProgress(test_bindings, max_examples) as progress:
             for task in progress.fuzz_tasks:
@@ -370,21 +360,23 @@ def deploy_and_prove(
         Raises:
             KSorobanError if a proof fails
         """
-        bindings = self.contract_bindings(contract_wasm)
-        has_init = 'init' in (b.name for b in bindings)
+        test_bindings, has_init = self.read_bindings(contract_wasm, id)
 
         contract_kast = self.kast_from_wasm(contract_wasm)
         child_kasts = tuple(self.kast_from_wasm(c) for c in child_wasms)
 
         conf, subst = self.deploy_test(contract_kast, child_kasts, has_init)
 
-        test_bindings = [b for b in bindings if b.name.startswith('test_') and (id is None or b.name == id)]
-
-        for binding in test_bindings:
-            print(binding.name)
-            proof = self.run_prove(conf, subst, binding, always_allocate, proof_dir, bug_report)
-            if proof.status == ProofStatus.FAILED:
-                raise KSorobanError(proof.summary)
+        with FuzzProgress(test_bindings, 1) as progress:
+            for task in progress.fuzz_tasks:
+                task.start()
+                proof = self.run_prove(conf, subst, task.binding, always_allocate, proof_dir, bug_report)
+                if proof.status == ProofStatus.PASSED:
+                    task.advance()
+                    task.end()
+                else:
+                    task.fail()
+                    raise KSorobanError(proof.summary)
 
     def prove_raw(
         self,
@@ -411,10 +403,27 @@ def prove_raw(
                 if proof.status == ProofStatus.FAILED:
                     raise KSorobanError(proof)
             else:
-                proof = run_claim(claim.label, claim, proof_dir, bug_report)
+                proof = run_claim(claim.label, claim, self.extra_module, proof_dir, bug_report)
                 if proof.status == ProofStatus.FAILED:
                     raise KSorobanError(proof.summary)
 
+    def read_bindings(self, contract_wasm: Path, id: str | None) -> tuple[list[ContractBinding], bool]:
+        print(f'Processing contract: {contract_wasm.stem}')
+
+        bindings = self.contract_bindings(contract_wasm)
+        has_init = 'init' in (b.name for b in bindings)
+        test_bindings = [b for b in bindings if b.name.startswith('test_') and (id is None or b.name == id)]
+
+        if id is None:
+            print(f'Discovered {len(test_bindings)} test functions:')
+        elif not test_bindings:
+            raise KeyError(f'Test function {id!r} not found.')
+        else:
+            print('Selected a single test function:')
+        print()
+
+        return test_bindings, has_init
+
 
 @dataclass(frozen=True)
 class ContractBinding:

src/komet/komet.py

@@ -28,10 +28,10 @@
     from collections.abc import Iterator
     from subprocess import CompletedProcess
 
+    from pyk.kast.outer import KFlatModule
     from pyk.utils import BugReport
 
-
-sys.setrecursionlimit(4000)
+sys.setrecursionlimit(8000)
 
 
 class Backend(Enum):
@@ -59,6 +59,7 @@ def main() -> None:
                 dir_path=args.directory,
                 wasm=wasm,
                 id=args.id,
+                extra_module=args.extra_module,
                 always_allocate=args.always_allocate,
                 proof_dir=args.proof_dir,
                 bug_report=args.bug_report,
@@ -67,10 +68,24 @@ def main() -> None:
             assert args.proof_dir is not None
             _exec_prove_view(proof_dir=args.proof_dir, id=args.id)
 
+        if args.prove_command == 'view-node':
+            assert args.proof_dir is not None
+            assert args.id is not None
+            assert args.node is not None
+            _exec_prove_view_node(proof_dir=args.proof_dir, id=args.id, node=args.node)
+        if args.prove_command == 'remove-node':
+            assert args.proof_dir is not None
+            assert args.id is not None
+            assert args.node is not None
+            _exec_prove_remove_node(proof_dir=args.proof_dir, id=args.id, node=args.node)
     elif args.command == 'prove-raw':
         assert args.claim_file is not None
         _exec_prove_raw(
-            claim_file=args.claim_file, label=args.label, proof_dir=args.proof_dir, bug_report=args.bug_report
+            claim_file=args.claim_file,
+            label=args.label,
+            extra_module=args.extra_module,
+            proof_dir=args.proof_dir,
+            bug_report=args.bug_report,
         )
 
     raise AssertionError()
@@ -89,10 +104,11 @@ def _exec_prove_raw(
     *,
     claim_file: Path,
     label: str | None,
+    extra_module: KFlatModule | None,
     proof_dir: Path | None,
     bug_report: BugReport | None = None,
 ) -> None:
-    kasmer = Kasmer(symbolic_definition)
+    kasmer = Kasmer(symbolic_definition, extra_module)
     try:
         kasmer.prove_raw(claim_file, label, proof_dir, bug_report)
         exit(0)
@@ -157,12 +173,13 @@ def _exec_prove_run(
     dir_path: Path | None,
     wasm: Path | None,
     id: str | None,
+    extra_module: KFlatModule | None,
     always_allocate: bool,
     proof_dir: Path | None,
     bug_report: BugReport | None = None,
 ) -> None:
     dir_path = Path.cwd() if dir_path is None else dir_path
-    kasmer = Kasmer(symbolic_definition)
+    kasmer = Kasmer(symbolic_definition, extra_module)
 
     child_wasms: tuple[Path, ...] = ()
 
@@ -202,6 +219,20 @@ def _exec_prove_view(*, proof_dir: Path, id: str) -> None:
     sys.exit(0)
 
 
+def _exec_prove_view_node(*, proof_dir: Path, id: str, node: int) -> None:
+    proof = APRProof.read_proof_data(proof_dir, id)
+    config = proof.kcfg.node(node).cterm.config
+    print(symbolic_definition.krun.pretty_print(config))
+    sys.exit(0)
+
+
+def _exec_prove_remove_node(*, proof_dir: Path, id: str, node: int) -> None:
+    proof = APRProof.read_proof_data(proof_dir, id)
+    proof.prune(node)
+    proof.write_proof_data()
+    sys.exit(0)
+
+
 @contextmanager
 def _preprocessed(program: Path) -> Iterator[Path]:
     program_text = program.read_text()
@@ -219,6 +250,14 @@ def _exit_with_output(cp: CompletedProcess) -> None:
     sys.exit(status)
 
 
+def extra_module_arg(extra_module: str) -> KFlatModule:
+    extra_module_file, extra_module_name, *_ = extra_module.split(':')
+    extra_module_path = Path(extra_module_file)
+    if not extra_module_path.is_file():
+        raise ValueError(f'Supplied --extra-module path is not a file: {extra_module_path}')
+    return symbolic_definition.parse_lemmas_module(extra_module_path, extra_module_name)
+
+
 def _argument_parser() -> ArgumentParser:
     parser = ArgumentParser(prog='komet')
     command_parser = parser.add_subparsers(dest='command', required=True)
@@ -243,24 +282,22 @@ def _argument_parser() -> ArgumentParser:
     prove_parser.add_argument(
         'prove_command',
         default='run',
-        choices=('run', 'view'),
+        choices=('run', 'view', 'view-node', 'remove-node'),
         metavar='COMMAND',
         help='Proof command to run. One of (%(choices)s)',
     )
-    prove_parser.add_argument('--proof-dir', type=ensure_dir_path, default=None, help='Output directory for proofs')
-    prove_parser.add_argument('--bug-report', type=bug_report_arg, default=None, help='Bug report directory for proofs')
+    prove_parser.add_argument('--node', type=int)
+    _add_common_prove_arguments(prove_parser)
+
     _add_common_test_arguments(prove_parser)
 
     prove_raw_parser = command_parser.add_parser(
         'prove-raw',
         help='Prove K claims directly from a file, bypassing the usual test contract structure; intended for development and advanced users.',
     )
     prove_raw_parser.add_argument('claim_file', metavar='CLAIM_FILE', type=file_path, help='path to claim file')
-    prove_raw_parser.add_argument('--proof-dir', type=ensure_dir_path, default=None, help='Output directory for proofs')
-    prove_raw_parser.add_argument(
-        '--bug-report', type=bug_report_arg, default=None, help='Bug report directory for proofs'
-    )
     prove_raw_parser.add_argument('--label', help='Label of the K claim in the file')
+    _add_common_prove_arguments(prove_raw_parser)
 
     return parser
 
@@ -280,3 +317,18 @@ def _add_common_test_arguments(parser: ArgumentParser) -> None:
         default=None,
         help='The working directory for the command (defaults to the current working directory).',
     )
+
+
+def _add_common_prove_arguments(parser: ArgumentParser) -> None:
+    parser.add_argument('--proof-dir', type=ensure_dir_path, default=None, help='Output directory for proofs')
+    parser.add_argument('--bug-report', type=bug_report_arg, default=None, help='Bug report directory for proofs')
+    parser.add_argument(
+        '--extra-module',
+        dest='extra_module',
+        default=None,
+        type=extra_module_arg,
+        help=(
+            'Extra module with user-defined lemmas to include for verification (which must import KASMER module).'
+            'Format is <file>:<module name>.'
+        ),
+    )

src/komet/proof.py

@@ -18,7 +18,7 @@
     from pathlib import Path
 
     from pyk.kast.inner import KInner
-    from pyk.kast.outer import KClaim
+    from pyk.kast.outer import KClaim, KFlatModule
     from pyk.utils import BugReport
 
 
@@ -37,14 +37,20 @@ def _explore_context(id: str, bug_report: BugReport | None) -> Iterator[KCFGExpl
 class SorobanSemantics(DefaultSemantics): ...
 
 
-def run_claim(id: str, claim: KClaim, proof_dir: Path | None = None, bug_report: BugReport | None = None) -> APRProof:
+def run_claim(
+    id: str,
+    claim: KClaim,
+    extra_module: KFlatModule | None = None,
+    proof_dir: Path | None = None,
+    bug_report: BugReport | None = None,
+) -> APRProof:
     if proof_dir is not None and APRProof.proof_data_exists(id, proof_dir):
         proof = APRProof.read_proof_data(proof_dir, id)
     else:
         proof = APRProof.from_claim(symbolic_definition.kdefinition, claim=claim, logs={}, proof_dir=proof_dir)
 
     with _explore_context(id, bug_report) as kcfg_explore:
-        prover = APRProver(kcfg_explore)
+        prover = APRProver(kcfg_explore, extra_module=extra_module)
         prover.advance_proof(proof)
 
     proof.write_proof_data()

src/komet/utils.py

@@ -1,24 +1,29 @@
 from __future__ import annotations
 
+import logging
 from functools import cached_property
+from subprocess import CalledProcessError
 from typing import TYPE_CHECKING
 
-from pyk.kast.outer import read_kast_definition
+from pyk.kast.outer import KRule, read_kast_definition
 from pyk.kdist import kdist
 from pyk.konvert import kast_to_kore
 from pyk.ktool.kompile import DefinitionInfo
 from pyk.ktool.kprove import KProve
 from pyk.ktool.krun import KRun
+from pyk.utils import single
 
 if TYPE_CHECKING:
     from pathlib import Path
     from subprocess import CompletedProcess
-    from typing import Any
+    from typing import Any, Final
 
     from pyk.kast.inner import KInner, KSort
-    from pyk.kast.outer import KDefinition
+    from pyk.kast.outer import KDefinition, KFlatModule
     from pyk.ktool.kompile import KompileBackend
 
+_LOGGER: Final = logging.getLogger(__name__)
+
 
 class KSorobanError(RuntimeError): ...
 
@@ -67,6 +72,22 @@ def krun_with_kast(self, pgm: KInner, sort: KSort | None = None, **kwargs: Any)
         kore_term = kast_to_kore(self.kdefinition, pgm, sort=sort)
         return self.krun.run_process(kore_term, **kwargs)
 
+    def parse_lemmas_module(self, module_path: Path, module_name: str) -> KFlatModule:
+        try:
+            modules = self.kprove.parse_modules(module_path, module_name=module_name)
+        except CalledProcessError as e:
+            _LOGGER.error('Could not parse extra module:')
+            _LOGGER.error(e.stderr)
+            raise e
+
+        module = single(module for module in modules.modules if module.name == module_name)
+
+        non_rule_sentences = [sent for sent in module.sentences if not isinstance(sent, KRule)]
+        if non_rule_sentences:
+            raise ValueError(f'Supplied --extra-module contains non-Rule sentences: {non_rule_sentences}')
+
+        return module
+
 
 concrete_definition = SorobanDefinition(kdist.get('soroban-semantics.llvm'))
 library_definition = SorobanDefinition(kdist.get('soroban-semantics.llvm-library'))