Teach kmir to reduce SMIR K definitions over a set of CFG roots (#845)

Co-authored-by: Stevengre <zhaojianhong96@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: 3 people

kmir/src/kmir/__main__.py

@@ -23,6 +23,7 @@
     LinkOpts,
     ProveOpts,
     PruneOpts,
+    ReduceOpts,
     RunOpts,
     SectionEdgeOpts,
     ShowOpts,
@@ -251,6 +252,17 @@ def _kmir_link(opts: LinkOpts) -> None:
     result.dump(opts.output_file)
 
 
+def _kmir_reduce(opts: ReduceOpts) -> None:
+    smir_info = SMIRInfo.from_file(opts.smir_file)
+    original = len(smir_info.items)
+    reduced = smir_info.reduce_to(opts.roots)
+    reduced.dump(opts.output_file)
+    _LOGGER.info(
+        f'Reduced {original} -> {len(reduced.items)} items'
+        f' ({original - len(reduced.items)} pruned), written to {opts.output_file}'
+    )
+
+
 def kmir(args: Sequence[str]) -> None:
     ns = _arg_parser().parse_args(args)
     opts = _parse_args(ns)
@@ -272,6 +284,8 @@ def kmir(args: Sequence[str]) -> None:
             _kmir_prove(opts)
         case LinkOpts():
             _kmir_link(opts)
+        case ReduceOpts():
+            _kmir_reduce(opts)
         case _:
             raise AssertionError()
 
@@ -575,6 +589,27 @@ def _arg_parser() -> ArgumentParser:
         default='linker_output.smir.json',
     )
 
+    reduce_parser = command_parser.add_parser(
+        'reduce',
+        help='Reduce SMIR to functions reachable from given roots',
+        parents=[kcli_args.logging_args],
+    )
+    reduce_parser.add_argument('smir_file', metavar='SMIR_JSON', help='SMIR JSON file to reduce')
+    reduce_parser.add_argument(
+        '--roots',
+        '-r',
+        required=True,
+        metavar='ROOTS',
+        help='Comma-separated root function names, or @file for newline-separated file',
+    )
+    reduce_parser.add_argument(
+        '--output-file',
+        '-o',
+        metavar='FILE',
+        help='Output file (default: reduced.smir.json)',
+        default='reduced.smir.json',
+    )
+
     return parser
 
 
@@ -677,6 +712,12 @@ def _parse_args(ns: Namespace) -> KMirOpts:
                 smir_files=ns.smir_files,
                 output_file=ns.output_file,
             )
+        case 'reduce':
+            return ReduceOpts(
+                smir_file=ns.smir_file,
+                roots=ns.roots,
+                output_file=ns.output_file,
+            )
         case _:
             raise AssertionError()
 

kmir/src/kmir/_prove.py

@@ -82,7 +82,6 @@ def _prove(opts: ProveOpts, target_path: Path, label: str) -> APRProof:
             if 'MonoItemFn' in item['mono_item_kind'] and item['mono_item_kind']['MonoItemFn'].get('body') is None
         ]
         has_missing = len(missing_body_syms) > 0
-        _LOGGER.info(f'Reduced items table size {len(smir_info.items)}')
         if has_missing:
             _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/src/kmir/options.py

@@ -323,3 +323,20 @@ class LinkOpts(KMirOpts):
     def __init__(self, smir_files: list[str], output_file: str | None = None) -> None:
         self.smir_files = [Path(f) for f in smir_files]
         self.output_file = Path(output_file) if output_file is not None else Path('linker_output.smir.json')
+
+
+@dataclass
+class ReduceOpts(KMirOpts):
+    smir_file: Path
+    output_file: Path
+    roots: list[str]
+
+    def __init__(self, smir_file: str, roots: str, output_file: str | None = None) -> None:
+        self.smir_file = Path(smir_file)
+        self.output_file = Path(output_file) if output_file is not None else Path('reduced.smir.json')
+        # Support @file syntax for reading roots from a file
+        if roots.startswith('@'):
+            roots_file = Path(roots[1:])
+            self.roots = list(filter(None, [r.strip() for r in roots_file.read_text().splitlines()]))
+        else:
+            self.roots = [r.strip() for r in roots.split(',') if r.strip()]

kmir/src/kmir/smir.py

@@ -10,6 +10,7 @@
 from .ty import EnumT, RefT, StructT, Ty, TypeMetadata, UnionT
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
     from pathlib import Path
     from typing import Final
 
@@ -180,13 +181,19 @@ def spans(self) -> dict[int, tuple[Path, int, int, int, int]]:
     def _is_func(item: dict[str, dict]) -> bool:
         return 'MonoItemFn' in item['mono_item_kind']
 
-    def reduce_to(self, start_name: str) -> SMIRInfo:
-        # returns a new SMIRInfo with all _items_ removed that are not reachable from the named function
-        start_ty = self.function_tys[start_name]
+    def reduce_to(self, start_symbols: str | Sequence[str]) -> SMIRInfo:
+        # returns a new SMIRInfo with all _items_ removed that are not reachable from the named function(s)
+        match start_symbols:
+            case str(symbol):
+                start_tys = [Ty(self.function_tys[symbol])]
+            case [*symbols] if symbols and all(isinstance(sym, str) for sym in symbols):
+                start_tys = [Ty(self.function_tys[sym]) for sym in symbols]
+            case _:
+                raise ValueError('SMIRInfo.reduce_to() received an invalid start_symbol')
 
-        _LOGGER.debug(f'Reducing items, starting at {start_ty}. Call Edges {self.call_edges}')
+        _LOGGER.debug(f'Reducing items, starting at {start_tys}. Call Edges {self.call_edges}')
 
-        reachable = compute_closure(Ty(start_ty), self.call_edges)
+        reachable = compute_closure(start_tys, self.call_edges)
 
         _LOGGER.debug(f'Reducing to reachable Tys {reachable}')
 
@@ -249,8 +256,8 @@ def call_edges(self) -> dict[Ty, set[Ty]]:
         return result
 
 
-def compute_closure(start: Ty, edges: dict[Ty, set[Ty]]) -> set[Ty]:
-    work = deque([start])
+def compute_closure(start_nodes: Sequence[Ty], edges: dict[Ty, set[Ty]]) -> set[Ty]:
+    work = deque(start_nodes)
     reached = set()
     finished = False
     while not finished:

kmir/src/tests/integration/test_integration.py

@@ -656,3 +656,40 @@ def test_schema_kapply_parse(
     json_data, expected_term, expected_sort = test_case
 
     assert parser.parse_mir_json(json_data, expected_sort.name) == (expected_term, expected_sort)
+
+
+ARITH_SMIR = PROVE_DIR / 'arith.smir.json'
+
+
+def test_reduce_standalone() -> None:
+    """Test that kmir reduce correctly prunes SMIR items by reachability."""
+    smir_data = json.loads(ARITH_SMIR.read_text())
+    info = SMIRInfo(smir_data)
+    assert len(info.items) == 11
+
+    # Single root 'add' — should keep 1 item
+    reduced_add = info.reduce_to('add')
+    assert len(reduced_add.items) == 1
+
+    # Single root 'mul' — should keep 1 item (independent from add)
+    reduced_mul = info.reduce_to('mul')
+    assert len(reduced_mul.items) == 1
+
+    # Multiple roots — should keep strictly more than either alone
+    reduced_multi = info.reduce_to(['add', 'mul'])
+    assert len(reduced_multi.items) == 2
+
+    # 'main' calls both add and mul — should keep all 3
+    reduced_main = info.reduce_to('main')
+    assert len(reduced_main.items) == 3
+
+    # Roundtrip: save reduced SMIR and reload it
+    with tempfile.NamedTemporaryFile(suffix='.smir.json', delete=False, mode='w') as f:
+        f.write(json.dumps(reduced_multi._smir))
+        reduced_path = Path(f.name)
+
+    try:
+        reloaded = SMIRInfo(json.loads(reduced_path.read_text()))
+        assert len(reloaded.items) == 2
+    finally:
+        reduced_path.unlink()

kmir/src/tests/unit/test_unit.py

@@ -59,5 +59,5 @@
 def test_compute_closure(test_case: tuple[str, Ty, dict[Ty, set[Ty]], list[Ty]]) -> None:
     _, start, edges, expected = test_case
 
-    result = compute_closure(start, edges)
+    result = compute_closure([start], edges)
     assert result == expected