Make the test_change_quorum property ~3x faster

kmxwasm/src/kmxwasm/ast/generic.py

@@ -1,3 +1,5 @@
+from dataclasses import dataclass
+
 from pyk.kast.inner import KApply, KInner, KSequence, bottom_up
 
 
@@ -101,7 +103,12 @@ def replace_contents_with_path(root: KInner, path: list[str], replacement: KInne
     return root.let(args=new_args)
 
 
-def get_with_path(root: KInner, path: list[str]) -> KInner:
+@dataclass(frozen=True)
+class ComponentNotFound:
+    element: str
+
+
+def find_with_path_internal(root: KInner, path: list[str]) -> KInner | ComponentNotFound:
     assert path
     for element in path:
         assert isinstance(root, KApply)
@@ -116,11 +123,25 @@ def get_with_path(root: KInner, path: list[str]) -> KInner:
                 raise ValueError(f'Path component found twice: {element!r}')
             found_arg = arg
         if not found_arg:
-            raise ValueError(f'Path component not found: {found_arg!r}')
+            return ComponentNotFound(element)
         root = found_arg
     return root
 
 
+def get_with_path(root: KInner, path: list[str]) -> KInner:
+    result = find_with_path_internal(root, path)
+    if isinstance(result, ComponentNotFound):
+        raise ValueError(f'Path component not found: {result.element!r}')
+    return result
+
+
+def find_with_path(root: KInner, path: list[str]) -> KInner | None:
+    result = find_with_path_internal(root, path)
+    if isinstance(result, ComponentNotFound):
+        return None
+    return result
+
+
 def get_single_argument_kapply_contents_path(root: KInner, path: list[str]) -> KInner:
     node = get_with_path(root, path)
     if not isinstance(node, KApply):

kmxwasm/src/kmxwasm/ast/mx.py

@@ -3,25 +3,36 @@
 from pyk.kast.inner import KApply, KInner, KSequence, KToken, collect
 
 from .collections import cell_map, full_list, k_map, simple_list
-from .generic import get_with_path, replace_with_path, set_ksequence_cell_contents, set_single_argument_kapply_contents
+from .generic import (
+    find_with_path,
+    get_with_path,
+    replace_with_path,
+    set_ksequence_cell_contents,
+    set_single_argument_kapply_contents,
+)
 
 COMMANDS_CELL_NAME = '<commands>'
 INSTRS_CELL_NAME = '<instrs>'
 K_CELL_NAME = '<k>'
 WASM_CELL_NAME = '<wasm>'
 CONTRACT_MOD_IDX_CELL_NAME = '<contractModIdx>'
+MAIN_STORE_CELL_NAME = '<mainStore>'
+FUNCS_CELL_NAME = '<funcs>'
 
 MANDOS_CELL_PATH = ['<foundry>', '<mandos>']
 NODE_CELL_PATH = MANDOS_CELL_PATH + ['<elrond>', '<node>']
 CALL_STATE_PATH = NODE_CELL_PATH + ['<callState>']
 CALL_STACK_PATH = NODE_CELL_PATH + ['<callStack>']
+INTERIM_STATES_PATH = NODE_CELL_PATH + ['<interimStates>']
 
 COMMANDS_CELL_PATH = NODE_CELL_PATH + [COMMANDS_CELL_NAME]
 K_CELL_PATH = MANDOS_CELL_PATH + [K_CELL_NAME]
 WASM_CELL_PATH = CALL_STATE_PATH + [WASM_CELL_NAME]
 INSTRS_CELL_PATH = WASM_CELL_PATH + [INSTRS_CELL_NAME]
 CONTRACT_MOD_IDX_CELL_PATH = CALL_STATE_PATH + [CONTRACT_MOD_IDX_CELL_NAME]
 
+FUNCS_PATH = WASM_CELL_PATH + [MAIN_STORE_CELL_NAME, FUNCS_CELL_NAME]
+
 
 # TODO: Move these to the elrond-semantics repository.
 def listBytes(items: Iterable[KInner]) -> KInner:  # noqa: N802
@@ -94,22 +105,32 @@ def get_first_command_name(root: KInner) -> str | None:
     return first.label.name
 
 
-def get_instrs_cell(root: KInner) -> KApply:
-    result = get_with_path(root, INSTRS_CELL_PATH)
+def get_instrs_cell(root: KInner) -> KApply | None:
+    result = find_with_path(root, INSTRS_CELL_PATH)
+    if not result:
+        return None
     assert isinstance(result, KApply)
     return result
 
 
-def instrs_cell_contents(root: KInner) -> KSequence:
+def instrs_cell_contents(root: KInner) -> KSequence | None:
     instrs = get_instrs_cell(root)
+    if not instrs:
+        return None
     assert len(instrs.args) == 1
     seq = instrs.args[0]
     assert isinstance(seq, KSequence)
     return seq
 
 
+def replace_instrs_cell(root: KInner, replacement: KSequence) -> KInner:
+    return replace_with_path(root, INSTRS_CELL_PATH, replacement=KApply('<instrs>', replacement))
+
+
 def get_first_instr(root: KInner) -> KApply | None:
     seq = instrs_cell_contents(root)
+    if not seq:
+        return None
     if not seq.items:
         return None
     first = seq.items[0]
@@ -152,6 +173,16 @@ def cfg_changes_call_stack(root: KInner) -> bool:
     return instr == 'endFoundryImmediately'
 
 
+def cfg_changes_interim_states(root: KInner) -> bool:
+    command = get_first_command_name(root)
+    if command in ['pushWorldState', 'popWorldState', 'dropWorldState']:
+        return True
+    instr = get_first_instr_name(root)
+    if not instr:
+        return False
+    return instr == 'endFoundryImmediately'
+
+
 def set_commands_cell_contents(root: KInner, contents: KSequence) -> KInner:
     return set_ksequence_cell_contents(root, COMMANDS_CELL_NAME, contents)
 

kmxwasm/src/kmxwasm/ast/wasm.py

@@ -1,11 +1,12 @@
 from typing import Iterable
 
-from pyk.kast.inner import KApply, KInner, KSequence
+from pyk.kast.inner import KApply, KInner, KSequence, KSort
 
 from .collections import cell_map, simple_list
 from .generic import set_ksequence_cell_contents
 
 # TODO: Move these to the wasm-semantics repository
+FUNCS_CELL_MAP = KSort('FuncDefCellMap')
 
 
 def funcDefCellMap(func_defs: Iterable[KInner]) -> KInner:  # noqa: N802

kmxwasm/src/kmxwasm/property.py

@@ -5,21 +5,23 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from pyk.kast.inner import KApply
+from pyk.cterm import CTerm
+from pyk.kast.inner import KApply, KSequence, KVariable
 from pyk.kcfg import KCFG
 from pyk.kcfg.show import KCFGShow
 from pyk.kore.rpc import KoreClientError
 from pyk.prelude.utils import token
 
+from .ast.mx import replace_instrs_cell, set_accounts_cell_content, set_call_stack_cell_content, set_interim_states_cell_content
 from .build import HASKELL, kbuild_semantics
 from .json import load_json_kcfg, load_json_kclaim, write_kcfg_json
 from .property_testing.paths import KBUILD_DIR, KBUILD_ML_PATH, ROOT
 from .property_testing.printers import print_node
-from .property_testing.running import RunException, Stuck, Success, run_claim, split_edge
+from .property_testing.running import RunException, Stuck, Success, profile_step, run_claim, split_edge
 from .property_testing.wasm_krun_initializer import WasmKrunInitializer
 from .timing import Timer
 
-sys.setrecursionlimit(4000)
+sys.setrecursionlimit(8000)
 
 
 def usage_error() -> None:
@@ -30,6 +32,7 @@ def usage_error() -> None:
     print('  python3 -m src.kmxwasm.property --tree')
     print('  python3 -m src.kmxwasm.property --bisect-after <id>')
     print('  python3 -m src.kmxwasm.property --show-node <id>')
+    print('  python3 -m src.kmxwasm.property --profile <node-id> [--remove <node-id-csv>] [--step <number>]')
     sys.exit(-1)
 
 
@@ -177,6 +180,71 @@ def run(self) -> None:
             raise NotImplementedError(f'Unknown run_claim result: {type(result)}')
 
 
+@dataclass(frozen=True)
+class Profile(Action):
+    node_id: int
+    remove: list[int]
+    depth: int
+    kcfg_path: Path
+    booster: bool
+
+    def run(self) -> None:
+        with kbuild_semantics(
+            output_dir=KBUILD_DIR, config_file=KBUILD_ML_PATH, target=HASKELL, booster=self.booster
+        ) as tools:
+            t = Timer('Loading kcfg')
+            kcfg = load_json_kcfg(self.kcfg_path)
+            t.measure()
+
+            t = Timer('Removing nodes')
+            for node_id in self.remove:
+                kcfg.remove_node(node_id)
+            t.measure()
+
+            t = Timer('Prepare profile node')
+            node = kcfg.get_node(self.node_id)
+            assert node
+            instrs = KSequence([KApply('aNop')] * self.depth)
+            new_config = replace_instrs_cell(node.cterm.config, instrs)
+            new_config = set_call_stack_cell_content(new_config, KVariable('CallStackVar'))
+            new_config = set_interim_states_cell_content(new_config, KVariable('InterimStatesVar'))
+            new_config = set_accounts_cell_content(new_config, KVariable('AccountsVar'))
+            kcfg.replace_node(node.id, cterm=CTerm(new_config, node.cterm.constraints))
+            t.measure()
+
+            result = profile_step(
+                tools,
+                restart_kcfg=kcfg,
+                node_id=self.node_id,
+                depth=self.depth,
+            )
+
+            if isinstance(result, Success):
+                print('Success')
+                return
+            if isinstance(result, RunException):
+                print('Exception')
+                show = KCFGShow(tools.printer)
+                for line in show.pretty(result.kcfg):
+                    print(line)
+                print('Last node:')
+                print('Printing: ', result.last_processed_node)
+                if result.last_processed_node != -1:
+                    node = result.kcfg.get_node(result.last_processed_node)
+                    if node is None:
+                        print(f'Node not found: {result.last_processed_node}')
+                    else:
+                        print_node(tools, node)
+                else:
+                    print('No node to print.')
+                if isinstance(result.exception, KoreClientError):
+                    print('code=', result.exception.code)
+                    print('message=', result.exception.message)
+                    print('data=', result.exception.data)
+                raise result.exception
+            raise NotImplementedError(f'Unknown run_claim result: {type(result)}')
+
+
 @dataclass(frozen=True)
 class ShowNode(Action):
     node_id: int
@@ -261,6 +329,13 @@ def read_flags() -> Action:
         default=None,
         help='CSV of node after which the edge should be split.',
     )
+    parser.add_argument(
+        '--profile',
+        dest='profile',
+        type=int,
+        required=False,
+        help='Run a single sequence of steps starting at the given node.',
+    )
     parser.add_argument(
         '--run',
         dest='run_node',
@@ -301,6 +376,13 @@ def read_flags() -> Action:
     if args.bisect_after:
         return BisectAfter(args.bisect_after, Path(args.kcfg), booster=args.booster)
 
+    to_remove = []
+    if args.remove_nodes:
+        to_remove = [int(n_id) for n_id in args.remove_nodes.split(',')]
+
+    if args.profile:
+        return Profile(args.profile, remove=to_remove, depth=args.step, kcfg_path=Path(args.kcfg), booster=args.booster)
+
     if args.claimfile is None:
         usage_error()
         sys.exit(-1)
@@ -309,10 +391,6 @@ def read_flags() -> Action:
         print(f'Input file ({claim_path}) does not exist.')
         sys.exit(-1)
 
-    to_remove = []
-    if args.remove_nodes:
-        to_remove = [int(n_id) for n_id in args.remove_nodes.split(',')]
-
     run: int | None = None
     if args.run_node != -1:
         run = args.run_node

kmxwasm/src/kmxwasm/property_testing/implication.py

@@ -3,14 +3,14 @@
 from ..ast.mx import commands_cell_contents, instrs_cell_contents, k_cell_contents
 
 
-def quick_ksequence_implication_check(antecedent: KSequence, consequent: KSequence) -> bool:
-    if antecedent.arity == 0:
-        if consequent.arity == 0:
+def quick_ksequence_implication_check(antecedent: KSequence | None, consequent: KSequence | None) -> bool:
+    if not antecedent or antecedent.arity == 0:
+        if not consequent or consequent.arity == 0:
             return True
         if isinstance(consequent.items[0], KApply):
             return False
         return True
-    if consequent.arity == 0:
+    if not consequent or consequent.arity == 0:
         if isinstance(antecedent.items[0], KApply):
             return False
         return True