Add caching to WasmKrunInitializer (#19)

* update mx-semantics

* rename imported K modules `foundry => kasmer`

* add cache

* cache `current_module`

* update mx-semantics

kmxwasm/src/kmxwasm/property_testing/wasm_krun_initializer.py

@@ -1,5 +1,5 @@
 from pyk.cterm import CTerm
-from pyk.kast.inner import KInner, KSequence
+from pyk.kast.inner import KApply, KInner, KSequence, KToken
 from pyk.kcfg import KCFG
 from pyk.prelude.collections import list_of, map_empty
 from pyk.prelude.utils import token
@@ -33,13 +33,39 @@ class WasmKrunInitializer:
     def __init__(self, tools: Tools):
         self.__tools = tools
         self.__first_wasm_cell: KInner | None = None
+        self.__cache: dict[str, tuple[KInner, KInner]] = {}  # address -> (<wasm>, <contractModIdx>)
 
     def initialize(self, kcfg: KCFG, start_node: KCFG.Node) -> None:
         start_cell = start_node.cterm.config
         commands_contents = commands_cell_contents(start_cell)
         assert len(commands_contents) > 0
         first = commands_contents.items[0]
 
+        assert isinstance(first, KApply)
+        assert isinstance(first.args[0], KToken)
+        acct_address = first.args[0].token
+
+        if acct_address in self.__cache:
+            print('*' * 30, 'Reading WASM from cache:', acct_address)
+            wasm_cell, current_module = self.__cache[acct_address]
+        else:
+            krun_cell = self.__prepare_krun_cell(start_cell, first)
+            print('*' * 30, 'Initializing WASM:', acct_address)
+            krun_result = self.__tools.krun(krun_cell)
+
+            wasm_cell = get_wasm_cell(krun_result)
+            current_module = get_contract_mod_idx_cell(krun_result)
+            self.__cache[acct_address] = (wasm_cell, current_module)
+
+        final_cell = replace_wasm_cell(start_cell, wasm_cell)
+        final_cell = replace_contract_mod_idx_cell(final_cell, current_module)
+        final_cell = set_commands_cell_contents(final_cell, KSequence(commands_contents.items[1:]))
+
+        final_cterm = CTerm(final_cell, start_node.cterm.constraints)
+        final_node = kcfg.create_node(final_cterm)
+        kcfg.create_edge(source_id=start_node.id, target_id=final_node.id, depth=1)
+
+    def __prepare_krun_cell(self, start_cell: KInner, first: KInner) -> KInner:
         krun_cell = set_call_stack_cell_content(start_cell, list_of([]))
         if not self.__first_wasm_cell:
             # This must run after set_call_stack_cell_content because the call
@@ -59,16 +85,4 @@ def initialize(self, kcfg: KCFG, start_node: KCFG.Node) -> None:
 
         krun_cell = replace_wasm_cell(krun_cell, self.__first_wasm_cell)
 
-        # TODO: Figure out if it's possible to cache the initialization result.
-        print('*' * 30, 'Initializing WASM.')
-        krun_result = self.__tools.krun(krun_cell)
-
-        wasm_cell = get_wasm_cell(krun_result)
-        final_cell = replace_wasm_cell(start_cell, wasm_cell)
-        current_module = get_contract_mod_idx_cell(krun_result)
-        final_cell = replace_contract_mod_idx_cell(final_cell, current_module)
-        final_cell = set_commands_cell_contents(final_cell, KSequence(commands_contents.items[1:]))
-
-        final_cterm = CTerm(final_cell, start_node.cterm.constraints)
-        final_node = kcfg.create_node(final_cterm)
-        kcfg.create_edge(source_id=start_node.id, target_id=final_node.id, depth=1)
+        return krun_cell