Allow compilation of a <file>:<module> via --add-module.

docs/add-module.md

@@ -0,0 +1,90 @@
+# Adding Extra Modules to Proofs
+
+## Overview
+
+The `--add-module` option for `kmir prove` allows including extra K rules in a proof without compiling them into the base K definition. This is useful for proof-specific lemmas, simplification rules, or summarized rules that are only needed for certain proofs.
+
+Rules from the added module are injected into the Haskell backend at prove time, alongside the SMIR-generated function equations.
+
+## Usage
+
+### K source files (`.k` or `.md`)
+
+To include rules from a K source file, use the `FILE:MODULE` format:
+
+```bash
+kmir prove test.rs --add-module lemmas.k:MY-LEMMAS
+kmir prove test.rs --add-module lemmas.md:MY-LEMMAS
+```
+
+The file is parsed against the compiled K definition using `kprove --dry-run`. The named module is extracted and its rules are converted to Kore axioms and added to the proof.
+
+If the file contains multiple modules, only the named module's rules are included. For example, given a file with modules `HELPERS` and `MY-LEMMAS`:
+
+```bash
+# Only rules from MY-LEMMAS are included
+kmir prove test.rs --add-module lemmas.k:MY-LEMMAS
+```
+
+### Pre-exported JSON files
+
+To include rules from a JSON module file (generated by `kmir show --to-module`):
+
+```bash
+kmir prove test.rs --add-module module.json
+```
+
+No module name is needed because JSON files are expected to contain exactly one serialized module (as produced by `--to-module`). Multi-module JSON inputs are not supported.
+
+## Module file requirements
+
+### K source files
+
+- The file must have a `.k` or `.md` extension.
+- The module can `imports` other modules that are already in the compiled definition. However, it must not `requires` files that would conflict with the compiled definition (i.e., do not re-require files already compiled in).
+
+For simple rules the dynamically loading pipeline should work reliably. However, there are some known limitations that restrict the kind of rules that can be expressed in a dynamically loaded module. We plan to address them in the future.
+
+### JSON files
+
+- The file must have a `.json` extension.
+- The file must contain a serialized `KFlatModule` (as produced by `kmir show --to-module output.json`).
+
+## Limitations
+
+The dynamic module loading pipeline has known limitations with complex rules. See the summary below, that includes a non-exhaustive list of the identified issues:
+
+| Limitation | Description | Workaround |
+|-----------|-------------|------------|
+| Limited syntax declarations | `kprove` rejects `syntax` productions in proof modules other than for tokens for existing sorts | Compile syntax into the base definition |
+| No `#as` patterns | pyk's `krule_to_kore` does not support `KAs` | Duplicate the matched pattern on the RHS |
+| Function rules may fail | The booster may reject dynamically-added function rules | Compile function rules into the base definition |
+| Sort injection issues | Complex rules with subsort relationships may crash the booster | Compile such rules into the base definition |
+
+## Examples
+
+### Simple simplification lemma
+
+```k
+// lemma.k
+module MY-SIMPLIFICATION
+    imports INT
+
+    rule [zero-add]: 0 +Int X => X [simplification]
+
+endmodule
+```
+
+```bash
+kmir prove test.rs --add-module lemma.k:MY-SIMPLIFICATION
+```
+
+### Proof summary from a previous run
+
+```bash
+# Export a proof node transition as a module
+kmir show proof-id --to-module summary.json --proof-dir ./proofs
+
+# Use it in a new proof
+kmir prove test.rs --add-module summary.json
+```

kmir/src/kmir/__main__.py

@@ -555,9 +555,9 @@ def _arg_parser() -> ArgumentParser:
     )
     prove_parser.add_argument(
         '--add-module',
-        type=Path,
-        metavar='FILE',
-        help='K module file to include (.json format from --to-module)',
+        type=str,
+        metavar='MODULE',
+        help='K module to include. Formats: FILE.k:MODULE or FILE.md:MODULE (K source), FILE.json (from --to-module). See docs/add-module.md for details.',
     )
     prove_parser.add_argument(
         '--max-workers', metavar='N', type=int, help='Maximum number of workers for parallel exploration'

kmir/src/kmir/kmir.py

@@ -57,7 +57,7 @@ def from_kompiled_kore(
         smir_info: SMIRInfo,
         target_dir: Path,
         *,
-        extra_module: Path | None = None,
+        extra_module: str | None = None,
         bug_report: Path | None = None,
         symbolic: bool = True,
         llvm_target: str | None = None,

kmir/src/kmir/kompile.py

@@ -5,19 +5,20 @@
 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, KToken, KVariable
 from pyk.kast.prelude.kint import intToken
 from pyk.kast.prelude.string import stringToken
 from pyk.kdist import kdist
 from pyk.kore.syntax import App, EVar, SortApp, String, Symbol, SymbolDecl
+from pyk.utils import single
 
 from .kmir import KMIR
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
-    from pathlib import Path
     from typing import Any, Final
 
     from pyk.kast.inner import KInner
@@ -162,27 +163,48 @@ def _add_exists_quantifiers(axiom: Axiom) -> Axiom:
     return Axiom(vars=axiom.vars, pattern=new_pattern, attrs=axiom.attrs)
 
 
-def _load_extra_module_rules(kmir: KMIR, module_path: Path) -> list[Sentence]:
-    """Load a K module from JSON and convert rules to Kore axioms.
+def _load_extra_module_rules(kmir: KMIR, module_spec: str) -> list[Sentence]:
+    """Load a K module and convert rules to Kore axioms.
+
+    Supports two formats:
+    - JSON file path: ``path/to/module.json`` (from ``--to-module``)
+    - K source with module name: ``path/to/file.k:MODULE_NAME`` or ``path/to/file.md:MODULE_NAME``
 
     Args:
         kmir: KMIR instance with the definition
-        module_path: Path to JSON module file (from --to-module output.json)
+        module_spec: Module specification string
 
     Returns:
         List of Kore axioms converted from the module rules
     """
     from pyk.kast.outer import KFlatModule, KRule
     from pyk.konvert import krule_to_kore
 
-    _LOGGER.info(f'Loading extra module rules: {module_path}')
+    _LOGGER.info(f'Loading extra module rules: {module_spec}')
 
-    if module_path.suffix != '.json':
-        _LOGGER.warning(f'Only JSON format is supported for --add-module: {module_path}')
-        return []
+    if ':' in module_spec:
+        # K source format: file.k:MODULE_NAME or file.md:MODULE_NAME
+        file_str, module_name = module_spec.rsplit(':', 1)
+        file_path = Path(file_str)
+        if not file_path.is_file():
+            raise ValueError(f'Supplied module path is not a file: {file_path}')
+        if file_path.suffix not in ('.k', '.md'):
+            raise ValueError(f'K source module must be a .k or .md file, got: {file_path}')
 
-    module_dict = json.loads(module_path.read_text())
-    k_module = KFlatModule.from_dict(module_dict)
+        include_dirs = (kdist.which('mir-semantics.haskell').parent,)
+        module_list = kmir.parse_modules(file_path, module_name=module_name, include_dirs=include_dirs)
+        k_module = single(module for module in module_list.modules if module.name == module_name)
+    else:
+        # JSON format: path/to/module.json
+        file_path = Path(module_spec)
+        if not file_path.is_file():
+            raise ValueError(f'Supplied module path is not a file: {file_path}')
+        if file_path.suffix != '.json':
+            raise ValueError(
+                f'Expected .json file or file.k:MODULE_NAME / file.md:MODULE_NAME format, got: {file_path}'
+            )
+        module_dict = json.loads(file_path.read_text())
+        k_module = KFlatModule.from_dict(module_dict)
 
     axioms: list[Sentence] = []
     for sentence in k_module.sentences:
@@ -203,7 +225,7 @@ def kompile_smir(
     target_dir: Path,
     *,
     bug_report: Path | None = None,
-    extra_module: Path | None = None,
+    extra_module: str | None = None,
     symbolic: bool = True,
     llvm_target: str | None = None,
     llvm_lib_target: str | None = None,

kmir/src/kmir/options.py

@@ -76,7 +76,7 @@ class ProveOpts(KMirOpts):
     smir: bool
     parsed_smir: dict | None
     start_symbol: str
-    add_module: Path | None
+    add_module: str | None
     break_on_calls: bool
     break_on_function_calls: bool
     break_on_intrinsic_calls: bool
@@ -127,7 +127,7 @@ def __init__(
         break_every_terminator: bool = False,
         break_every_step: bool = False,
         terminate_on_thunk: bool = False,
-        add_module: Path | None = None,
+        add_module: str | None = None,
         break_on_function: list[str] | None = None,
     ) -> None:
         self.rs_file = rs_file

kmir/src/tests/integration/data/modules/test-add-module-multiple.k

@@ -0,0 +1,47 @@
+module TEST-ADD-MODULE
+    imports KMIR
+
+    rule [BASIC-BLOCK-1-TO-3]: <kmir>
+           <k>
+             ( #execTerminator ( terminator ( ... kind: terminatorKindCall ( ... func: operandConstant ( constOperand ( ... span: span ( 0 ) , userTy: noUserTypeAnnotationIndex , const: mirConst ( ... kind: constantKindZeroSized , ty: ty ( -1 ) , id: mirConstId ( 0 ) ) ) ) , args: .Operands , destination: place ( ... local: local ( 0 ) , projection: .ProjectionElems ) , target: noBasicBlockIdx , unwind: unwindActionContinue ) , span: span ( 0 ) ) ) => #EndProgram ~> .K )
+           </k>
+           <retVal>
+             _RETVAL_CELL
+           </retVal>
+           <currentFunc>
+             ( CURRENTFUNC_CELL => ty ( -1 ) )
+           </currentFunc>
+           <currentFrame>
+             <currentBody>
+               ( _CURRENTBODY_CELL => ListItem ( basicBlock ( ... statements: .Statements , terminator: terminator ( ... kind: terminatorKindReturn , span: span ( 50 ) ) ) ) )
+             </currentBody>
+             <caller>
+               ( CALLER_CELL => CURRENTFUNC_CELL:Ty )
+             </caller>
+             <dest>
+               ( DEST_CELL => place ( ... local: local ( 0 ) , projection: .ProjectionElems ) )
+             </dest>
+             <target>
+               ( TARGET_CELL => noBasicBlockIdx )
+             </target>
+             <unwind>
+               ( UNWIND_CELL => unwindActionContinue )
+             </unwind>
+             <locals>
+               ListItem ( newLocal ( ty ( ( 0 => 1 ) ) , ( mutabilityNot => mutabilityMut ) ) )
+             </locals>
+           </currentFrame>
+           <stack>
+             ( .List => ListItem ( StackFrame ( CALLER_CELL:Ty , DEST_CELL:Place , TARGET_CELL:MaybeBasicBlockIdx , UNWIND_CELL:UnwindAction , ListItem ( newLocal ( ty ( 0 ) , mutabilityNot ) ) ) ) )
+           </stack>
+         </kmir>
+      [priority(20), label(BASIC-BLOCK-1-TO-3)]
+
+endmodule
+
+module SKIPPED-MODULE
+    imports INT
+
+    rule [test-simplification]: X -Int 0 => X [simplification]
+
+endmodule

kmir/src/tests/integration/data/modules/test-add-module.k

@@ -0,0 +1,40 @@
+module TEST-ADD-MODULE
+    imports KMIR
+
+    rule [BASIC-BLOCK-1-TO-3]: <kmir>
+           <k>
+             ( #execTerminator ( terminator ( ... kind: terminatorKindCall ( ... func: operandConstant ( constOperand ( ... span: span ( 0 ) , userTy: noUserTypeAnnotationIndex , const: mirConst ( ... kind: constantKindZeroSized , ty: ty ( -1 ) , id: mirConstId ( 0 ) ) ) ) , args: .Operands , destination: place ( ... local: local ( 0 ) , projection: .ProjectionElems ) , target: noBasicBlockIdx , unwind: unwindActionContinue ) , span: span ( 0 ) ) ) => #EndProgram ~> .K )
+           </k>
+           <retVal>
+             _RETVAL_CELL
+           </retVal>
+           <currentFunc>
+             ( CURRENTFUNC_CELL => ty ( -1 ) )
+           </currentFunc>
+           <currentFrame>
+             <currentBody>
+               ( _CURRENTBODY_CELL => ListItem ( basicBlock ( ... statements: .Statements , terminator: terminator ( ... kind: terminatorKindReturn , span: span ( 50 ) ) ) ) )
+             </currentBody>
+             <caller>
+               ( CALLER_CELL => CURRENTFUNC_CELL:Ty )
+             </caller>
+             <dest>
+               ( DEST_CELL => place ( ... local: local ( 0 ) , projection: .ProjectionElems ) )
+             </dest>
+             <target>
+               ( TARGET_CELL => noBasicBlockIdx )
+             </target>
+             <unwind>
+               ( UNWIND_CELL => unwindActionContinue )
+             </unwind>
+             <locals>
+               ListItem ( newLocal ( ty ( ( 0 => 1 ) ) , ( mutabilityNot => mutabilityMut ) ) )
+             </locals>
+           </currentFrame>
+           <stack>
+             ( .List => ListItem ( StackFrame ( CALLER_CELL:Ty , DEST_CELL:Place , TARGET_CELL:MaybeBasicBlockIdx , UNWIND_CELL:UnwindAction , ListItem ( newLocal ( ty ( 0 ) , mutabilityNot ) ) ) ) )
+           </stack>
+         </kmir>
+      [priority(20), label(BASIC-BLOCK-1-TO-3)]
+
+endmodule

kmir/src/tests/integration/data/modules/test-add-module.md

@@ -0,0 +1,44 @@
+# Test module for --add-module
+
+```k
+module TEST-ADD-MODULE
+    imports KMIR
+
+    rule [BASIC-BLOCK-1-TO-3]: <kmir>
+           <k>
+             ( #execTerminator ( terminator ( ... kind: terminatorKindCall ( ... func: operandConstant ( constOperand ( ... span: span ( 0 ) , userTy: noUserTypeAnnotationIndex , const: mirConst ( ... kind: constantKindZeroSized , ty: ty ( -1 ) , id: mirConstId ( 0 ) ) ) ) , args: .Operands , destination: place ( ... local: local ( 0 ) , projection: .ProjectionElems ) , target: noBasicBlockIdx , unwind: unwindActionContinue ) , span: span ( 0 ) ) ) => #EndProgram ~> .K )
+           </k>
+           <retVal>
+             _RETVAL_CELL
+           </retVal>
+           <currentFunc>
+             ( CURRENTFUNC_CELL => ty ( -1 ) )
+           </currentFunc>
+           <currentFrame>
+             <currentBody>
+               ( _CURRENTBODY_CELL => ListItem ( basicBlock ( ... statements: .Statements , terminator: terminator ( ... kind: terminatorKindReturn , span: span ( 50 ) ) ) ) )
+             </currentBody>
+             <caller>
+               ( CALLER_CELL => CURRENTFUNC_CELL:Ty )
+             </caller>
+             <dest>
+               ( DEST_CELL => place ( ... local: local ( 0 ) , projection: .ProjectionElems ) )
+             </dest>
+             <target>
+               ( TARGET_CELL => noBasicBlockIdx )
+             </target>
+             <unwind>
+               ( UNWIND_CELL => unwindActionContinue )
+             </unwind>
+             <locals>
+               ListItem ( newLocal ( ty ( ( 0 => 1 ) ) , ( mutabilityNot => mutabilityMut ) ) )
+             </locals>
+           </currentFrame>
+           <stack>
+             ( .List => ListItem ( StackFrame ( CALLER_CELL:Ty , DEST_CELL:Place , TARGET_CELL:MaybeBasicBlockIdx , UNWIND_CELL:UnwindAction , ListItem ( newLocal ( ty ( 0 ) , mutabilityNot ) ) ) ) )
+           </stack>
+         </kmir>
+      [priority(20), label(BASIC-BLOCK-1-TO-3)]
+
+endmodule
+```