feat: support for snapshotStorage(address) cheatcode

examples/simple/test/SimpleState.t.sol

@@ -38,10 +38,28 @@ contract SimpleStateTest is SymTest, Test {
 
         // note: a total of 253 feasible paths are generated, of which only 10 unique states exist
         for (uint i = 0; i < 10; i++) {
-            (success,) = address(target).call(svm.createCalldata("SimpleState"));
+            (success,) = address(target).call(svm.createCalldata("SimpleState")); // excluding view functions
             vm.assume(success);
         }
 
         assertFalse(target.buggy());
     }
+
+    function check_buggy_with_snapshot() public {
+        bool success;
+
+        // take the initial storage snapshot
+        uint id = svm.snapshotStorage(address(target));
+
+        // note: a total of 253 feasible paths are generated, of which only 10 unique states exist
+        for (uint i = 0; i < 10; i++) {
+            (success,) = address(target).call(svm.createCalldata("SimpleState", true)); // including view functions
+            vm.assume(success);
+            uint nid = svm.snapshotStorage(address(target));
+            vm.assume(nid > id); // ignore if no state changes
+            id = nid;
+        }
+
+        assertFalse(target.buggy());
+    }
 }

src/halmos/cheatcodes.py

@@ -230,6 +230,19 @@ def symbolic_storage(ex, arg, sevm, stack, step_id):
     return ByteVec()  # empty return data
 
 
+def snapshot_storage(ex, arg, sevm, stack, step_id):
+    account = uint160(arg.get_word(4))
+    account_alias = sevm.resolve_address_alias(
+        ex, account, stack, step_id, allow_branching=False
+    )
+
+    if account_alias is None:
+        error_msg = f"snapshotStorage() is not allowed for a nonexistent account: {hexify(account)}"
+        raise HalmosException(error_msg)
+
+    return ByteVec(uint256(ex.storage[account_alias].cnt))
+
+
 def create_calldata_contract(ex, arg, sevm, stack, step_id):
     contract_name = name_of(extract_string_argument(arg, 0))
     return create_calldata_generic(ex, sevm, contract_name)
@@ -447,6 +460,7 @@ class halmos_cheat_code:
         0x3B0FA01B: create_address,  # createAddress(string)
         0x6E0BB659: create_bool,  # createBool(string)
         0xDC00BA4D: symbolic_storage,  # enableSymbolicStorage(address)
+        0x5DBB8438: snapshot_storage,  # snapshotStorage(address)
         0xBE92D5A2: create_calldata_contract,  # createCalldata(string)
         0xDEEF391B: create_calldata_contract_bool,  # createCalldata(string,bool)
         0x88298B32: create_calldata_file_contract,  # createCalldata(string,string)

src/halmos/sevm.py

@@ -833,6 +833,17 @@ def extend_path(self, path):
 class StorageData:
     symbolic: bool = False
     mapping: dict = field(default_factory=dict)
+    cnt: int = 0
+
+    def __getitem__(self, key) -> ArrayRef | BitVecRef:
+        return self.mapping[key]
+
+    def __setitem__(self, key, value) -> None:
+        self.mapping[key] = value
+        self.cnt += 1
+
+    def __contains__(self, key) -> bool:
+        return key in self.mapping
 
 
 class Exec:  # an execution path
@@ -1327,26 +1338,25 @@ def init(
         """
         assert_address(addr)
 
-        storage_data = ex.storage[addr]
-        mapping = storage_data.mapping[slot][num_keys]
+        storage = ex.storage[addr]
 
-        if size_keys in mapping:
+        if (slot, num_keys, size_keys) in storage:
             return
 
         if size_keys > 0:
             # do not use z3 const array `K(BitVecSort(size_keys), ZERO)` when not ex.symbolic
             # instead use normal smt array, and generate emptyness axiom; see load()
-            mapping[size_keys] = cls.empty(addr, slot, keys)
+            storage[slot, num_keys, size_keys] = cls.empty(addr, slot, keys)
             return
 
         # size_keys == 0
-        mapping[size_keys] = (
+        storage[slot, num_keys, size_keys] = (
             BitVec(
                 # note: uuid is excluded to be deterministic
                 f"storage_{id_str(addr)}_{slot}_{num_keys}_{size_keys}_00",
                 BitVecSort256,
             )
-            if storage_data.symbolic
+            if storage.symbolic
             else ZERO
         )
 
@@ -1356,44 +1366,43 @@ def load(cls, ex: Exec, addr: Any, loc: Word) -> Word:
 
         cls.init(ex, addr, slot, keys, num_keys, size_keys)
 
-        storage_data = ex.storage[addr]
-        mapping = storage_data.mapping[slot][num_keys]
+        storage = ex.storage[addr]
+        storage_chunk = storage[slot, num_keys, size_keys]
 
         if num_keys == 0:
-            return mapping[size_keys]
+            return storage_chunk
 
-        symbolic = storage_data.symbolic
+        symbolic = storage.symbolic
         concat_keys = concat(keys)
 
         if not symbolic:
             # generate emptyness axiom for each array index, instead of using quantified formula; see init()
             default_value = Select(cls.empty(addr, slot, keys), concat_keys)
             ex.path.append(default_value == ZERO)
 
-        return ex.select(mapping[size_keys], concat_keys, ex.storages, symbolic)
+        return ex.select(storage_chunk, concat_keys, ex.storages, symbolic)
 
     @classmethod
     def store(cls, ex: Exec, addr: Any, loc: Any, val: Any) -> None:
         (slot, keys, num_keys, size_keys) = cls.get_key_structure(ex, loc)
 
         cls.init(ex, addr, slot, keys, num_keys, size_keys)
 
-        storage_data = ex.storage[addr]
-        mapping = storage_data.mapping[slot][num_keys]
+        storage = ex.storage[addr]
 
         if num_keys == 0:
-            mapping[size_keys] = val
+            storage[slot, num_keys, size_keys] = val
             return
 
         new_storage_var = Array(
             f"storage_{id_str(addr)}_{slot}_{num_keys}_{size_keys}_{uid()}_{1+len(ex.storages):>02}",
             BitVecSorts[size_keys],
             BitVecSort256,
         )
-        new_storage = Store(mapping[size_keys], concat(keys), val)
+        new_storage = Store(storage[slot, num_keys, size_keys], concat(keys), val)
         ex.path.append(new_storage_var == new_storage)
 
-        mapping[size_keys] = new_storage_var
+        storage[slot, num_keys, size_keys] = new_storage_var
         ex.storages[new_storage_var] = new_storage
 
     @classmethod
@@ -1497,10 +1506,10 @@ def init(cls, ex: Exec, addr: Any, loc: BitVecRef, size_keys: int) -> None:
         """
         assert_address(addr)
 
-        mapping = ex.storage[addr].mapping
+        storage = ex.storage[addr]
 
-        if size_keys not in mapping:
-            mapping[size_keys] = cls.empty(addr, loc)
+        if size_keys not in storage:
+            storage[size_keys] = cls.empty(addr, loc)
 
     @classmethod
     def load(cls, ex: Exec, addr: Any, loc: Word) -> Word:
@@ -1509,16 +1518,15 @@ def load(cls, ex: Exec, addr: Any, loc: Word) -> Word:
 
         cls.init(ex, addr, loc, size_keys)
 
-        storage_data = ex.storage[addr]
-        mapping = storage_data.mapping
-        symbolic = storage_data.symbolic
+        storage = ex.storage[addr]
+        symbolic = storage.symbolic
 
         if not symbolic:
             # generate emptyness axiom for each array index, instead of using quantified formula; see init()
             default_value = Select(cls.empty(addr, loc), loc)
             ex.path.append(default_value == ZERO)
 
-        return ex.select(mapping[size_keys], loc, ex.storages, symbolic)
+        return ex.select(storage[size_keys], loc, ex.storages, symbolic)
 
     @classmethod
     def store(cls, ex: Exec, addr: Any, loc: Any, val: Any) -> None:
@@ -1527,17 +1535,17 @@ def store(cls, ex: Exec, addr: Any, loc: Any, val: Any) -> None:
 
         cls.init(ex, addr, loc, size_keys)
 
-        mapping = ex.storage[addr].mapping
+        storage = ex.storage[addr]
 
         new_storage_var = Array(
             f"storage_{id_str(addr)}_{size_keys}_{uid()}_{1+len(ex.storages):>02}",
             BitVecSorts[size_keys],
             BitVecSort256,
         )
-        new_storage = Store(mapping[size_keys], loc, val)
+        new_storage = Store(storage[size_keys], loc, val)
         ex.path.append(new_storage_var == new_storage)
 
-        mapping[size_keys] = new_storage_var
+        storage[size_keys] = new_storage_var
         ex.storages[new_storage_var] = new_storage
 
     @classmethod

tests/expected/simple.json

@@ -234,6 +234,15 @@
                 "num_paths": null,
                 "time": null,
                 "num_bounded_loops": null
+            },
+            {
+                "name": "check_buggy_with_snapshot()",
+                "exitcode": 1,
+                "num_models": 1,
+                "models": null,
+                "num_paths": null,
+                "time": null,
+                "num_bounded_loops": null
             }
         ],
         "test/TotalPrice.t.sol:TotalPriceTest": [