Support --start-symbol argument and lift common functionality into

Tools

Author: gtrepta

kmir/src/kmir/__main__.py

@@ -6,8 +6,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from pyk.kast.inner import KSort, Subst
-
 from kmir.build import semantics
 from kmir.convert_from_definition.v2parser import parse_json
 
@@ -23,6 +21,7 @@ class KMirOpts: ...
 class RunOpts(KMirOpts):
     input_file: Path
     depth: int
+    start_symbol: str
 
 
 def _kmir_run(opts: RunOpts) -> None:
@@ -35,10 +34,7 @@ def _kmir_run(opts: RunOpts) -> None:
 
     kmir_kast, _ = parse_result
 
-    subst = Subst({'$PGM': kmir_kast})
-    init_config = subst.apply(tools.definition.init_config(KSort('GeneratedTopCell')))
-    init_kore = tools.krun.kast_to_kore(init_config, KSort('GeneratedTopCell'))
-    result = tools.krun.run_pattern(init_kore, depth=opts.depth)
+    result = tools.run_parsed(kmir_kast, opts.start_symbol, opts.depth)
 
     print(tools.kprint.kore_to_pretty(result))
 
@@ -60,6 +56,9 @@ def _arg_parser() -> ArgumentParser:
     run_parser = command_parser.add_parser('run', help='run stable MIR programs')
     run_parser.add_argument('input_file', metavar='FILE', help='MIR program to run')
     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'
+    )
 
     return parser
 
@@ -69,7 +68,7 @@ def _parse_args(args: Sequence[str]) -> KMirOpts:
 
     match ns.command:
         case 'run':
-            return RunOpts(input_file=Path(ns.input_file).resolve(), depth=ns.depth)
+            return RunOpts(input_file=Path(ns.input_file).resolve(), depth=ns.depth, start_symbol=ns.start_symbol)
         case _:
             raise AssertionError()
 

kmir/src/kmir/tools.py

@@ -2,15 +2,19 @@
 
 from typing import TYPE_CHECKING
 
+from pyk.kast.inner import KSort, Subst
 from pyk.kast.outer import read_kast_definition
 from pyk.ktool.krun import KRun
+from pyk.prelude.string import stringToken
 
 from .kparse import KParse
 
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from pyk.kast import KInner
     from pyk.kast.outer import KDefinition
+    from pyk.kore.syntax import Pattern
     from pyk.ktool.kprint import KPrint
 
 
@@ -42,3 +46,14 @@ def krun(self) -> KRun:
     @property
     def definition(self) -> KDefinition:
         return self.__definition
+
+    def run_parsed(self, parsed_smir: KInner, start_symbol: KInner | str = 'main', depth: int | None = None) -> Pattern:
+        if isinstance(start_symbol, str):
+            start_symbol = stringToken(start_symbol)
+
+        subst = Subst({'$PGM': parsed_smir, '$STARTSYM': start_symbol})
+        init_config = subst.apply(self.definition.init_config(KSort('GeneratedTopCell')))
+        init_kore = self.krun.kast_to_kore(init_config, KSort('GeneratedTopCell'))
+        result = self.krun.run_pattern(init_kore, depth=depth)
+
+        return result

kmir/src/tests/integration/test_integration.py

@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING
 
 import pytest
-from pyk.kast.inner import KApply, KSort, KToken, Subst
+from pyk.kast.inner import KApply, KSort, KToken
 
 from kmir.convert_from_definition.v2parser import Parser
 
@@ -187,10 +187,7 @@ def test_exec_smir(
     assert parsed is not None
     kmir_kast, _ = parsed
 
-    subst = Subst({'$PGM': kmir_kast})
-    init_config = subst.apply(tools.definition.init_config(KSort('GeneratedTopCell')))
-    init_kore = tools.krun.kast_to_kore(init_config, KSort('GeneratedTopCell'))
-    result = tools.krun.run_pattern(init_kore, depth=depth)
+    result = tools.run_parsed(kmir_kast, depth=depth)
 
     with output_kast.open('r') as f:
         expected = f.read().rstrip()