Constraints for storage types

package/version

@@ -1 +1 @@
-0.1.314
+0.1.316

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "kontrol"
-version = "0.1.314"
+version = "0.1.316"
 description = "Foundry integration for KEVM"
 authors = [
     "Runtime Verification, Inc. <[email protected]>",

src/kontrol/__init__.py

@@ -5,4 +5,4 @@
 if TYPE_CHECKING:
     from typing import Final
 
-VERSION: Final = '0.1.314'
+VERSION: Final = '0.1.316'

src/kontrol/prove.py

@@ -30,9 +30,11 @@
 from .foundry import Foundry, foundry_to_xml
 from .hevm import Hevm
 from .options import ConfigType, TraceOptions
-from .solc_to_k import Contract, hex_string_to_int
+from .solc_to_k import Contract, hex_string_to_int, _range_predicate
 from .utils import parse_test_version_tuple
 
+from .solc_to_k import StorageFieldMappingType
+
 if TYPE_CHECKING:
     from collections.abc import Iterable
     from typing import Final
@@ -551,7 +553,8 @@ def _method_to_initialized_cfg(
     kcfg.let_node(target_node_id, cterm=target_cterm)
 
     _LOGGER.info(f'Simplifying KCFG for test: {test.name}')
-    kcfg_explore.simplify(kcfg, {})
+#      kcfg_explore.simplify(kcfg, {})
+    print(kcfg.node(1).cterm)
 
     return kcfg, init_node_id, target_node_id
 
@@ -914,6 +917,7 @@ def _init_cterm(
 
     for constraint in storage_constraints:
         init_cterm = init_cterm.add_constraint(constraint)
+        print(constraint)
 
     # The calling contract is assumed to be in the present accounts for non-tests
     if not (config_type == ConfigType.TEST_CONFIG or active_symbolik):
@@ -982,6 +986,13 @@ def _create_cse_accounts(
     new_accounts.append(Foundry.symbolic_account(contract_name, contract_code, storage_map))
 
     for field in storage_fields:
+
+        for constraint in field.data_type.compute_constraints(storage_map):
+            new_account_constraints.append(constraint)
+            print(field.data_type.slot_vars())
+
+        if isinstance(field.data_type, StorageFieldMappingType):
+            ...
         if field.data_type == 'string':
             lookup = KEVM.lookup(storage_map, intToken(field.slot))
             length_byte_lt32 = ltInt(
@@ -1030,8 +1041,8 @@ def _create_cse_accounts(
             new_account_constraints.append(mlEqualsTrue(lowest_bit_not_set))
             new_account_constraints.append(mlEqualsTrue(length_byte_lt32))
             new_account_constraints.append(mlEqualsTrue(length_byte_positive))
-        if field.data_type.startswith('contract '):
-            contract_type = field.data_type.split(' ')[1]
+        if field.data_type.name.startswith('contract '):
+            contract_type = field.data_type.name.split(' ')[1]
             for full_contract_name, contract_obj in foundry.contracts.items():
                 # TODO: this is not enough, it is possible that the same contract comes with
                 # src% and test%, in which case we don't know automatically which one to choose

src/kontrol/solc_to_k.py

@@ -7,6 +7,8 @@
 from functools import cached_property
 from subprocess import CalledProcessError
 from typing import TYPE_CHECKING, NamedTuple
+from pyk.prelude.k import K_ITEM, GENERATED_TOP_CELL
+from pyk.prelude.ml import mlEqualsTrue
 
 from kevm_pyk.kevm import KEVM
 from pyk.kast.att import Atts, KAtt
@@ -322,13 +324,136 @@ def parse_devdoc(tag: str, devdoc: dict | None) -> dict:
     return natspecs
 
 
-class StorageField(NamedTuple):
+@dataclass
+class StorageFieldType:
+    name: str
+    slot: KInner
+
+    def compute_slots(self, base_slot: KInner) -> None:
+        self.slot = base_slot
+
+    def slots(self) -> list[tuple[KInner, str]]:
+        return [(self.slot, self.name)]
+
+    def slot_vars(self) -> list[KInner]:
+        return []
+
+    def compute_constraints(self, storage_map: KInner) -> list[KInner]:
+
+        constraints = []
+        for slot, slot_type in self.slots():
+            constraint = _range_predicate(term=KEVM.lookup(storage_map, slot.args[1]), type_label=slot_type)
+            if constraint is not None:
+                constraint = mlEqualsTrue(constraint)
+                for slot_var in self.slot_vars():
+                    constraint = KLabel('#Forall', K_ITEM, GENERATED_TOP_CELL)(slot_var, constraint)
+                constraints.append(constraint)
+        return constraints
+
+@dataclass
+class StorageFieldArrayType(StorageFieldType):
+    base_type: StorageFieldType
+
+
+@dataclass
+class StorageFieldMappingType(StorageFieldType):
+    key_type: StorageFieldType
+    val_type: StorageFieldType
+    slot_var: list[KInner] = ()
+
+    def slot_vars(self) -> list[KInner]:
+        return list(self.slot_var) + self.val_type.slot_vars()
+
+    def compute_slots(self, base_slot: KInner) -> None:
+        variable = KVariable('V_' + hash_str(base_slot)[0:8])
+        self.slot_var = (variable,)
+        self.slot = base_slot
+        self.val_type.compute_slots(
+#              KLabel('#Forall', K_ITEM, GENERATED_TOP_CELL)(variable,
+                KApply('buf', [intToken(32),
+                KApply('keccak(_)_SERIALIZATION_Int_Bytes', [
+                    KEVM.bytes_append(
+                        variable,
+                        base_slot,
+                    )
+                ])
+                ])
+#              )
+        )
+
+    def slots(self) -> list[tuple[KInner, str]]:
+        return [(self.slot, self.name)] + self.val_type.slots()
+
+@dataclass
+class StorageFieldStructType(StorageFieldType):
+    members: tuple[StorageField, ...]
+
+@dataclass
+class StorageField():
     label: str
-    data_type: str
+    data_type: StorageFieldType
     slot: int
     offset: int
 
 
+
+#  def constraints_for_storage_field(sf: StorageField, storage_map: KInner) -> list[KInner]:
+#      lookup = KEVM.lookup(storage_map, sf.data_type.get_slot(intToken(sf.slot)))
+#      term = sf.data_type.get_slot(KApply('asWord', [sf.slot]))
+#      print(term)
+#      range_pred = _range_predicate(type_label=sf.data_type.name, term=term)
+
+
+def process_data_type(dct: dict, types_dct: dict) -> StorageFieldType:
+    if 'key' in dct:
+        # Mapping
+        return StorageFieldMappingType(
+            name=dct['label'], 
+            key_type=process_data_type(types_dct[dct['key']], types_dct),
+            val_type=process_data_type(types_dct[dct['value']], types_dct),
+            slot=None,
+        )
+    elif 'members' in dct:
+        # Struct
+        return StorageFieldStructType(
+            name=dct['label'],
+            members=tuple(storage_field_from_dict(member, types_dct) for member in dct['members']),
+            slot=None,
+        )
+    elif 'base' in dct:
+        # Array
+        return StorageFieldArrayType(
+            name=dct['label'],
+            base_type=process_data_type(types_dct[dct['base']], types_dct),
+            slot=None,
+        )
+    else:
+        # Other
+        return StorageFieldType(name=dct['label'], slot=None)
+
+def storage_field_from_dict(dct: dict, types_dct: dict) -> StorageField:
+    label= dct['label']
+    slot = dct['slot']
+    offset = dct['offset']
+    data_type = process_data_type(types_dct[dct['type']], types_dct)
+    return StorageField(label, data_type, slot, offset)
+
+
+def process_storage_layout(storage_layout: dict) -> tuple[StorageField, ...]:
+    storage = storage_layout.get('storage', [])
+    types = storage_layout.get('types', {})
+
+    fields_list: list[StorageField] = []
+    for field in storage:
+        storage_field = storage_field_from_dict(field, types)
+        fields_list.append(storage_field)
+        storage_field.data_type.compute_slots(KApply('buf', [intToken(32), intToken(storage_field.slot)]))
+
+
+    return tuple(fields_list)
+
+
+
 @dataclass
 class Contract:
     @dataclass
@@ -1325,22 +1450,3 @@ def _find_function_calls(node: dict) -> None:
     return function_calls
 
 
-def process_storage_layout(storage_layout: dict) -> tuple[StorageField, ...]:
-    storage = storage_layout.get('storage', [])
-    types = storage_layout.get('types', {})
-
-    fields_list: list[StorageField] = []
-    for field in storage:
-        try:
-            type_info = types.get(field['type'], {})
-            storage_field = StorageField(
-                label=field['label'],
-                data_type=type_info.get('label', field['type']),
-                slot=int(field['slot']),
-                offset=int(field['offset']),
-            )
-            fields_list.append(storage_field)
-        except (KeyError, ValueError) as e:
-            _LOGGER.error(f'Error processing field {field}: {e}')
-
-    return tuple(fields_list)

src/tests/integration/test-data/foundry/test/Mapping.t.sol

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity =0.8.13;
+
+import "forge-std/Test.sol";
+import "kontrol-cheatcodes/KontrolCheats.sol";
+
+contract MappingContract {
+    mapping(address => mapping(address => uint256)) public balances;
+    mapping(address => uint256) public single_map;
+
+    function get_mapping_val(address a) public returns (uint256) {
+        return balances[a][a];
+    }
+
+    function get_mapping_val2(address a) public returns (uint256) {
+        return single_map[a];
+    }
+}
+
+contract MappingTest is Test {
+  uint256 val;
+  uint256 val2;
+  MappingContract c;
+
+  constructor() public {
+      c = new MappingContract();
+  }
+
+  function my_internal() internal { }
+
+  function test_mapping(address a) public {
+      val = c.get_mapping_val(a);
+      val2 = c.get_mapping_val2(a);
+      my_internal();
+      assert(val < 256);
+      assert(val2 < 256);
+  }
+
+}