Add cheatcodes that generate symbolic variables with custom names

src/kontrol/foundry.py

@@ -20,10 +20,19 @@
 from kevm_pyk.kevm import KEVM, KEVMNodePrinter, KEVMSemantics
 from kevm_pyk.utils import byte_offset_to_lines, legacy_explore, print_failure_info, print_model
 from pyk.cterm import CTerm
-from pyk.kast.inner import KApply, KInner, KSort, KToken, KVariable
-from pyk.kast.manip import cell_label_to_var_name, collect, extract_lhs, flatten_label, minimize_term, top_down
+from pyk.kast.inner import KApply, KInner, KSequence, KSort, KToken, KVariable, Subst
+from pyk.kast.manip import (
+    cell_label_to_var_name,
+    collect,
+    extract_lhs,
+    flatten_label,
+    minimize_term,
+    set_cell,
+    top_down,
+)
 from pyk.kast.outer import KDefinition, KFlatModule, KImport, KRequire
 from pyk.kcfg import KCFG
+from pyk.kcfg.kcfg import Step
 from pyk.kcfg.minimize import KCFGMinimizer
 from pyk.prelude.bytes import bytesToken
 from pyk.prelude.collections import map_empty
@@ -44,6 +53,7 @@
     _read_digest_file,
     append_to_file,
     empty_lemmas_file_contents,
+    ensure_name_is_unique,
     foundry_toml_extra_contents,
     kontrol_file_contents,
     kontrol_toml_file_contents,
@@ -57,6 +67,7 @@
 
     from pyk.kast.outer import KAst
     from pyk.kcfg.kcfg import NodeIdLike
+    from pyk.kcfg.semantics import KCFGExtendResult
     from pyk.kcfg.tui import KCFGElem
     from pyk.proof.implies import RefutationProof
     from pyk.proof.show import NodePrinter
@@ -82,6 +93,77 @@
 _LOGGER: Final = logging.getLogger(__name__)
 
 
+class KontrolSemantics(KEVMSemantics):
+
+    @staticmethod
+    def cut_point_rules(
+        break_on_jumpi: bool,
+        break_on_jump: bool,
+        break_on_calls: bool,
+        break_on_storage: bool,
+        break_on_basic_blocks: bool,
+        break_on_load_program: bool,
+    ) -> list[str]:
+        return ['FOUNDRY-CHEAT-CODES.rename'] + KEVMSemantics.cut_point_rules(
+            break_on_jumpi,
+            break_on_jump,
+            break_on_calls,
+            break_on_storage,
+            break_on_basic_blocks,
+            break_on_load_program,
+        )
+
+    def _check_rename_pattern(self, cterm: CTerm) -> bool:
+        """Given a CTerm, check if the rule 'FOUNDRY-CHEAT-CODES.rename' is at the top of the K_CELL.
+
+        :param cterm: The CTerm representing the current state of the proof node.
+        :return: `True` if the pattern matches and a custom step can be made; `False` otherwise.
+        """
+        abstract_pattern = KSequence(
+            [
+                KApply('foundry_rename', [KVariable('###RENAME_TARGET'), KVariable('###NEW_NAME')]),
+                KVariable('###CONTINUATION'),
+            ]
+        )
+        self._cached_subst = abstract_pattern.match(cterm.cell('K_CELL'))
+        return self._cached_subst is not None
+
+    def _exec_rename_custom_step(self, cterm: CTerm) -> KCFGExtendResult | None:
+        subst = self._cached_subst
+        assert subst is not None
+
+        # Extract the target var and new name from the substitution
+        target_var = subst['###RENAME_TARGET']
+        name_token = subst['###NEW_NAME']
+        assert type(target_var) is KVariable
+        assert type(name_token) is KToken
+
+        # Ensure the name is unique
+        name_str = name_token.token[1:-1]
+        if len(name_str) == 0:
+            _LOGGER.warning('Name of symbolic variable cannot be empty. Reverting to the default name.')
+            return None
+        name = ensure_name_is_unique(name_str, cterm.config)
+
+        # Replace var in configuration and constraints
+        rename_subst = Subst({target_var.name: KVariable(name, target_var.sort)})
+        config = rename_subst(cterm.config)
+        constraints = [rename_subst(constraint) for constraint in cterm.constraints]
+        new_cterm = CTerm.from_kast(set_cell(config, 'K_CELL', KSequence(subst['###CONTINUATION'])))
+
+        _LOGGER.info(f'Renaming {target_var.name} to {name}')
+        return Step(CTerm(new_cterm.config, constraints), 1, (), ['foundry_rename'], cut=True)
+
+    def custom_step(self, cterm: CTerm) -> KCFGExtendResult | None:
+        if self._check_rename_pattern(cterm):
+            return self._exec_rename_custom_step(cterm)
+        else:
+            return super().custom_step(cterm)
+
+    def can_make_custom_step(self, cterm: CTerm) -> bool:
+        return self._check_rename_pattern(cterm) or super().can_make_custom_step(cterm)
+
+
 class FoundryKEVM(KEVM):
     foundry: Foundry
 
@@ -841,7 +923,7 @@ def foundry_show(
     if options.failure_info:
         with legacy_explore(
             foundry.kevm,
-            kcfg_semantics=KEVMSemantics(),
+            kcfg_semantics=KontrolSemantics(),
             id=test_id,
             smt_timeout=options.smt_timeout,
             smt_retry_limit=options.smt_retry_limit,
@@ -908,7 +990,7 @@ def _collect_klabel(_k: KInner) -> None:
         ]
         sentences = [sent for sent in sentences if not _contains_foundry_klabel(sent.body)]
         sentences = [
-            sent for sent in sentences if not KEVMSemantics().is_terminal(CTerm.from_kast(extract_lhs(sent.body)))
+            sent for sent in sentences if not KontrolSemantics().is_terminal(CTerm.from_kast(extract_lhs(sent.body)))
         ]
         if len(sentences) == 0:
             _LOGGER.warning(f'No claims or rules retained for proof {proof.id}')
@@ -1098,7 +1180,7 @@ def foundry_simplify_node(
 
     with legacy_explore(
         foundry.kevm,
-        kcfg_semantics=KEVMSemantics(),
+        kcfg_semantics=KontrolSemantics(),
         id=apr_proof.id,
         bug_report=options.bug_report,
         kore_rpc_command=kore_rpc_command,
@@ -1146,7 +1228,7 @@ def check_cells_equal(cell: str, nodes: Iterable[KCFG.Node]) -> bool:
     check_cells_ne = [check_cell for check_cell in check_cells if not check_cells_equal(check_cell, nodes)]
 
     if check_cells_ne:
-        if not all(KEVMSemantics().same_loop(nodes[0].cterm, nd.cterm) for nd in nodes):
+        if not all(KontrolSemantics().same_loop(nodes[0].cterm, nd.cterm) for nd in nodes):
             raise ValueError(f'Nodes {options.nodes} cannot be merged because they differ in: {check_cells_ne}')
 
     anti_unification = nodes[0].cterm
@@ -1186,7 +1268,7 @@ def foundry_step_node(
 
     with legacy_explore(
         foundry.kevm,
-        kcfg_semantics=KEVMSemantics(),
+        kcfg_semantics=KontrolSemantics(),
         id=apr_proof.id,
         bug_report=options.bug_report,
         kore_rpc_command=kore_rpc_command,
@@ -1262,7 +1344,7 @@ def foundry_section_edge(
 
     with legacy_explore(
         foundry.kevm,
-        kcfg_semantics=KEVMSemantics(),
+        kcfg_semantics=KontrolSemantics(),
         id=apr_proof.id,
         bug_report=options.bug_report,
         kore_rpc_command=kore_rpc_command,
@@ -1313,7 +1395,7 @@ def foundry_get_model(
 
     with legacy_explore(
         foundry.kevm,
-        kcfg_semantics=KEVMSemantics(),
+        kcfg_semantics=KontrolSemantics(),
         id=proof.id,
         bug_report=options.bug_report,
         kore_rpc_command=kore_rpc_command,