pyk: fix add_cell_map_items over-wrapping, add --parser to pyk run

pyk/pyproject.toml

@@ -79,7 +79,7 @@ include = ["src/pyk"]
 [tool.pydocstyle]
 convention = "google"
 add-select = "D204,D401,D404"
-add-ignore = "D1"
+add-ignore = "D1,D202"
 
 [tool.isort]
 profile = "black"

pyk/regression-new/issue-582/1.test.out

@@ -1,3 +1,8 @@
-<k>
-  7 ~> .K
-</k>
+<generatedTop>
+  <k>
+    7 ~> .K
+  </k>
+  <generatedCounter>
+    0
+  </generatedCounter>
+</generatedTop>

pyk/regression-new/krun-deserialize/sum.kore.imp.out

@@ -1,9 +1,13 @@
-<T>
-  <k>
-    .K
-  </k>
-  <state>
-    n |-> 0
-    sum |-> 5050
-  </state>
-</T>
+<generatedTop>
+  <T>
+    <k>
+      .K
+    </k>
+    <state>
+      sum |-> 5050 n |-> 0
+    </state>
+  </T>
+  <generatedCounter>
+    0
+  </generatedCounter>
+</generatedTop>

pyk/regression-new/skipped

@@ -51,7 +51,6 @@ issue-3446
 issue-3520-freshConfig
 issue-3647-debugTokens
 issue-3672-debugParse
-issue-582
 ite-bug
 itp
 json-input
@@ -69,7 +68,6 @@ kprove-error-status
 kprove-java
 kprove-markdown
 kprove-var-equals
-krun-deserialize
 llvm-kompile-type
 llvm-krun
 locations
@@ -97,7 +95,6 @@ search-bound
 set-symbolic-tests
 set_unification
 spec-rule-application
-star-multiplicity
 trace
 unification-lemmas
 useless

pyk/regression-new/star-multiplicity/1.test.out

@@ -5,18 +5,21 @@
   <cells>
     <cell>
       <num>
-        0
+        5
       </num>
       <data>
-        1
+        6
       </data>
     </cell> <cell>
       <num>
-        5
+        0
       </num>
       <data>
-        6
+        1
       </data>
     </cell>
   </cells>
+  <generatedCounter>
+    0
+  </generatedCounter>
 </generatedTop>

pyk/src/pyk/__main__.py

@@ -373,7 +373,7 @@ def exec_run(options: RunOptions) -> None:
     else:
         kompiled_directory = options.definition_dir
     krun = KRun(kompiled_directory)
-    rc, res = krun.krun(pgm_file)
+    rc, res = krun.krun(pgm_file, parser=options.parser)
     print(krun.pretty_print(res))
     sys.exit(rc)
 

pyk/src/pyk/cli/pyk.py

@@ -354,11 +354,13 @@ def get_argument_type() -> dict[str, Callable]:
 class RunOptions(LoggingOptions):
     pgm_file: str
     definition_dir: Path | None
+    parser: str | None
 
     @staticmethod
     def default() -> dict[str, Any]:
         return {
             'definition_dir': None,
+            'parser': None,
         }
 
     @staticmethod
@@ -481,6 +483,7 @@ def create_argument_parser() -> ArgumentParser:
     )
     run_args.add_argument('pgm_file', type=str, help='File program to run in it.')
     run_args.add_argument('--definition', type=dir_path, dest='definition_dir', help='Path to definition to use.')
+    run_args.add_argument('--parser', type=str, help='Command to use as a custom parser.')
 
     prove_args = pyk_args_command.add_parser(
         'prove',

pyk/src/pyk/kast/outer.py

@@ -1135,6 +1135,25 @@ def cell_collection_productions(self) -> tuple[KProduction, ...]:
         """Returns the `KProduction` which are cell collection declarations transitively imported by the main module of this definition."""
         return tuple(prod for module in self.modules for prod in module.cell_collection_productions)
 
+    @cached_property
+    def cell_map_item_info(self) -> dict[str, tuple[str, KSort]]:
+        """Map from cell label to (element-constructor label, cell-map sort).
+
+        Derived from cell-collection productions that carry both ELEMENT and WRAP_ELEMENT
+        attributes and whose element constructor is a known symbol.
+        """
+        # example:
+        # syntax AccountCellMap [cellCollection, hook(MAP.Map)]
+        # syntax AccountCellMap ::= AccountCellMap AccountCellMap [assoc, avoid, cellCollection, comm, element(AccountCellMapItem), function, hook(MAP.concat), unit(.AccountCellMap), wrapElement(<account>)]
+        result: dict[str, tuple[str, KSort]] = {}
+        for ccp in self.cell_collection_productions:
+            if Atts.ELEMENT in ccp.att and Atts.WRAP_ELEMENT in ccp.att:
+                cell_label = ccp.att[Atts.WRAP_ELEMENT]
+                element_ctor = ccp.att[Atts.ELEMENT]
+                if element_ctor in self.symbols:
+                    result[cell_label] = (element_ctor, self.symbols[element_ctor].sort)
+        return result
+
     @cached_property
     def rules(self) -> tuple[KRule, ...]:
         """Returns the `KRule` sentences transitively imported by the main module of this definition."""
@@ -1511,36 +1530,37 @@ def _add_sort_params(_k: KInner) -> KInner:
 
     def add_cell_map_items(self, kast: KInner) -> KInner:
         """Wrap cell-map items in the syntactical wrapper that the frontend generates for them (see `KDefinition.remove_cell_map_items`)."""
-        # example:
-        # syntax AccountCellMap [cellCollection, hook(MAP.Map)]
-        # syntax AccountCellMap ::= AccountCellMap AccountCellMap [assoc, avoid, cellCollection, comm, element(AccountCellMapItem), function, hook(MAP.concat), unit(.AccountCellMap), wrapElement(<account>)]
-
-        cell_wrappers = {}
-        for ccp in self.cell_collection_productions:
-            if Atts.ELEMENT in ccp.att and Atts.WRAP_ELEMENT in ccp.att:
-                cell_wrappers[ccp.att[Atts.WRAP_ELEMENT]] = ccp.att[Atts.ELEMENT]
 
+        # Wrapping is correct only when the parent production expects the cell MAP sort (e.g.
+        # EntryCellMap), not when it expects the individual cell element sort (e.g. EntryCell).
+        # For example, EntryCellMapKey(<entry>(...)) takes EntryCell — the <entry> must NOT be
+        # wrapped, whereas _EntryCellMap_(<entry>(...), ...) expects EntryCellMap — wrapping is needed.
         def _wrap_elements(_k: KInner) -> KInner:
-            if type(_k) is KApply and _k.label.name in cell_wrappers:
-                return KApply(cell_wrappers[_k.label.name], [_k.args[0], _k])
-            return _k
+            if not isinstance(_k, KApply) or _k.label.name not in self.symbols:
+                return _k
+            arg_sorts = self.symbols[_k.label.name].argument_sorts
+            if not any(isinstance(arg, KApply) and arg.label.name in self.cell_map_item_info for arg in _k.args):
+                return _k
+            new_args: list[KInner] = []
+            for arg_sort, arg in zip(arg_sorts, _k.args, strict=True):
+                new_arg = arg
+                if isinstance(arg, KApply):
+                    if arg.label.name in self.cell_map_item_info:
+                        element_ctor, element_ctor_sort = self.cell_map_item_info[arg.label.name]
+                        if arg_sort == element_ctor_sort:
+                            new_arg = KApply(element_ctor, [arg.args[0], arg])
+                new_args.append(new_arg)
+            return _k.let(args=new_args)
 
         # To ensure we don't get duplicate wrappers.
         _kast = self.remove_cell_map_items(kast)
         return bottom_up(_wrap_elements, _kast)
 
     def remove_cell_map_items(self, kast: KInner) -> KInner:
         """Remove cell-map syntactical wrapper items that the frontend generates (see `KDefinition.add_cell_map_items`)."""
-        # example:
-        # syntax AccountCellMap [cellCollection, hook(MAP.Map)]
-        # syntax AccountCellMap ::= AccountCellMap AccountCellMap [assoc, avoid, cellCollection, comm, element(AccountCellMapItem), function, hook(MAP.concat), unit(.AccountCellMap), wrapElement(<account>)]
+        cell_wrappers = {ctor: label for label, (ctor, _) in self.cell_map_item_info.items()}
 
-        cell_wrappers = {}
-        for ccp in self.cell_collection_productions:
-            if Atts.ELEMENT in ccp.att and Atts.WRAP_ELEMENT in ccp.att:
-                cell_wrappers[ccp.att[Atts.ELEMENT]] = ccp.att[Atts.WRAP_ELEMENT]
-
-        def _wrap_elements(_k: KInner) -> KInner:
+        def _unwrap_elements(_k: KInner) -> KInner:
             if (
                 type(_k) is KApply
                 and _k.label.name in cell_wrappers
@@ -1551,7 +1571,7 @@ def _wrap_elements(_k: KInner) -> KInner:
                 return _k.args[1]
             return _k
 
-        return bottom_up(_wrap_elements, kast)
+        return bottom_up(_unwrap_elements, kast)
 
     def empty_config(self, sort: KSort) -> KInner:
         """Given a cell-sort, compute an "empty" configuration for it (all the constructor structure of the configuration in place, but variables in cell positions)."""

pyk/src/pyk/ktool/krun.py

@@ -294,8 +294,8 @@ def run_pattern(
         assert parser.eof
         return res
 
-    def krun(self, input_file: Path) -> tuple[int, KInner]:
-        result = _krun(input_file=input_file, definition_dir=self.definition_dir, output=KRunOutput.KORE)
+    def krun(self, input_file: Path, parser: str | None = None) -> tuple[int, KInner]:
+        result = _krun(input_file=input_file, definition_dir=self.definition_dir, output=KRunOutput.KORE, parser=parser)
         kore = KoreParser(result.stdout).pattern()
         kast = self.kore_to_kast(kore)
         return (result.returncode, kast)

pyk/src/tests/unit/kast/test_definition.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from pyk.kast.att import Atts, KAtt
+from pyk.kast.inner import KApply, KSort, KVariable
+from pyk.kast.outer import KDefinition, KFlatModule, KNonTerminal, KProduction, KTerminal
+
+if TYPE_CHECKING:
+    from typing import Final
+
+    from pyk.kast.inner import KInner
+
+
+# ---------------------------------------------------------------------------
+# Minimal test definition
+#
+# Cell map fragment:
+#   AccountCellMap ::= AccountCellMap AccountCellMap  [cellCollection, element(AccountCellMapItem), wrapElement(<account>)]
+#   AccountCellMap ::= AccountCellMapItem(Int, AccountCell)
+#   AccountCell    ::= <account>(Int, Int)
+#   AccountCell    ::= getEntry(AccountCell)           -- takes element sort, NOT map sort
+# ---------------------------------------------------------------------------
+
+INT: Final = KSort('Int')
+ACCOUNT_CELL_MAP: Final = KSort('AccountCellMap')
+ACCOUNT_CELL: Final = KSort('AccountCell')
+
+_ACCT_MAP_CONCAT: Final = KProduction(
+    sort=ACCOUNT_CELL_MAP,
+    items=[KNonTerminal(ACCOUNT_CELL_MAP), KNonTerminal(ACCOUNT_CELL_MAP)],
+    klabel='_AccountCellMap_',
+    att=KAtt(entries=[Atts.CELL_COLLECTION(None), Atts.ELEMENT('AccountCellMapItem'), Atts.WRAP_ELEMENT('<account>')]),
+)
+
+_ACCT_MAP_ITEM: Final = KProduction(
+    sort=ACCOUNT_CELL_MAP,
+    items=[
+        KTerminal('AccountCellMapItem'),
+        KTerminal('('),
+        KNonTerminal(INT),
+        KTerminal(','),
+        KNonTerminal(ACCOUNT_CELL),
+        KTerminal(')'),
+    ],
+    klabel='AccountCellMapItem',
+)
+
+_ACCOUNT_CELL: Final = KProduction(
+    sort=ACCOUNT_CELL,
+    items=[
+        KTerminal('<account>'),
+        KTerminal('('),
+        KNonTerminal(INT),
+        KTerminal(','),
+        KNonTerminal(INT),
+        KTerminal(')'),
+    ],
+    klabel='<account>',
+)
+
+_GET_ENTRY: Final = KProduction(
+    sort=ACCOUNT_CELL,
+    items=[KTerminal('getEntry'), KTerminal('('), KNonTerminal(ACCOUNT_CELL), KTerminal(')')],
+    klabel='getEntry',
+)
+
+DEFN: Final = KDefinition(
+    'TEST',
+    [KFlatModule('TEST', [_ACCT_MAP_CONCAT, _ACCT_MAP_ITEM, _ACCOUNT_CELL, _GET_ENTRY])],
+)
+
+# ---------------------------------------------------------------------------
+# KDefinition.add_cell_map_items
+# ---------------------------------------------------------------------------
+
+_ACCT_1: Final = KApply('<account>', [KVariable('X', sort=INT), KVariable('Y', sort=INT)])
+_ACCT_2: Final = KApply('<account>', [KVariable('A', sort=INT), KVariable('B', sort=INT)])
+
+ADD_CELL_MAP_ITEMS_DATA: Final = (
+    # Parent expects AccountCellMap (the map sort) — children are wrapped in AccountCellMapItem.
+    (
+        'wraps_when_parent_expects_cell_map_sort',
+        KApply('_AccountCellMap_', [_ACCT_1, _ACCT_2]),
+        KApply(
+            '_AccountCellMap_',
+            [
+                KApply('AccountCellMapItem', [KVariable('X', sort=INT), _ACCT_1]),
+                KApply('AccountCellMapItem', [KVariable('A', sort=INT), _ACCT_2]),
+            ],
+        ),
+    ),
+    # Parent expects AccountCell (the element sort) — the <account> child must NOT be wrapped.
+    # Before the guard fix, _wrap_elements would incorrectly wrap here too.
+    (
+        'no_wrap_when_parent_expects_cell_element_sort',
+        KApply('getEntry', [_ACCT_1]),
+        KApply('getEntry', [_ACCT_1]),
+    ),
+)
+
+
+@pytest.mark.parametrize(
+    'test_id,term,expected',
+    ADD_CELL_MAP_ITEMS_DATA,
+    ids=[test_id for test_id, *_ in ADD_CELL_MAP_ITEMS_DATA],
+)
+def test_add_cell_map_items(test_id: str, term: KInner, expected: KInner) -> None:
+    assert DEFN.add_cell_map_items(term) == expected