Add cheatcodes that generate symbolic variables with custom names (#902)

* add #rename rule and new cheatcodes

* update custom_step

* Add tests

* format

* formatting

* remove prints

* add test to end to end

* remove unimplemented random*(*,string) cheatcodes

* use free_vars

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,
@@ -55,8 +65,10 @@
     from collections.abc import Iterable
     from typing import Any, Final
 
+    from pyk.cterm import CTermSymbolic
     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 +94,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)
+
+        # 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, _cterm_symbolic: CTermSymbolic) -> KCFGExtendResult | None:
+        if self._check_rename_pattern(cterm):
+            return self._exec_rename_custom_step(cterm)
+        else:
+            return super().custom_step(cterm, _cterm_symbolic)
+
+    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 +924,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 +991,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 +1181,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 +1229,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 +1269,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 +1345,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 +1396,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,

src/kontrol/kdist/cheatcodes.md

@@ -339,6 +339,10 @@ This rule then takes the address using `#asWord(#range(ARGS, 0, 32))` and makes
          <k> #cheatcode_call SELECTOR ARGS => #loadAccount #asWord(ARGS) ~> #setSymbolicStorage #asWord(ARGS) ... </k>
       requires SELECTOR ==Int selector ( "symbolicStorage(address)" )
         orBool SELECTOR ==Int selector ( "setArbitraryStorage(address)")
+
+    rule [cheatcode.call.withName.symbolicStorage]:
+         <k> #cheatcode_call SELECTOR ARGS => #loadAccount #asWord(#range(ARGS,0,32)) ~> #setSymbolicStorage #asWord(#range(ARGS,0,32)) Bytes2String(#range(ARGS, 96, #asWord(#range(ARGS, 64, 32)))) ... </k>
+      requires SELECTOR ==Int selector ( "symbolicStorage(address,string)" )
 ```
 
 #### `copyStorage` - Copies the storage of one account into another.
@@ -376,6 +380,14 @@ This rule returns a symbolic integer of up to the bit width that was sent as an
        andBool 0 <Int #asWord(ARGS) andBool #asWord(ARGS) <=Int 32
        ensures 0 <=Int ?WORD andBool ?WORD <Int 2 ^Int (8 *Int #asWord(ARGS))
        [preserves-definedness]
+
+    rule [cheatcode.call.withName.freshUInt]:
+         <k> #cheatcode_call SELECTOR ARGS => #rename(?WORD, Bytes2String(#range(ARGS, 96, #asWord(#range(ARGS, 64, 32))))) ... </k>
+         <output> _ => #buf(32, ?WORD) </output>
+      requires SELECTOR ==Int selector ( "freshUInt(uint8,string)" )
+       andBool 0 <Int #asWord(#range(ARGS, 0, 32)) andBool #asWord(#range(ARGS, 0, 32)) <=Int 32
+       ensures 0 <=Int ?WORD andBool ?WORD <Int 2 ^Int (8 *Int #asWord(#range(ARGS, 0, 32)))
+       [preserves-definedness]
 ```
 
 #### `randomUint` - Returns a single symbolic unsigned integer of a given size.
@@ -441,6 +453,13 @@ This rule returns a symbolic boolean value being either 0 (false) or 1 (true).
         orBool SELECTOR ==Int selector ( "randomBool()" )
        ensures #rangeBool(?WORD)
        [preserves-definedness]
+
+    rule [cheatcode.call.withName.freshBool]:
+         <k> #cheatcode_call SELECTOR ARGS => #rename(?WORD, Bytes2String(#range(ARGS, 64, #asWord(#range(ARGS, 32, 32))))) ... </k>
+         <output> _ => #buf(32, ?WORD) </output>
+      requires SELECTOR ==Int selector ( "freshBool(string)" )
+       ensures #rangeBool(?WORD)
+       [preserves-definedness]
 ```
 
 #### `freshBytes` - Returns a fully symbolic byte array value of the given length.
@@ -466,6 +485,16 @@ This rule returns a fully symbolic byte array value of the given length.
          orBool SELECTOR ==Int selector ( "randomBytes(uint256)" )
        ensures lengthBytes(?BYTES) ==Int #asWord(ARGS)
       [preserves-definedness]
+
+    rule [cheatcode.call.withName.freshBytes]:
+         <k> #cheatcode_call SELECTOR ARGS => #rename(?BYTES, Bytes2String(#range(ARGS, 96, #asWord(#range(ARGS, 64, 32))))) ... </k>
+         <output> _ =>
+            #buf(32, 32) +Bytes #buf(32, #asWord(#range(ARGS, 0, 32))) +Bytes ?BYTES
+            +Bytes #buf ( ( ( notMaxUInt5 &Int ( #asWord(#range(ARGS, 0, 32)) +Int maxUInt5 ) ) -Int #asWord(#range(ARGS, 0, 32)) ) , 0 )
+         </output>
+      requires SELECTOR ==Int selector ( "freshBytes(uint256,string)" )
+       ensures lengthBytes(?BYTES) ==Int #asWord(#range(ARGS, 0, 32))
+      [preserves-definedness]
 ```
 
 This rule returns a fully symbolic byte array value of length 4.
@@ -512,6 +541,13 @@ This rule returns a symbolic address value.
         orBool SELECTOR ==Int selector ( "randomAddress()" )
        ensures #rangeAddress(?WORD) andBool ?WORD =/=Int #address(FoundryTest) andBool ?WORD =/=Int #address(FoundryCheat)
        [preserves-definedness]
+
+    rule [foundry.call.withName.freshAddress]:
+         <k> #cheatcode_call SELECTOR ARGS => #rename(?WORD, Bytes2String(#range(ARGS, 64, #asWord(#range(ARGS, 32, 32))))) ... </k>
+         <output> _ => #buf(32, ?WORD) </output>
+      requires SELECTOR ==Int selector ( "freshAddress(string)" )
+       ensures #rangeAddress(?WORD) andBool ?WORD =/=Int #address(FoundryTest) andBool ?WORD =/=Int #address(FoundryCheat)
+       [preserves-definedness]
 ```
 
 Expecting the next call to revert
@@ -1109,6 +1145,15 @@ Mock functions
 Utils
 -----
 
+ - Defining a new production `#rename` for all the types for which we generate symbolic values.
+We don't care about the values because they will be processed in the `custom_step` function in Python.
+
+```k
+    syntax RenameTarget ::= Int | Bytes | Map
+    syntax KItem ::= #rename ( RenameTarget , String ) [symbol(foundry_rename)]
+ // ---------------------------------------------------------------------------
+    rule [rename]: <k> #rename(_,_) => .K ... </k>
+```
 - `#loadAccount ACCT` creates a new, empty account for `ACCT` if it does not already exist. Otherwise, it has no effect.
 
 ```k
@@ -1202,7 +1247,9 @@ Utils
 - `#setSymbolicStorage ACCTID` takes a given account and makes its storage fully symbolic.
 
 ```k
-     syntax KItem ::= "#setSymbolicStorage" Int [symbol(foundry_setSymbolicStorage)]
+    syntax KItem ::= "#setSymbolicStorage" Int [symbol(foundry_setSymbolicStorage)]
+                   | "#setSymbolicStorage" Int String [symbol(foundry_setSymbolicStorageWithName)]
+ // ----------------------------------------------------------------------------------------------
 ```
 
 ```{.k .symbolic}
@@ -1213,6 +1260,14 @@ Utils
            <origStorage> _ => ?STORAGE </origStorage>
            ...
          </account>
+
+    rule <k> #setSymbolicStorage ACCTID NAME => #rename(?STORAGE, NAME) ... </k>
+         <account>
+           <acctID> ACCTID </acctID>
+           <storage> _ => ?STORAGE </storage>
+           <origStorage> _ => ?STORAGE </origStorage>
+           ...
+         </account>
 ```
 
 - `#copyStorage ACCTFROM ACCTTO` copies the storage of ACCTFROM to be the storage of ACCTTO
@@ -1669,18 +1724,23 @@ Selectors for **implemented** cheat code functions.
     rule ( selector ( "expectEmit(bool,bool,bool,bool,address)" )  => 2176505587 )
     rule ( selector ( "sign(uint256,bytes32)" )                    => 3812747940 )
     rule ( selector ( "symbolicStorage(address)" )                 => 769677742  )
+    rule ( selector ( "symbolicStorage(address,string)" )          => 745143816  )
     rule ( selector ( "setArbitraryStorage(address)" )             => 3781367863 )
     rule ( selector ( "freshUInt(uint8)" )                         => 625253732  )
+    rule ( selector ( "freshUInt(uint8,string)" )                  => 1530912521 )
     rule ( selector ( "randomUint(uint256)" )                      => 3481396892 )
     rule ( selector ( "randomUint()" )                             => 621954864  )
     rule ( selector ( "randomUint(uint256,uint256)" )              => 3592095003 )
     rule ( selector ( "freshBool()" )                              => 2935720297 )
+    rule ( selector ( "freshBool(string)" )                        => 525694724  )
     rule ( selector ( "randomBool()" )                             => 3451987645 )
     rule ( selector ( "freshBytes(uint256)" )                      => 1389402351 )
+    rule ( selector ( "freshBytes(uint256,string)" )               => 390682600  )
     rule ( selector ( "randomBytes(uint256)" )                     => 1818047145 )
     rule ( selector ( "randomBytes4()" )                           => 2608649593 )
     rule ( selector ( "randomBytes8()" )                           => 77050021   )
     rule ( selector ( "freshAddress()" )                           => 2363359817 )
+    rule ( selector ( "freshAddress(string)" )                     => 1202084987 )
     rule ( selector ( "randomAddress()" )                          => 3586058741 )
     rule ( selector ( "prank(address)" )                           => 3395723175 )
     rule ( selector ( "prank(address,address)" )                   => 1206193358 )

src/kontrol/prove.py

@@ -11,7 +11,7 @@
 from typing import TYPE_CHECKING, Any, ContextManager, NamedTuple
 
 from kevm_pyk.cli import EVMChainOptions
-from kevm_pyk.kevm import KEVM, KEVMSemantics, _process_jumpdests
+from kevm_pyk.kevm import KEVM, _process_jumpdests
 from kevm_pyk.utils import KDefinition__expand_macros, abstract_cell_vars, run_prover
 from multiprocess.pool import Pool  # type: ignore
 from pyk.cterm import CTerm, CTermSymbolic
@@ -35,7 +35,7 @@
 from pyk.utils import hash_str, run_process_2, single, unique
 from rich.progress import Progress, SpinnerColumn, TaskID, TextColumn, TimeElapsedColumn
 
-from .foundry import Foundry, foundry_to_xml
+from .foundry import Foundry, KontrolSemantics, foundry_to_xml
 from .hevm import Hevm
 from .options import ConfigType, TraceOptions
 from .solc_to_k import Contract, hex_string_to_int
@@ -381,7 +381,7 @@ def create_kcfg_explore() -> KCFGExplore:
                 )
                 return KCFGExplore(
                     cterm_symbolic,
-                    kcfg_semantics=KEVMSemantics(auto_abstract_gas=options.auto_abstract_gas),
+                    kcfg_semantics=KontrolSemantics(auto_abstract_gas=options.auto_abstract_gas),
                     id=test.id,
                 )
 
@@ -425,7 +425,7 @@ def create_kcfg_explore() -> KCFGExplore:
                         }
                     ),
                 )
-            cut_point_rules = KEVMSemantics.cut_point_rules(
+            cut_point_rules = KontrolSemantics.cut_point_rules(
                 options.break_on_jumpi,
                 options.break_on_jump,
                 options.break_on_calls,
@@ -465,7 +465,7 @@ def create_kcfg_explore() -> KCFGExplore:
                 max_depth=options.max_depth,
                 max_iterations=options.max_iterations,
                 cut_point_rules=cut_point_rules,
-                terminal_rules=KEVMSemantics.terminal_rules(options.break_every_step),
+                terminal_rules=KontrolSemantics.terminal_rules(options.break_every_step),
                 counterexample_info=options.counterexample_info,
                 max_frontier_parallel=options.max_frontier_parallel,
                 fail_fast=options.fail_fast,

src/kontrol/utils.py

@@ -10,6 +10,7 @@
 
 if TYPE_CHECKING:
     from typing import Final
+    from pyk.cterm import CTerm
     from argparse import Namespace
 
 import os
@@ -25,6 +26,21 @@
 _LOGGER: Final = logging.getLogger(__name__)
 
 
+def ensure_name_is_unique(name: str, cterm: CTerm) -> str:
+    """Ensure that a given name for a KVariable is unique within the context of a CTerm.
+
+    :param name: name of a KVariable
+    :param cterm: cterm
+    :return: Returns the name if it's not used, otherwise appends a suffix.
+    :rtype: str
+    """
+    if name not in cterm.free_vars:
+        return name
+
+    index = next(i for i in range(len(cterm.free_vars) + 1) if f'{name}_{i}' not in cterm.free_vars)
+    return f'{name}_{index}'
+
+
 def check_k_version() -> None:
     expected_k_version = KVersion.parse(f'v{pyk.__version__}')
     actual_k_version = k_version()

src/tests/integration/test-data/end-to-end-prove-all

@@ -1,4 +1,5 @@
 CounterTest.test_Increment()
+RandomVarTest.test_custom_names()
 RandomVarTest.test_randomBool()
 RandomVarTest.test_randomAddress()
 RandomVarTest.test_randomUint()