Support for user-defined preconditions

pyproject.toml

@@ -15,6 +15,7 @@ python = "^3.10"
 kevm-pyk = { git = "https://github.com/runtimeverification/evm-semantics.git", tag = "v1.0.680", subdirectory = "kevm-pyk" }
 eth-utils = "^4.1.1"
 pycryptodome = "^3.20.0"
+antlr4-python3-runtime = "^4.13.1"
 
 [tool.poetry.group.dev.dependencies]
 autoflake = "*"
@@ -51,12 +52,21 @@ remove-all-unused-imports = true
 ignore-init-module-imports = true
 remove-duplicate-keys = true
 remove-unused-variables = true
-exclude = "src/tests/integration/test-data"
+exclude = [
+    "src/tests/integration/test-data",
+    "src/kontrol/solidity",
+]
 
 [tool.black]
 line-length = 120
 skip-string-normalization = true
 
+[tool.flake8]
+exclude = "src/kontrol/solidity"
+
 [tool.mypy]
 disallow_untyped_defs = true
-exclude = "src/tests/integration/test-data"
+exclude = [
+    "src/tests/integration/test-data",
+    "src/kontrol/solidity",
+]

src/kontrol/prove.py

@@ -9,6 +9,7 @@
 from subprocess import CalledProcessError
 from typing import TYPE_CHECKING, Any, ContextManager, NamedTuple
 
+from antlr4 import CommonTokenStream, InputStream  # type: ignore
 from kevm_pyk.kevm import KEVM, KEVMSemantics, _process_jumpdests
 from kevm_pyk.utils import KDefinition__expand_macros, abstract_cell_vars, run_prover
 from multiprocess.pool import Pool  # type: ignore
@@ -34,6 +35,9 @@
 from .hevm import Hevm
 from .options import ConfigType, TraceOptions
 from .solc_to_k import Contract, hex_string_to_int
+from .solidity.AnnotationVisitor import AnnotationVisitor
+from .solidity.SolidityLexer import SolidityLexer
+from .solidity.SolidityParser import SolidityParser
 from .state_record import StateDiffEntry, StateDumpEntry
 from .utils import console, parse_test_version_tuple
 
@@ -1034,7 +1038,12 @@ def _init_cterm(
         else:
             # Symbolic accounts of all relevant contracts
             accounts, storage_constraints = _create_cse_accounts(
-                foundry, storage_fields, contract_account_name, contract_code
+                foundry, storage_fields, contract_account_name, contract_code, method
+            )
+
+            # Add precondition constraints
+            storage_constraints.extend(
+                _create_precondition_constraints(foundry, storage_fields, contract_account_name, method)
             )
 
         accounts.append(KVariable('ACCOUNTS_REST', sort=KSort('AccountCellMap')))
@@ -1078,6 +1087,11 @@ def _init_cterm(
     for constraint in storage_constraints:
         init_cterm = init_cterm.add_constraint(constraint)
 
+    # TODO(palina): we don't need that as long as we're doing it in _cse_...
+    # if method.preconditions is not None:
+    #     for precondition in method.preconditions:
+    #         init_cterm = init_cterm.add_constraint(mlEqualsTrue(precondition.to_kapply))
+
     # 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(
@@ -1120,11 +1134,52 @@ def _create_initial_account_list(
     return init_account_list
 
 
+def _create_precondition_constraints(
+    foundry: Foundry,
+    storage_fields: tuple[StorageField, ...],
+    contract_name: str,
+    method: Contract.Method,
+) -> list[KApply]:
+
+    precondition_constraints: list[KApply] = []
+
+    def precondition_to_kapply(precondition: str) -> KApply:
+        """
+        Converts the precondition to a KApply term.
+
+        - Constants are translated into `intToken` terms
+        - Variables should be searched in method's inputs or in the contract's storage fields
+        - Operators are translated into the corresponding KLabel application, e.g., `leInt`, `eqInt`, etc.
+        TBD
+        """
+
+        # Parse the input expression
+        input_stream = InputStream(precondition)
+        lexer = SolidityLexer(input_stream)
+
+        stream = CommonTokenStream(lexer)
+        parser = SolidityParser(stream)
+        tree = parser.expression()
+
+        # Evaluate the expression
+        evaluator = AnnotationVisitor(method, foundry, storage_fields, contract_name)
+        result = evaluator.visit(tree)
+
+        return result
+
+    if method.preconditions is not None:
+        for p in method.preconditions:
+            precondition_constraints.append(mlEqualsTrue(precondition_to_kapply(p.precondition)))
+
+    return precondition_constraints
+
+
 def _create_cse_accounts(
     foundry: Foundry,
     storage_fields: tuple[StorageField, ...],
     contract_name: str,
     contract_code: KInner,
+    method: Contract.Method,
 ) -> tuple[list[KInner], list[KApply]]:
     """
     Recursively generates a list of new accounts corresponding to `contract` fields, each having <code> and <storage> cell (partially) set up.
@@ -1293,7 +1348,7 @@ def extend_storage(map: KInner, slot: int, value: KInner) -> KInner:
                         )
 
                     contract_accounts, contract_constraints = _create_cse_accounts(
-                        foundry, contract_obj.fields, field_name, contract_account_code
+                        foundry, contract_obj.fields, field_name, contract_account_code, method
                     )
                     new_accounts.extend(contract_accounts)
                     new_account_constraints.extend(contract_constraints)

src/kontrol/solc_to_k.py

@@ -337,6 +337,32 @@ def parse_devdoc(tag: str, devdoc: dict | None) -> dict:
     return natspecs
 
 
+def parse_annotations(
+    devdoc: str | None, method: Contract.Method | Contract.Constructor
+) -> tuple[Precondition, ...] | None:
+    """
+    Parse developer documentation (devdoc) to extract user-provided preconditions.
+
+    Returns a tuple of Precondition objects, each representing a single precondition and a method to which it belongs.
+    """
+
+    if devdoc is None:
+        return None
+
+    preconditions: list[Precondition] = []
+    for precondition in devdoc.split(','):
+        # Trim whitespace and skip if empty
+        precondition = precondition.strip()
+        if not precondition:
+            continue
+
+        # Create a Precondition object and add it to the list
+        new_precondition = Precondition(precondition, method)
+        preconditions.append(new_precondition)
+
+    return tuple(preconditions)
+
+
 class StorageField(NamedTuple):
     label: str
     data_type: str
@@ -345,6 +371,19 @@ class StorageField(NamedTuple):
     linked_interface: str | None
 
 
+@dataclass
+class Precondition:
+    precondition: str
+    method: Contract.Method | Contract.Constructor
+
+    def __init__(self, precondition: str, method: Contract.Method | Contract.Constructor):
+        """
+        Initializes a new instance of the Precondition class.
+        """
+        self.precondition = precondition
+        self.method = method
+
+
 @dataclass
 class Contract:
     @dataclass
@@ -356,6 +395,7 @@ class Constructor:
         contract_storage_digest: str
         payable: bool
         signature: str
+        preconditions: tuple[Precondition, ...] | None
 
         def __init__(
             self,
@@ -369,7 +409,7 @@ def __init__(
             self.contract_name = contract_name
             self.contract_digest = contract_digest
             self.contract_storage_digest = contract_storage_digest
-            # TODO: support NatSpec comments for dynamic types
+            # TODO: support NatSpec comments for dynamic types, including preconditions
             self.inputs = tuple(inputs_from_abi(abi['inputs'], None))
             self.sort = sort
             # TODO: Check that we're handling all state mutability cases
@@ -487,6 +527,7 @@ class Method:
         ast: dict | None
         natspec_values: dict | None
         function_calls: tuple[str, ...] | None
+        preconditions: tuple[Precondition, ...] | None
 
         def __init__(
             self,
@@ -517,6 +558,9 @@ def __init__(
             natspec_tags = ['custom:kontrol-array-length-equals', 'custom:kontrol-bytes-length-equals']
             self.natspec_values = {tag.split(':')[1]: parse_devdoc(tag, devdoc) for tag in natspec_tags}
             self.inputs = tuple(inputs_from_abi(abi['inputs'], self.natspec_values))
+            self.preconditions = (
+                parse_annotations(devdoc.get('custom:kontrol-precondition', None), self) if devdoc is not None else None
+            )
             self.function_calls = tuple(function_calls) if function_calls is not None else None
 
         @property

src/kontrol/solidity/AnnotationVisitor.py

@@ -0,0 +1,148 @@
+from kevm_pyk.kevm import KEVM
+from pyk.kast import KInner
+from pyk.kast.inner import KApply, KSort, KVariable
+from pyk.kast.manip import abstract_term_safely
+from pyk.prelude.kbool import FALSE, TRUE
+from pyk.prelude.kint import intToken
+
+from ..foundry import Foundry
+from ..solc_to_k import Contract, StorageField
+from .SolidityParser import SolidityParser
+from .SolidityVisitor import SolidityVisitor
+
+
+class AnnotationVisitor(SolidityVisitor):
+    def __init__(
+        self,
+        method: Contract.Method | Contract.Constructor,
+        foundry: Foundry,
+        storage_fields: tuple[StorageField, ...],
+        contract_name: str,
+    ):
+        self.method = method
+        self.foundry = foundry
+        self.storage_fields = storage_fields
+        self.contract_name = contract_name
+
+    # def visitAndExpression(self, ctx: SolidityParser.AndExpressionContext):
+    #     left = self.visit(ctx.booleanExpression(0))
+    #     right = self.visit(ctx.booleanExpression(1))
+    #     return left and right
+
+    # def visitOrExpression(self, ctx: SolidityParser.OrExpressionContext):
+    #     left = self.visit(ctx.booleanExpression(0))
+    #     right = self.visit(ctx.booleanExpression(1))
+    #     return left or right
+
+    # def visitNotExpression(self, ctx: SolidityParser.NotExpressionContext):
+    #     value = self.visit(ctx.booleanExpression())
+    #     return not value
+
+    def visitRelationalExpression(self, ctx: SolidityParser.RelationalExpressionContext) -> KApply:
+        left = self.visit(ctx.arithmeticExpression(0))
+        right = self.visit(ctx.arithmeticExpression(1))
+
+        op = ctx.RelOp().getText()
+
+        # Map operators to KLabel applications
+        operator_mapping = {
+            '<=': '_<=Int_',
+            '>=': '_>=Int_',
+            '==': '_==Int_',
+            '!=': '_=/=Int_',
+            '<': '_<Int_',
+            '>': '_>Int_',
+        }
+
+        if op in operator_mapping:
+            operator_label = operator_mapping[op]
+        else:
+            raise ValueError(f'Unsupported operator in a precondition: {op}')
+
+        return KApply(operator_label, left, right)
+
+    def visitBooleanLiteral(self, ctx: SolidityParser.BooleanLiteralContext) -> KInner:
+        return TRUE if ctx.getText() == 'true' else FALSE
+
+    def visitVariable(self, ctx: SolidityParser.VariableContext) -> KInner:
+        var_name = ctx.getText()
+        # Search for matches in function inputs
+        for input in self.method.inputs:
+            if input.name == var_name:
+                # TODO: add support for complex types
+                return abstract_term_safely(KVariable('_###SOLIDITY_ARG_VAR###_'), base_name=f'V{input.arg_name}')
+
+        # Search for matches in contract storage fields
+        for field in self.storage_fields:
+            if field.label == var_name:
+                storage_map: KInner = KVariable(self.contract_name + '_STORAGE', sort=KSort('Map'))
+                return KEVM.lookup(storage_map, intToken(field.slot))
+
+        raise ValueError(f'Variable {var_name} not found in function inputs or storage fields of {self.method.name}.')
+
+    def visitContractVariableAccess(self, ctx: SolidityParser.ContractVariableAccessContext):
+        contract_field_name: str = ctx.contractVariableAccessExpr().VariableName(0).getText()
+        var_name: str = ctx.contractVariableAccessExpr().VariableName(1).getText()
+
+        for field in self.storage_fields:
+            if field.data_type.startswith('contract ') and field.label == contract_field_name:
+                contract_type = field.data_type.split(' ')[1]
+
+                # TODO: it is possible for a contact to have an interface annotation, `linked_interface`
+                for full_contract_name, contract_obj in self.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
+                    if full_contract_name.split('%')[-1] == contract_type:
+                        for field in contract_obj.fields:
+                            if field.label == var_name:
+                                storage_map: KInner = KVariable(
+                                    self.contract_name + '_' + contract_field_name.upper() + '_STORAGE', sort=KSort('Map')
+                                )
+                                return KEVM.lookup(storage_map, intToken(field.slot))
+
+        raise ValueError(f'Variable {contract_field_name}.{var_name} not found.')
+
+    def visitIntegerLiteral(self, ctx: SolidityParser.IntegerLiteralContext) -> KInner:
+        return intToken(ctx.getText())
+
+    def visitAddExpression(self, ctx: SolidityParser.AddExpressionContext):
+        left = self.visit(ctx.arithmeticExpression(0))
+        right = self.visit(ctx.arithmeticExpression(1))
+        return KApply('_+Int_', left, right)
+
+    def visitSubtractExpression(self, ctx: SolidityParser.SubtractExpressionContext):
+        left = self.visit(ctx.arithmeticExpression(0))
+        right = self.visit(ctx.arithmeticExpression(1))
+        return KApply('_-Int_', left, right)
+
+    def visitMultiplyExpression(self, ctx: SolidityParser.MultiplyExpressionContext):
+        left = self.visit(ctx.arithmeticExpression(0))
+        right = self.visit(ctx.arithmeticExpression(1))
+        return KApply('_*Int_', left, right)
+
+    def visitDivideExpression(self, ctx: SolidityParser.DivideExpressionContext):
+        left = self.visit(ctx.arithmeticExpression(0))
+        right = self.visit(ctx.arithmeticExpression(1))
+        return KApply('_/Int_', left, right)
+
+    def visitPowExpression(self, ctx: SolidityParser.PowExpressionContext):
+        left = self.visit(ctx.arithmeticExpression(0))
+        right = self.visit(ctx.arithmeticExpression(1))
+        return KApply('_^Int_', left, right)
+
+    # def visitLengthAccess(self, ctx: SolidityParser.LengthAccessContext):
+    #     var_name = ctx.variableName().getText()
+    #     return len(self.context.get(var_name, ""))
+
+    # def visitArrayElement(self, ctx: SolidityParser.ArrayElementContext):
+    #     var_name = ctx.variableName().getText()
+    #     index = int(ctx.INTEGER().getText())
+    #     return self.context.get(var_name, [])[index]
+
+    # def visitMappingElement(self, ctx: SolidityParser.MappingElementContext):
+    #     var_name = ctx.variableName().getText()
+    #     key = ctx.variableName().getText()
+    #     return self.context.get(var_name, {}).get(key, 0)
+
+    # def visitAddressLiteral(self, ctx: SolidityParser.AddressLiteralContext):
+    #     return ctx.getText()