Record Solidity/K mapping for input names

src/kontrol/foundry.py

@@ -21,7 +21,15 @@
 from kevm_pyk.utils import byte_offset_to_lines, legacy_explore, print_failure_info, print_model
 from pyk.cterm import CTerm
 from pyk.kast.inner import KApply, KInner, KSort, KToken, KVariable
-from pyk.kast.manip import cell_label_to_var_name, collect, extract_lhs, flatten_label, minimize_term, top_down
+from pyk.kast.manip import (
+    abstract_term_safely,
+    cell_label_to_var_name,
+    collect,
+    extract_lhs,
+    flatten_label,
+    minimize_term,
+    top_down,
+)
 from pyk.kast.outer import KDefinition, KFlatModule, KImport, KRequire
 from pyk.kcfg import KCFG
 from pyk.prelude.bytes import bytesToken
@@ -211,6 +219,10 @@ def main_file(self) -> Path:
     def contracts_file(self) -> Path:
         return self.kompiled / 'contracts.k'
 
+    @property
+    def input_mapping_file(self) -> Path:
+        return self.kompiled / 'input-mapping.json'
+
     @cached_property
     def kevm(self) -> KEVM:
         use_directory = self.out / 'tmp'
@@ -385,6 +397,41 @@ def build(self, no_metadata: bool) -> None:
         except CalledProcessError as err:
             raise RuntimeError("Couldn't forge build!") from err
 
+    def record_input_name_mapping(self, input_mapping_file: Path) -> None:
+        """
+        Saves a JSON file storing a mapping between the names of Method's inputs and their K representation.
+        """
+
+        input_mapping = {}
+
+        # Create the mapping for each contract's method
+        for contract in self.contracts.values():
+            contract_mapping = {}
+
+            # Create the mapping for each method's arguments
+            for method in contract.methods:
+                method_input_mapping = {
+                    abstract_term_safely(KVariable('_###SOLIDITY_ARG_VAR###_'), base_name=f'V{arg_name}').name: name
+                    for arg_name, name in method.arg_names.items()
+                }
+
+                contract_mapping[method.signature] = method_input_mapping
+
+            input_mapping[contract.name_with_path] = contract_mapping
+
+        # Add environment variables to the mapping
+        env_variables_mapping = {
+            'TIMESTAMP_CELL': 'block.timestamp',
+            'NUMBER_CELL': 'block.number',
+            'ORIGIN_ID': 'tx.origin',
+            'CALLER_ID': 'msg.sender',
+        }
+
+        input_mapping['env'] = env_variables_mapping
+
+        # Write resulting mapping to JSON file
+        input_mapping_file.write_text(json.dumps(input_mapping, indent=4))
+
     @cached_property
     def all_tests(self) -> list[str]:
         test_dir = os.path.join(self.profile.get('test', 'test'), '')
@@ -1289,7 +1336,13 @@ def foundry_get_model(
             res_lines.append('')
             res_lines.append(f'Node id: {node_id}')
             node = proof.kcfg.node(node_id)
-            res_lines.extend(print_model(node, kcfg_explore))
+
+            contract_name, test_name = test_id.split(':')[0].split('.')
+            input_mapping = json.loads(foundry.input_mapping_file.read_text())
+            proof_input_mapping = input_mapping.get(contract_name, {}).get(test_name, {})
+            proof_input_mapping.update(input_mapping.get('env', {}))
+
+            res_lines.extend(print_model(node, kcfg_explore, proof_input_mapping))
 
     return '\n'.join(res_lines)
 

src/kontrol/kompile.py

@@ -132,6 +132,8 @@ def foundry_kompile(
         _LOGGER.info(f'Wrote file: {foundry_contracts_file}')
         foundry.main_file.write_text(kevm.pretty_print(contract_main_definition) + '\n')
         _LOGGER.info(f'Wrote file: {foundry.main_file}')
+        foundry.record_input_name_mapping(foundry.input_mapping_file)
+        _LOGGER.info(f'Wrote file: {foundry.input_mapping_file}')
 
     def kompilation_digest() -> str:
         k_files = list(options.requires) + [foundry_contracts_file, foundry.main_file]

src/kontrol/prove.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import logging
 import time
 from abc import abstractmethod
@@ -587,6 +588,10 @@ def method_to_apr_proof(
         config_type=config_type,
     )
 
+    input_mapping = json.loads(foundry.input_mapping_file.read_text())
+    proof_input_mapping = input_mapping.get(test.contract.name_with_path, {}).get(test.method.signature, {})
+    proof_input_mapping.update(input_mapping.get('env', []))
+
     apr_proof = APRProof(
         test.id,
         kcfg,
@@ -597,6 +602,7 @@ def method_to_apr_proof(
         bmc_depth=bmc_depth,
         proof_dir=foundry.proofs_dir,
         subproof_ids=summary_ids,
+        variable_names_mapping=proof_input_mapping,
     )
 
     return apr_proof

src/kontrol/solc_to_k.py

@@ -576,17 +576,17 @@ def flat_inputs(self) -> tuple[Input, ...]:
             return tuple(input for sub_inputs in self.inputs for input in sub_inputs.flattened())
 
         @cached_property
-        def arg_names(self) -> tuple[str, ...]:
-            arg_names: list[str] = []
+        def arg_names(self) -> dict[str, str]:
+            arg_names: dict[str, str] = {}
             for input in self.inputs:
                 if input.type.endswith('[]') and not input.type.startswith('tuple'):
                     if input.array_lengths is None:
                         raise ValueError(f'Array length bounds missing for {input.name}')
                     length = input.array_lengths[0]
-                    arg_names.extend(f'{input.arg_name}_{i}' for i in range(length))
+                    arg_names.update({f'{input.arg_name}_{i}': f'{input.name}_{i}' for i in range(length)})
                 else:
-                    arg_names.extend([sub_input.arg_name for sub_input in input.flattened()])
-            return tuple(arg_names)
+                    arg_names.update({sub_input.arg_name: sub_input.name for sub_input in input.flattened()})
+            return arg_names
 
         @cached_property
         def arg_types(self) -> tuple[str, ...]:
@@ -651,7 +651,7 @@ def rule(
             self, contract: KInner, application_label: KLabel, contract_name: str, enums: dict[str, int]
         ) -> KRule | None:
             prod_klabel = self.unique_klabel
-            arg_vars = [KVariable(name) for name in self.arg_names]
+            arg_vars = [KVariable(name) for name in self.arg_names.keys()]
             args: list[KInner] = []
             conjuncts: list[KInner] = []
             for input in self.inputs:
@@ -728,7 +728,7 @@ def application(self) -> KInner:
             assert klabel is not None
             args = [
                 abstract_term_safely(KVariable('_###SOLIDITY_ARG_VAR###_'), base_name=f'V{name}')
-                for name in self.arg_names
+                for name in self.arg_names.keys()
             ]
             return klabel(args)