Implement kevm.forgetBranch cheatcode

src/kontrol/foundry.py

@@ -20,14 +20,24 @@
 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
+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.kdist import kdist
 from pyk.prelude.bytes import bytesToken
 from pyk.prelude.collections import map_empty
-from pyk.prelude.k import DOTS
+from pyk.prelude.k import DOTS, GENERATED_TOP_CELL
 from pyk.prelude.kbool import notBool
 from pyk.prelude.kint import INT, intToken
 from pyk.prelude.ml import mlEqualsFalse, mlEqualsTrue
@@ -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,75 @@
 _LOGGER: Final = logging.getLogger(__name__)
 
 
+class KontrolSemantics(KEVMSemantics):
+
+    def _check_forget_pattern(self, cterm: CTerm) -> bool:
+        """Given a CTerm, check if the rule 'FOUNDRY-ACCOUNTS.forget' is at the top of the K_CELL.
+
+        This method checks if the 'FOUNDRY-ACCOUNTS.forget' rule is at the top of the `K_CELL` in the given `cterm`.
+        If the rule matches, the resulting substitution is cached in `_cached_subst` for later use in `custom_step`
+        :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('cheatcode_forget', [KVariable('###TERM1'), KVariable('###OPERATOR'), KVariable('###TERM2')]),
+                KVariable('###CONTINUATION'),
+            ]
+        )
+        self._cached_subst = abstract_pattern.match(cterm.cell('K_CELL'))
+        return self._cached_subst is not None
+
+    def _exec_forget_custom_step(self, cterm: CTerm, cterm_symbolic: CTermSymbolic) -> KCFGExtendResult | None:
+        """Given a CTerm, remove the constraint that's placed at the top of the K_CELL in the 'FOUNDRY-ACCOUNTS.forget' cut-rule.
+
+        :param cterm: CTerm of a proof node.
+        :param cterm_symbolic: CTermSymbolic instance
+        :return: If the K_CELL matches the forget_pattern and the constraint is successfully removed, returns a Step with depth 1 along with the new configuration, registering that the `cheatcode_forget` rule has been applied. If the constraint is not found, returns None.
+        """
+        _operators = ['_==Int_', '_!=Int_', '_<=Int_', '_<Int_', '_>=Int_', '_>Int_']
+        _LOGGER.info('Custom step: Forget pattern confirmed')
+        subst = self._cached_subst
+        assert subst is not None
+
+        # Extract the terms and operator from the substitution
+        fst_term = subst['###TERM1']
+        snd_term = subst['###TERM2']
+        operator = subst['###OPERATOR']
+        assert isinstance(operator, KToken)
+        constraint = mlEqualsTrue(KApply(_operators[int(operator.token)], fst_term, snd_term))
+
+        # Simplify the constraint
+        kevm = KEVM(kdist.get('kontrol.foundry'))
+        empty_config = kevm.definition.empty_config(GENERATED_TOP_CELL)
+        simplification_cterm = CTerm.from_kast(empty_config).add_constraint(constraint)
+        result_cterm, _ = cterm_symbolic.simplify(simplification_cterm)
+        target_constraint = single(result_cterm.constraints)
+
+        # Remove the target constraint from the current constraints
+        new_constraints: tuple[KInner, ...] = tuple(c for c in cterm.constraints if c != target_constraint)
+
+        # Check if the constraints have changed
+        if new_constraints == cterm.constraints:
+            _LOGGER.info('Custom step: Target constraint not found in path conditions; returning None')
+            return None
+
+        _LOGGER.info(f'Custom step: Removing constraint: {target_constraint}')
+
+        # Update the K_CELL with the continuation
+        new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'K_CELL', KSequence(subst['###CONTINUATION'])))
+        return Step(CTerm(new_cterm.config, new_constraints), 1, (), ['cheatcode_forget'], cut=True)
+
+    def custom_step(self, cterm: CTerm, cterm_symbolic: CTermSymbolic) -> KCFGExtendResult | None:
+        if self._check_forget_pattern(cterm):
+            return self._exec_forget_custom_step(cterm, cterm_symbolic)
+        else:
+            return super().custom_step(cterm, cterm_symbolic)
+
+    def can_make_custom_step(self, cterm: CTerm) -> bool:
+        return self._check_forget_pattern(cterm) or super().can_make_custom_step(cterm)
+
+
 class FoundryKEVM(KEVM):
     foundry: Foundry
 
@@ -841,7 +922,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 +989,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 +1179,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 +1227,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 +1267,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 +1343,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 +1394,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

@@ -1106,6 +1106,26 @@ Mock functions
       [priority(30)]
 ```
 
+
+Abstraction functions
+---------------------
+
+#### `forgetBranch` - removes a given path constraint.
+
+```
+    function forgetBranch(uint256 op1, ComparisonOperator op, uint256 op2) external;
+```
+
+```k
+    rule [cheatcode.call.abstract]:
+         <k> #cheatcode_call SELECTOR ARGS
+          => #forget ( #asWord(#range(ARGS,0,32)), #asWord(#range(ARGS,32,32)), #asWord(#range(ARGS,64,32)))
+          ...
+         </k>
+      requires SELECTOR ==Int selector ( "forgetBranch(uint256,uint8,uint256)" )
+
+```
+
 Utils
 -----
 
@@ -1691,6 +1711,8 @@ Selectors for **implemented** cheat code functions.
     rule ( selector ( "mockCall(address,bytes,bytes)" )            => 3110212580 )
     rule ( selector ( "mockFunction(address,address,bytes)" )      => 2918731041 )
     rule ( selector ( "copyStorage(address,address)" )             => 540912653  )
+    rule ( selector ( "replaceConstraint(uint256,uint8,uint256,bool)" ) => 1408652492 )
+    rule ( selector ( "forgetBranch(uint256,uint8,uint256)" )      => 1720990067 )
 ```
 
 Selectors for **unimplemented** cheat code functions.

src/kontrol/kdist/foundry.md

@@ -78,6 +78,10 @@ Then, we define helpers in K which can:
              <storage> STORAGE => STORAGE [ #loc(FoundryCheat . Failed) <- 1 ] </storage>
              ...
            </account>
+
+    syntax KItem ::= #forget ( Int , Int , Int ) [symbol(cheatcode_forget)]
+ // -----------------------------------------------------------------------
+    rule [forget]: <k> #forget(_C1,_OP,_C2) => .K ... </k>
 ```
 
 #### Structure of execution

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,
@@ -441,6 +441,8 @@ def create_kcfg_explore() -> KCFGExplore:
                     rule.label for rule in foundry.kevm.definition.all_modules_dict['KONTROL-ASSERTIONS'].rules
                 )
 
+            cut_point_rules.append('FOUNDRY-ACCOUNTS.forget')
+
             extra_lemmas_module: KFlatModule | None = None
             if options.extra_module:
                 extra_module_file, extra_module_name, *_ = options.extra_module.split(':')
@@ -465,7 +467,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/tests/integration/test-data/foundry/test/ForgetBranch.t.sol

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.13;
+
+import "forge-std/Test.sol";
+import "kontrol-cheatcodes/KontrolCheats.sol";
+
+contract ForgetBranchTest is Test, KontrolCheats {
+
+    function test_forgetBranch(uint256 x) public {
+        uint256 y;
+
+        vm.assume(x > 200);
+        kevm.forgetBranch(200, 2, x);
+        if(x > 200){
+            y = 21;
+        } else {
+            y = 42;
+        }
+    }
+}