Contract storage fields in CSE

src/kontrol/prove.py

@@ -21,7 +21,7 @@
 from pyk.prelude.collections import list_empty, map_empty, map_of, set_empty
 from pyk.prelude.k import GENERATED_TOP_CELL
 from pyk.prelude.kbool import FALSE, TRUE, notBool
-from pyk.prelude.kint import intToken
+from pyk.prelude.kint import eqInt, intToken
 from pyk.prelude.ml import mlEqualsFalse, mlEqualsTrue
 from pyk.prelude.string import stringToken
 from pyk.proof import ProofStatus
@@ -32,7 +32,7 @@
 from .cli import FoundryOptions, RpcOptions, TraceOptions
 from .foundry import Foundry, foundry_to_xml
 from .hevm import Hevm
-from .solc_to_k import Contract, hex_string_to_int
+from .solc_to_k import Contract, hex_string_to_int, type_constraint
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -43,6 +43,7 @@
     from pyk.kore.rpc import KoreServer
 
     from .deployment import DeploymentStateEntry
+    from .solc_to_k import StorageField
 
 _LOGGER: Final = logging.getLogger(__name__)
 
@@ -619,6 +620,7 @@ def _method_to_cfg(
         is_constructor=isinstance(method, Contract.Constructor),
         active_symbolik=active_symbolik,
         trace_options=trace_options,
+        storage_fields=contract.fields,
     )
     new_node_ids = []
 
@@ -781,6 +783,7 @@ def _init_cterm(
     callvalue: KInner | None = None,
     deployment_state_entries: Iterable[DeploymentStateEntry] | None = None,
     trace_options: TraceOptions | None = None,
+    storage_fields: tuple[StorageField, ...],
 ) -> CTerm:
     schedule = KApply('SHANGHAI_EVM')
 
@@ -800,7 +803,7 @@ def _init_cterm(
         'LOCALMEM_CELL': bytesToken(b''),
         'ACTIVE_CELL': FALSE,
         'MEMORYUSED_CELL': intToken(0),
-        'WORDSTACK_CELL': KApply('.WordStack_EVM-TYPES_WordStack'),
+        'WORDSTACK_CELL': KEVM.wordstack_empty(),
         'PC_CELL': intToken(0),
         'GAS_CELL': KEVM.inf_gas(KVariable('VGAS')),
         'K_CELL': KSequence([KEVM.sharp_execute(), KVariable('CONTINUATION')]),
@@ -819,9 +822,11 @@ def _init_cterm(
         'TRACEWORDSTACK_CELL': TRUE if trace_options.trace_wordstack else FALSE,
         'TRACEMEMORY_CELL': TRUE if trace_options.trace_memory else FALSE,
         'RECORDEDTRACE_CELL': FALSE,
-        'TRACEDATA_CELL': KApply('.List'),
+        'TRACEDATA_CELL': list_empty(),
     }
 
+    storage_map, storage_field_constraints = _create_initial_storage(storage_fields)
+
     if is_test or is_setup or is_constructor or active_symbolik:
         init_account_list = _create_initial_account_list(program, deployment_state_entries)
         init_subst_test = {
@@ -843,7 +848,7 @@ def _init_cterm(
         # Status: Currently, only the executing contract
         # TODO: Add all other accounts belonging to relevant contracts
         accounts: list[KInner] = [
-            Foundry.symbolic_account(Foundry.symbolic_contract_prefix(), program),
+            Foundry.symbolic_account(Foundry.symbolic_contract_prefix(), program, storage_map),
             KVariable('ACCOUNTS_REST', sort=KSort('AccountCellMap')),
         ]
 
@@ -867,7 +872,7 @@ def _init_cterm(
         # The address of the executing contract, the calling contract, and the origin contract
         # is always guaranteed to not be the address of the cheatcode contract
         init_cterm = init_cterm.add_constraint(
-            mlEqualsFalse(KApply('_==Int_', [KVariable(contract_id, sort=KSort('Int')), Foundry.address_CHEATCODE()]))
+            mlEqualsFalse(eqInt(KVariable(contract_id, sort=KSort('Int')), Foundry.address_CHEATCODE()))
         )
 
     # The calling contract is assumed to be in the present accounts for non-tests
@@ -880,12 +885,71 @@ def _init_cterm(
                 )
             )
         )
+        for c in storage_field_constraints:
+            init_cterm = init_cterm.add_constraint(c)
 
     init_cterm = KEVM.add_invariant(init_cterm)
 
     return init_cterm
 
 
+def _create_initial_storage(storage_fields: tuple[StorageField, ...]) -> tuple[KInner, list[KApply]]:
+    """
+    Creates the KAst representation of the storage layout for a given set of storage fields.
+
+    Args:
+        storage_fields (tuple[StorageField, ...]): A tuple of StorageField objects representing the contract's storage fields.
+
+    Returns:
+        tuple[KInner, list[KApply]]:
+            - A mapping of storage slots (IntToken) to their corresponding values (Bytes).
+            - A list of constraints for the symbolic variables used for each storage field.
+    """
+
+    storage_dict: dict[KInner, KInner] = {}
+    storage_constraints: list[KApply] = []
+    slot_data: KInner = intToken(0)
+    slot_length_left = 32
+    slot_index = None
+
+    for field in storage_fields:
+        if not (field.data_type in ['address', 'bool'] or field.data_type.startswith('uint')):
+            # skip not supported types
+            continue
+
+        symbolic_var = KVariable(f'SF_{field.label.upper()}', sort=KSort('Int'))
+        storage_constraints.append(type_constraint(field.data_type, symbolic_var))
+
+        if slot_index is None:
+            # initialize for the first field
+            slot_index = field.slot
+            slot_data = KEVM.buf(field.length, symbolic_var)
+            slot_length_left -= field.length
+        elif slot_index == field.slot:
+            # handle fields in the same slot but with different offsets
+            slot_data = KEVM.bytes_append(slot_data, KEVM.buf(field.length, symbolic_var))
+            slot_length_left -= field.length
+        else:
+            # when the slot value changes, save the previous data in the dict and compute the current one
+            if slot_length_left > 0:
+                slot_data = KEVM.bytes_append(slot_data, KEVM.buf(slot_length_left, intToken(0)))
+
+            storage_dict[intToken(slot_index)] = KApply('#asWord(_)_EVM-TYPES_Int_Bytes', slot_data)
+
+            # Reset for the new slot
+            slot_index = field.slot
+            slot_data = KEVM.buf(field.length, symbolic_var)
+            slot_length_left = 32 - field.length
+
+    # Ensure the last slot data is saved
+    if slot_index is not None:
+        if slot_length_left > 0:
+            slot_data = KEVM.bytes_append(slot_data, KEVM.buf(slot_length_left, intToken(0)))
+        storage_dict[intToken(slot_index)] = KApply('#asWord(_)_EVM-TYPES_Int_Bytes', slot_data)
+
+    return map_of(storage_dict), storage_constraints
+
+
 def _create_initial_account_list(
     program: KInner, deployment_state: Iterable[DeploymentStateEntry] | None
 ) -> list[KInner]:

src/kontrol/solc_to_k.py

@@ -18,6 +18,7 @@
 from pyk.kdist import kdist
 from pyk.prelude.kbool import TRUE, andBool
 from pyk.prelude.kint import eqInt, intToken
+from pyk.prelude.ml import mlEqualsTrue
 from pyk.prelude.string import stringToken
 from pyk.utils import hash_str, run_process, single
 
@@ -334,6 +335,7 @@ class StorageField(NamedTuple):
     data_type: str
     slot: int
     offset: int
+    length: int
 
 
 @dataclass
@@ -1289,9 +1291,25 @@ def process_storage_layout(storage_layout: dict) -> tuple[StorageField, ...]:
                 data_type=type_info.get('label', field['type']),
                 slot=int(field['slot']),
                 offset=int(field['offset']),
+                length=int(type_info.get('numberOfBytes', 32)),
             )
             fields_list.append(storage_field)
         except (KeyError, ValueError) as e:
             _LOGGER.error(f'Error processing field {field}: {e}')
 
     return tuple(fields_list)
+
+
+def type_constraint(data_type: str, symbolic_var: KVariable) -> KApply:
+    # TODO: fill with rest of the cases
+    data_type = data_type.lower()
+    if data_type.startswith('uint'):
+        return mlEqualsTrue(KEVM.range_uint(int(data_type[4:]), symbolic_var))
+
+    match data_type:
+        case 'address':
+            return mlEqualsTrue(KEVM.range_address(symbolic_var))
+        case 'bool':
+            return mlEqualsTrue(KEVM.range_bool(symbolic_var))
+        case _:
+            raise ValueError(f'Unimplemented: {data_type}')

src/tests/integration/test-data/show/AddConst.applyOp(uint256).cse.expected

@@ -6,7 +6,7 @@
 │   statusCode: STATUSCODE:StatusCode
 ┃
 ┃ (branch)
-┣━━┓ constraint: VV0_x_114b9705:Int <=Int ( maxUInt256 -Int #lookup ( CONTRACT_STORAGE:Map , 0 ) )
+┣━━┓ constraint: VV0_x_114b9705:Int <=Int ( maxUInt256 -Int SF_C:Int )
 ┃  │
 ┃  ├─ 8
 ┃  │   k: #execute ~> CONTINUATION:K
@@ -29,7 +29,7 @@
 ┃      callDepth: CALLDEPTH_CELL_5d410f2a:Int
 ┃      statusCode: STATUSCODE_FINAL:StatusCode
 ┃
-┗━━┓ constraint: ( maxUInt256 -Int #lookup ( CONTRACT_STORAGE:Map , 0 ) ) <Int VV0_x_114b9705:Int
+┗━━┓ constraint: ( maxUInt256 -Int SF_C:Int ) <Int VV0_x_114b9705:Int
    │
    ├─ 9
    │   k: #execute ~> CONTINUATION:K
@@ -76,7 +76,7 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
              <ethereum>
                <evm>
                  <output>
-                   ( _OUTPUT_CELL => #buf ( 32 , ( VV0_x_114b9705:Int +Int #lookup ( CONTRACT_STORAGE:Map , 0 ) ) ) )
+                   ( _OUTPUT_CELL => #buf ( 32 , ( VV0_x_114b9705:Int +Int SF_C:Int ) ) )
                  </output>
                  <statusCode>
                    ( _STATUSCODE => EVMC_SUCCESS )
@@ -98,7 +98,7 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
                      ( .WordStack => ( selector ( "applyOp(uint256)" ) : .WordStack ) )
                    </wordStack>
                    <localMem>
-                     ( b"" => b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" +Bytes #buf ( 32 , ( VV0_x_114b9705:Int +Int #lookup ( CONTRACT_STORAGE:Map , 0 ) ) ) )
+                     ( b"" => b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" +Bytes #buf ( 32 , ( VV0_x_114b9705:Int +Int SF_C:Int ) ) )
                    </localMem>
                    <memoryUsed>
                      0
@@ -138,7 +138,7 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
                        CONTRACT_BAL:Int
                      </balance>
                      <storage>
-                       CONTRACT_STORAGE:Map
+                       ( 0 |-> SF_C:Int )
                      </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
@@ -222,7 +222,8 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
              </traceData>
            </KEVMTracing>
          </foundry>
-      requires ( 0 <=Int CALLER_ID:Int
+      requires ( 0 <=Int SF_C:Int
+       andBool ( 0 <=Int CALLER_ID:Int
        andBool ( 0 <=Int ORIGIN_ID:Int
        andBool ( 0 <=Int CONTRACT_ID:Int
        andBool ( 0 <=Int NUMBER_CELL:Int
@@ -237,18 +238,19 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
        andBool ( CALLER_ID:Int <Int pow160
        andBool ( ORIGIN_ID:Int <Int pow160
        andBool ( CONTRACT_ID:Int <Int pow160
+       andBool ( SF_C:Int <Int pow256
        andBool ( NUMBER_CELL:Int <=Int maxSInt256
        andBool ( CONTRACT_BAL:Int <Int pow256
        andBool ( TIMESTAMP_CELL:Int <Int pow256
        andBool ( VV0_x_114b9705:Int <Int pow256
        andBool ( ( notBool <acctID>
           CONTRACT_ID:Int
         </acctID> in_keys ( ACCOUNTS_REST:AccountCellMap ) )
+       andBool ( VV0_x_114b9705:Int <=Int ( maxUInt256 -Int SF_C:Int )
        andBool ( ( notBool #range ( 0 < CALLER_ID:Int <= 9 ) )
        andBool ( ( notBool #range ( 0 < ORIGIN_ID:Int <= 9 ) )
        andBool ( ( notBool #range ( 0 < CONTRACT_ID:Int <= 9 ) )
-       andBool ( VV0_x_114b9705:Int <=Int ( maxUInt256 -Int #lookup ( CONTRACT_STORAGE:Map , 0 ) )
-               ))))))))))))))))))))))))
+               ))))))))))))))))))))))))))
       [priority(20), label(BASIC-BLOCK-8-TO-6)]
 
     rule [BASIC-BLOCK-9-TO-7]: <foundry>
@@ -288,7 +290,7 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
                      0
                    </callValue>
                    <wordStack>
-                     ( .WordStack => ( 0 : ( VV0_x_114b9705:Int : ( #lookup ( CONTRACT_STORAGE:Map , 0 ) : ( 118 : ( 0 : ( VV0_x_114b9705:Int : ( 70 : ( selector ( "applyOp(uint256)" ) : .WordStack ) ) ) ) ) ) ) ) )
+                     ( .WordStack => ( 0 : ( VV0_x_114b9705:Int : ( SF_C:Int : ( 118 : ( 0 : ( VV0_x_114b9705:Int : ( 70 : ( selector ( "applyOp(uint256)" ) : .WordStack ) ) ) ) ) ) ) ) )
                    </wordStack>
                    <localMem>
                      ( b"" => b"NH{q\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x11\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x80" )
@@ -331,7 +333,7 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
                        CONTRACT_BAL:Int
                      </balance>
                      <storage>
-                       CONTRACT_STORAGE:Map
+                       ( 0 |-> SF_C:Int )
                      </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
@@ -415,7 +417,8 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
              </traceData>
            </KEVMTracing>
          </foundry>
-      requires ( 0 <=Int CALLER_ID:Int
+      requires ( 0 <=Int SF_C:Int
+       andBool ( 0 <=Int CALLER_ID:Int
        andBool ( 0 <=Int ORIGIN_ID:Int
        andBool ( 0 <=Int CONTRACT_ID:Int
        andBool ( 0 <=Int NUMBER_CELL:Int
@@ -430,18 +433,19 @@ module SUMMARY-SRC%CSE%ADDCONST.APPLYOP(UINT256):0
        andBool ( CALLER_ID:Int <Int pow160
        andBool ( ORIGIN_ID:Int <Int pow160
        andBool ( CONTRACT_ID:Int <Int pow160
+       andBool ( SF_C:Int <Int pow256
        andBool ( NUMBER_CELL:Int <=Int maxSInt256
        andBool ( CONTRACT_BAL:Int <Int pow256
        andBool ( TIMESTAMP_CELL:Int <Int pow256
        andBool ( VV0_x_114b9705:Int <Int pow256
        andBool ( ( notBool <acctID>
           CONTRACT_ID:Int
         </acctID> in_keys ( ACCOUNTS_REST:AccountCellMap ) )
+       andBool ( ( maxUInt256 -Int SF_C:Int ) <Int VV0_x_114b9705:Int
        andBool ( ( notBool #range ( 0 < CALLER_ID:Int <= 9 ) )
        andBool ( ( notBool #range ( 0 < ORIGIN_ID:Int <= 9 ) )
        andBool ( ( notBool #range ( 0 < CONTRACT_ID:Int <= 9 ) )
-       andBool ( ( maxUInt256 -Int #lookup ( CONTRACT_STORAGE:Map , 0 ) ) <Int VV0_x_114b9705:Int
-               ))))))))))))))))))))))))
+               ))))))))))))))))))))))))))
       [priority(20), label(BASIC-BLOCK-9-TO-7)]
 
 endmodule

src/tests/integration/test-data/show/ArithmeticContract.add(uint256,uint256).cse.expected

@@ -137,6 +137,9 @@ module SUMMARY-SRC%ARITHMETICCONTRACT.ADD(UINT256,UINT256):0
                      <balance>
                        CONTRACT_BAL:Int
                      </balance>
+                     <storage>
+                       .Map
+                     </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
                      </nonce>
@@ -329,6 +332,9 @@ module SUMMARY-SRC%ARITHMETICCONTRACT.ADD(UINT256,UINT256):0
                      <balance>
                        CONTRACT_BAL:Int
                      </balance>
+                     <storage>
+                       .Map
+                     </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
                      </nonce>

src/tests/integration/test-data/show/ArithmeticContract.add_sub_external(uint256,uint256,uint256).cse.expected

@@ -197,6 +197,9 @@ module SUMMARY-SRC%ARITHMETICCONTRACT.ADD-SUB-EXTERNAL(UINT256,UINT256,UINT256):
                      <balance>
                        CONTRACT_BAL:Int
                      </balance>
+                     <storage>
+                       .Map
+                     </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
                      </nonce>
@@ -400,6 +403,9 @@ module SUMMARY-SRC%ARITHMETICCONTRACT.ADD-SUB-EXTERNAL(UINT256,UINT256,UINT256):
                      <balance>
                        CONTRACT_BAL:Int
                      </balance>
+                     <storage>
+                       .Map
+                     </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
                      </nonce>
@@ -604,6 +610,9 @@ module SUMMARY-SRC%ARITHMETICCONTRACT.ADD-SUB-EXTERNAL(UINT256,UINT256,UINT256):
                      <balance>
                        CONTRACT_BAL:Int
                      </balance>
+                     <storage>
+                       .Map
+                     </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
                      </nonce>
@@ -809,6 +818,9 @@ module SUMMARY-SRC%ARITHMETICCONTRACT.ADD-SUB-EXTERNAL(UINT256,UINT256,UINT256):
                      <balance>
                        CONTRACT_BAL:Int
                      </balance>
+                     <storage>
+                       .Map
+                     </storage>
                      <nonce>
                        CONTRACT_NONCE:Int
                      </nonce>