CSE: Include symbolic contracts corresponding to contract fields into accounts

package/version

@@ -1 +1 @@
-0.1.301
+0.1.302

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "kontrol"
-version = "0.1.301"
+version = "0.1.302"
 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.301'
+VERSION: Final = '0.1.302'

src/kontrol/prove.py

@@ -42,6 +42,7 @@
 
     from .deployment import DeploymentStateEntry
     from .options import ProveOptions
+    from .solc_to_k import StorageField
 
 _LOGGER: Final = logging.getLogger(__name__)
 
@@ -506,6 +507,7 @@ def _method_to_initialized_cfg(
 
     empty_config = foundry.kevm.definition.empty_config(GENERATED_TOP_CELL)
     kcfg, new_node_ids, init_node_id, target_node_id = _method_to_cfg(
+        foundry,
         empty_config,
         test.contract,
         test.method,
@@ -541,6 +543,7 @@ def _method_to_initialized_cfg(
 
 
 def _method_to_cfg(
+    foundry: Foundry,
     empty_config: KInner,
     contract: Contract,
     method: Contract.Method | Contract.Constructor,
@@ -567,9 +570,11 @@ def _method_to_cfg(
         program = contract_code
 
     init_cterm = _init_cterm(
+        foundry,
         empty_config,
         program=program,
         contract_code=bytesToken(contract_code),
+        storage_fields=contract.fields,
         use_gas=use_gas,
         deployment_state_entries=deployment_state_entries,
         calldata=calldata,
@@ -759,9 +764,11 @@ def _init_account(address: int) -> None:
 
 
 def _init_cterm(
+    foundry: Foundry,
     empty_config: KInner,
     program: bytes,
     contract_code: KInner,
+    storage_fields: tuple[StorageField, ...],
     use_gas: bool,
     config_type: ConfigType,
     active_symbolik: bool,
@@ -813,6 +820,8 @@ def _init_cterm(
         'TRACEDATA_CELL': KApply('.List'),
     }
 
+    storage_constraints = []
+
     if config_type == ConfigType.TEST_CONFIG or active_symbolik:
         init_account_list = _create_initial_account_list(contract_code, deployment_state_entries)
         init_subst_test = {
@@ -832,11 +841,21 @@ def _init_cterm(
     else:
         # Symbolic accounts of all relevant contracts
         # 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(), contract_code),
-            KVariable('ACCOUNTS_REST', sort=KSort('AccountCellMap')),
-        ]
+        # TODO: Add all other accounts belonging to relevant contracts``
+        accounts: list[KInner] = []
+
+        # TODO(palina): executed contract's storage
+        storage_map = KVariable(Foundry.symbolic_contract_prefix() + '_STORAGE', sort=KSort('Map'))
+        # TODO(palina): executed contract's name
+        contract_name = Foundry.symbolic_contract_prefix() + '_ID'
+        new_accounts, storage_constraints, storage_map = _create_cse_accounts(
+            foundry, storage_fields, contract_name, storage_map
+        )
+
+        accounts.extend(new_accounts)
+
+        accounts.append(Foundry.symbolic_account(Foundry.symbolic_contract_prefix(), program, storage_map))
+        accounts.append(KVariable('ACCOUNTS_REST', sort=KSort('AccountCellMap')))
 
         init_subst_accounts = {'ACCOUNTS_CELL': KEVM.accounts(accounts)}
         init_subst.update(init_subst_accounts)
@@ -865,6 +884,10 @@ def _init_cterm(
             mlEqualsFalse(KApply('_==Int_', [KVariable(contract_id, sort=KSort('Int')), Foundry.address_CHEATCODE()]))
         )
 
+    # TODO(palina): make sure these constraints are getting applied
+    for constraint in storage_constraints:
+        init_cterm.add_constraint(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):
         init_cterm.add_constraint(
@@ -902,6 +925,67 @@ def _create_initial_account_list(
     return init_account_list
 
 
+def _create_cse_accounts(
+    foundry: Foundry, storage_fields: tuple[StorageField, ...], contract_name: str, contract_storage: KApply
+) -> tuple[list[KApply], list[KApply], KApply]:
+    # TODO: call this function recursively
+    new_accounts = []
+    new_account_constraints = []
+
+    for field in storage_fields:
+        if field.data_type.startswith('contract '):
+            contract_type = field.data_type.split(' ')[1]
+            for contract_name, contract_obj in foundry.contracts.items():
+                if contract_name.split('%')[-1] == contract_type:
+                    contract_code = KEVM.bin_runtime(KApply(f'contract_{contract_obj.name_with_path}'))
+                    contract_account_name = 'CONTRACT-' + field.label.upper()
+                    new_account = Foundry.symbolic_account(contract_account_name, contract_code)
+                    new_accounts.append(new_account)
+
+                    # TODO: handle the case where the field has a non-zero offset
+                    # maxUInt160 &Int #lookup ( CONTRACT_STORAGE:Map , 0 )
+                    contract_storage_lookup = KApply(
+                        '_&Int_',
+                        [
+                            intToken(1461501637330902918203684832716283019655932542975),
+                            KEVM.lookup(contract_storage, intToken(field.slot)),
+                        ],
+                    )
+                    new_account_constraints.append(
+                        mlEqualsTrue(
+                            KApply('_==Int_', [contract_storage_lookup, KVariable(contract_account_name + '_ID')])
+                        )
+                    )
+
+                    # New contract account is not the cheatcode contract
+                    new_account_constraints.append(
+                        mlEqualsFalse(
+                            KApply(
+                                '_==Int_',
+                                [
+                                    KVariable(contract_account_name + '_ID', sort=KSort('Int')),
+                                    Foundry.address_CHEATCODE(),
+                                ],
+                            )
+                        )
+                    )
+
+                    # New contract account is also not the symbolic contract account being executed
+                    new_account_constraints.append(
+                        mlEqualsFalse(
+                            KApply(
+                                '_==Int_',
+                                [
+                                    KVariable(contract_account_name + '_ID', sort=KSort('Int')),
+                                    KVariable(Foundry.symbolic_contract_prefix(), sort=KSort('Int')),
+                                ],
+                            )
+                        )
+                    )
+
+    return new_accounts, new_account_constraints, contract_storage
+
+
 def _final_cterm(
     empty_config: KInner,
     program: KInner,

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

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity =0.8.13;
+
+contract Token {
+    uint256 public totalSupply;
+
+    constructor(uint256 _totalSupply) {
+        totalSupply = _totalSupply;
+    }
+}
+
+contract Escrow {
+    Token token;
+    uint256 tokenTotalSupply;
+
+    function getTokenTotalSupply() public {
+        uint256 tokenTS = token.totalSupply();
+        tokenTotalSupply = tokenTS + 15;
+    }
+}
+
+contract ContractFieldTest {
+    Escrow escrow; 
+
+    function setUp() public {
+        escrow = new Escrow();
+    }
+
+    /* Calling `getTokenTotalSupply` will summarize `totalSupply` and
+       include `Token token` into the list of accounts in `getTokenTotalSupply`'s summary
+    */
+    function getEscrowToken() public {
+        escrow.getTokenTotalSupply();
+    }
+}