From 03ca944f6eac27a1f9134c873bd0b06d9b8814b9 Mon Sep 17 00:00:00 2001
From: Andrei <16517508+anvacaru@users.noreply.github.com>
Date: Wed, 4 Dec 2024 14:37:28 +0200
Subject: [PATCH] KEVMSemantics: refactor custom_step and can_make_custom_step

---
 kevm-pyk/src/kevm_pyk/kevm.py | 42 +++++++++++++++++++++--------------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/kevm-pyk/src/kevm_pyk/kevm.py b/kevm-pyk/src/kevm_pyk/kevm.py
index 82f8c788d3..ebd100cf1a 100644
--- a/kevm-pyk/src/kevm_pyk/kevm.py
+++ b/kevm-pyk/src/kevm_pyk/kevm.py
@@ -37,6 +37,7 @@
     from pathlib import Path
     from typing import Final
 
+    from pyk.cterm import CTermSymbolic
     from pyk.kast.inner import KAst, Subst
     from pyk.kast.outer import KFlatModule
     from pyk.kcfg import KCFG
@@ -144,22 +145,11 @@ def _replace(term: KInner) -> KInner:
 
         return CTerm(config=bottom_up(_replace, cterm.config), constraints=cterm.constraints)
 
-    def custom_step(self, cterm: CTerm) -> KCFGExtendResult | None:
-        """Given a CTerm, update the JUMPDESTS_CELL and PROGRAM_CELL if the rule 'EVM.program.load' is at the top of the K_CELL.
-
-        :param cterm: CTerm of a proof node.
-        :return: If the K_CELL matches the load_pattern, a Step with depth 1 is returned together with the new configuration, also registering that the `EVM.program.load` rule has been applied. Otherwise, None is returned.
-        """
-        if self.can_make_custom_step(cterm):
-            subst = self._cached_subst
-            assert subst is not None
-            bytecode_sections = flatten_label('_+Bytes__BYTES-HOOKED_Bytes_Bytes_Bytes', subst['###BYTECODE'])
-            jumpdests_set = compute_jumpdests(bytecode_sections)
-            new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'JUMPDESTS_CELL', jumpdests_set))
-            new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'PROGRAM_CELL', subst['###BYTECODE']))
-            new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'K_CELL', KSequence(subst['###CONTINUATION'])))
-            return Step(new_cterm, 1, (), ['EVM.program.load'], cut=True)
-        return None
+    def custom_step(self, cterm: CTerm, _cterm_symbolic: CTermSymbolic) -> KCFGExtendResult | None:
+        if self._check_load_pattern(cterm):
+            return self._exec_load_custom_step(cterm)
+        else:
+            return None
 
     @staticmethod
     def cut_point_rules(
@@ -207,7 +197,7 @@ def terminal_rules(break_every_step: bool) -> list[str]:
             terminal_rules.append('EVM.step')
         return terminal_rules
 
-    def can_make_custom_step(self, cterm: CTerm) -> bool:
+    def _check_load_pattern(self, cterm: CTerm) -> bool:
         """Given a CTerm, check if the rule 'EVM.program.load' is at the top of the K_CELL.
 
         This method checks if the `EVM.program.load` rule is at the top of the `K_CELL` in the given `cterm`.
@@ -219,6 +209,24 @@ def can_make_custom_step(self, cterm: CTerm) -> bool:
         self._cached_subst = load_pattern.match(cterm.cell('K_CELL'))
         return self._cached_subst is not None
 
+    def _exec_load_custom_step(self, cterm: CTerm) -> KCFGExtendResult:
+        """Given a CTerm, update the JUMPDESTS_CELL and PROGRAM_CELL if the rule 'EVM.program.load' is at the top of the K_CELL.
+
+        :param cterm: CTerm of a proof node.
+        :return: If the K_CELL matches the load_pattern, a Step with depth 1 is returned together with the new configuration, also registering that the `EVM.program.load` rule has been applied.
+        """
+        subst = self._cached_subst
+        assert subst is not None
+        bytecode_sections = flatten_label('_+Bytes__BYTES-HOOKED_Bytes_Bytes_Bytes', subst['###BYTECODE'])
+        jumpdests_set = compute_jumpdests(bytecode_sections)
+        new_cterm = CTerm.from_kast(set_cell(cterm.kast, 'JUMPDESTS_CELL', jumpdests_set))
+        new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'PROGRAM_CELL', subst['###BYTECODE']))
+        new_cterm = CTerm.from_kast(set_cell(new_cterm.kast, 'K_CELL', KSequence(subst['###CONTINUATION'])))
+        return Step(new_cterm, 1, (), ['EVM.program.load'], cut=True)
+
+    def can_make_custom_step(self, cterm: CTerm) -> bool:
+        return self._check_load_pattern(cterm)
+
     def is_mergeable(self, ct1: CTerm, ct2: CTerm) -> bool:
         """Given two CTerms of Edges' targets, check if they are mergeable.