Add mypy

.github/workflows/python-package.yml

@@ -27,5 +27,7 @@ jobs:
         make install
     - name: Lint
       run: make lint
+    - name: Type check
+      run: make type
     - name: Test with pytest
       run: make test

Makefile

@@ -14,6 +14,9 @@ lint: ## Check whether the code is formated correctly
 	black . --check
 	mdformat specs/ --number --check
 
+type: ## Check the typing of the Python code
+	mypy src
+
 test: ## Run tests
 	pytest
 

setup.cfg

@@ -31,3 +31,4 @@ test =
 lint =
     black >= 22.1.0
     mdformat >= 0.7.13
+    mypy >= 0.931

src/zkevm_specs/__init__.py

@@ -1,3 +1,4 @@
+from . import bytecode
 from . import encoding
 from . import evm
 from . import opcode

src/zkevm_specs/bytecode.py

@@ -1,7 +1,7 @@
-from typing import Sequence, Union, Tuple, Set
+from typing import Sequence, Union, Tuple, Set, NamedTuple
 from collections import namedtuple
 from .util import keccak256, FQ, RLC
-from .evm.opcode import get_push_size
+from .evm import get_push_size, BytecodeTableRow
 from .encoding import U8, U256, is_circuit_code
 
 # Row in the circuit
@@ -10,7 +10,9 @@
     "q_first q_last hash index byte is_code push_data_left hash_rlc hash_length byte_push_size is_final padding",
 )
 # Unrolled bytecode
-UnrolledBytecode = namedtuple("UnrolledBytecode", "bytes rows")
+class UnrolledBytecode(NamedTuple):
+    bytes: bytes
+    rows: Sequence[BytecodeTableRow]
 
 
 @is_circuit_code
@@ -24,7 +26,7 @@ def select(
     when_true: U256,
     when_false: U256,
 ) -> U256:
-    return selector * when_true + (1 - selector) * when_false
+    return U256(selector * when_true + (1 - selector) * when_false)
 
 
 @is_circuit_code
@@ -107,21 +109,21 @@ def assign_bytecode_circuit(k: int, bytecodes: Sequence[UnrolledBytecode], rando
         for idx, row in enumerate(bytecode.rows):
             # Track which byte is an opcode and which is push data
             is_code = push_data_left == 0
-            byte_push_size = get_push_size(row[2])
+            byte_push_size = get_push_size(row.byte)
             push_data_left = byte_push_size if is_code else push_data_left - 1
 
             # Add the byte to the accumulator
-            hash_rlc = hash_rlc * randomness + row[2]
+            hash_rlc = hash_rlc * randomness + row.byte
 
             # Set the data for this row
             rows.append(
                 Row(
                     offset == 0,
                     offset == last_row_offset,
-                    row[0],
-                    row[1],
-                    row[2],
-                    row[3],
+                    row.bytecode_hash,
+                    row.index,
+                    row.byte,
+                    row.is_code,
                     push_data_left,
                     hash_rlc,
                     len(bytecode.bytes),
@@ -178,10 +180,10 @@ def assign_push_table():
 
 
 # Generate keccak table
-def assign_keccak_table(bytecodes: Sequence[bytes], randomness: int):
+def assign_keccak_table(bytecodes: Sequence[bytes], randomness: FQ):
     keccak_table = []
     for bytecode in bytecodes:
         hash = RLC(bytes(reversed(keccak256(bytecode))), randomness)
         rlc = RLC(bytes(reversed(bytecode)), randomness, len(bytecode))
-        keccak_table.append((rlc, len(bytecode), hash))
+        keccak_table.append((rlc.expr(), len(bytecode), hash.expr()))
     return _convert_table(keccak_table)

src/zkevm_specs/encoding/lookup.py

@@ -2,7 +2,7 @@
 
 
 class LookupTable:
-    columns: Tuple[str]
+    columns: Tuple[str, ...]
     rows: Set[Tuple[int, ...]]
 
     def __init__(self, columns: Sequence[str]) -> None:

src/zkevm_specs/encoding/utils.py

@@ -1,8 +1,8 @@
-from typing import Sequence, Tuple, List
+from typing import Sequence, Tuple
 from .typing import U8, U256, U64
 
 
-def is_circuit_code(func) -> object:
+def is_circuit_code(func):
     """
     A no-op decorator just to mark the function
     """
@@ -15,19 +15,19 @@ def wrapper(*args, **kargs):
 
 def u256_to_u8s(x: U256) -> Tuple[U8, ...]:
     assert 0 <= x < 2**256, "expect x is unsigned 256 bits"
-    return tuple((x >> 8 * i) & 0xFF for i in range(32))
+    return tuple(U8((x >> 8 * i) & 0xFF) for i in range(32))
 
 
 def u256_to_u64s(x: U256) -> Tuple[U64, ...]:
     assert 0 <= x < 2**256, "expect x is unsigned 256 bits"
-    return tuple((x >> 64 * i) & 0xFFFFFFFFFFFFFFFF for i in range(4))
+    return tuple(U64((x >> 64 * i) & 0xFFFFFFFFFFFFFFFF) for i in range(4))
 
 
 def u8s_to_u256(xs: Sequence[U8]) -> U256:
     assert len(xs) == 32
     for u8 in xs:
         assert 0 <= u8 <= 255
-    return sum(x * (2 ** (8 * i)) for i, x in enumerate(xs))
+    return U256(sum(x * (2 ** (8 * i)) for i, x in enumerate(xs)))
 
 
 # [u8;32]->[u64;4]
@@ -37,5 +37,5 @@ def u8s_to_u64s(xs: Sequence[U8]) -> Tuple[U64, ...]:
     A = [u64_0] * 4  # A = A3A2A1A0
     for i in range(4):
         for j in range(8):
-            A[i] += U64(xs[j + 8 * i] * (2 ** (8 * j)))
+            A[i] += xs[j + 8 * i] * (2 ** (8 * j))
     return tuple(A)

src/zkevm_specs/evm/__init__.py

@@ -6,4 +6,4 @@
 from .step import *
 from .table import *
 from .typing import *
-from . import util
+from .util import *

src/zkevm_specs/evm/execution/begin_tx.py

@@ -1,6 +1,6 @@
-from ...util import GAS_COST_TX, GAS_COST_CREATION_TX, EMPTY_CODE_HASH
+from ...util import GAS_COST_TX, GAS_COST_CREATION_TX, EMPTY_CODE_HASH, FQ, RLC, cast_expr
 from ..execution_state import ExecutionState
-from ..instruction import Instruction, Transition
+from ..instruction import Instruction, ReversionInfo, Transition
 from ..precompiled import PrecompiledAddress
 from ..table import CallContextFieldTag, TxContextFieldTag, AccountFieldTag
 
@@ -9,28 +9,23 @@ def begin_tx(instruction: Instruction):
     call_id = instruction.curr.rw_counter
 
     tx_id = instruction.call_context_lookup(CallContextFieldTag.TxId, call_id=call_id)
-    rw_counter_end_of_reversion = instruction.call_context_lookup(
-        CallContextFieldTag.RwCounterEndOfReversion, call_id=call_id
-    )
-    is_persistent = instruction.call_context_lookup(
-        CallContextFieldTag.IsPersistent, call_id=call_id
-    )
+    reversion_info = instruction.reversion_info(call_id=call_id)
 
     if instruction.is_first_step:
-        instruction.constrain_equal(instruction.curr.rw_counter, 1)
-        instruction.constrain_equal(tx_id, 1)
+        instruction.constrain_equal(instruction.curr.rw_counter, FQ(1))
+        instruction.constrain_equal(tx_id, FQ(1))
 
     tx_caller_address = instruction.tx_context_lookup(tx_id, TxContextFieldTag.CallerAddress)
     tx_callee_address = instruction.tx_context_lookup(tx_id, TxContextFieldTag.CalleeAddress)
     tx_is_create = instruction.tx_context_lookup(tx_id, TxContextFieldTag.IsCreate)
-    tx_value = instruction.tx_context_lookup(tx_id, TxContextFieldTag.Value)
+    tx_value = cast_expr(instruction.tx_context_lookup(tx_id, TxContextFieldTag.Value), RLC)
     tx_call_data_length = instruction.tx_context_lookup(tx_id, TxContextFieldTag.CallDataLength)
 
     # Verify nonce
     tx_nonce = instruction.tx_context_lookup(tx_id, TxContextFieldTag.Nonce)
     nonce, nonce_prev = instruction.account_write(tx_caller_address, AccountFieldTag.Nonce)
     instruction.constrain_equal(tx_nonce, nonce_prev)
-    instruction.constrain_equal(nonce, nonce_prev + 1)
+    instruction.constrain_equal(nonce, nonce_prev.expr() + 1)
 
     # TODO: Implement EIP 1559 (currently it supports legacy transaction format)
     # Calculate gas fee
@@ -42,24 +37,27 @@ def begin_tx(instruction: Instruction):
     # TODO: Handle gas cost of tx level access list (EIP 2930)
     tx_call_data_gas_cost = instruction.tx_context_lookup(tx_id, TxContextFieldTag.CallDataGasCost)
     gas_left = (
-        tx_gas
+        tx_gas.expr()
         - (GAS_COST_CREATION_TX if tx_is_create == 1 else GAS_COST_TX)
-        - tx_call_data_gas_cost
+        - tx_call_data_gas_cost.expr()
     )
     instruction.constrain_gas_left_not_underflow(gas_left)
 
     # Prepare access list of caller and callee
-    instruction.constrain_equal(instruction.add_account_to_access_list(tx_id, tx_caller_address), 1)
-    instruction.constrain_equal(instruction.add_account_to_access_list(tx_id, tx_callee_address), 1)
+    instruction.constrain_equal(
+        instruction.add_account_to_access_list(tx_id, tx_caller_address), FQ(1)
+    )
+    instruction.constrain_equal(
+        instruction.add_account_to_access_list(tx_id, tx_callee_address), FQ(1)
+    )
 
     # Verify transfer
     instruction.transfer_with_gas_fee(
         tx_caller_address,
         tx_callee_address,
         tx_value,
         gas_fee,
-        is_persistent,
-        rw_counter_end_of_reversion,
+        reversion_info,
     )
 
     if tx_is_create == 1:
@@ -79,16 +77,16 @@ def begin_tx(instruction: Instruction):
         # - TxId is checked from previous step or constraint to 1 if is_first_step
         # - IsSuccess, IsPersistent will be verified in the end of tx
         for (tag, value) in [
-            (CallContextFieldTag.Depth, 1),
+            (CallContextFieldTag.Depth, FQ(1)),
             (CallContextFieldTag.CallerAddress, tx_caller_address),
             (CallContextFieldTag.CalleeAddress, tx_callee_address),
-            (CallContextFieldTag.CallDataOffset, 0),
+            (CallContextFieldTag.CallDataOffset, FQ(0)),
             (CallContextFieldTag.CallDataLength, tx_call_data_length),
             (CallContextFieldTag.Value, tx_value),
-            (CallContextFieldTag.IsStatic, False),
-            (CallContextFieldTag.LastCalleeId, 0),
-            (CallContextFieldTag.LastCalleeReturnDataOffset, 0),
-            (CallContextFieldTag.LastCalleeReturnDataLength, 0),
+            (CallContextFieldTag.IsStatic, FQ(False)),
+            (CallContextFieldTag.LastCalleeId, FQ(0)),
+            (CallContextFieldTag.LastCalleeReturnDataOffset, FQ(0)),
+            (CallContextFieldTag.LastCalleeReturnDataLength, FQ(0)),
         ]:
             instruction.constrain_equal(
                 instruction.call_context_lookup(tag, call_id=call_id), value
@@ -104,9 +102,13 @@ def begin_tx(instruction: Instruction):
             state_write_counter=Transition.to(2),
         )
 
+        assert instruction.next is not None
+
         # Constrain either:
         # - is_empty_code and is_to_end_tx
         # - (not is_empty_code) and (not is_to_end_tx)
-        is_empty_code = instruction.is_equal(code_hash, instruction.int_to_rlc(EMPTY_CODE_HASH, 32))
+        is_empty_code = instruction.is_equal(
+            code_hash, RLC(EMPTY_CODE_HASH, instruction.randomness)
+        )
         is_to_end_tx = instruction.is_equal(instruction.next.execution_state, ExecutionState.EndTx)
         instruction.constrain_equal(is_empty_code + is_to_end_tx, 2 * is_empty_code * is_to_end_tx)

src/zkevm_specs/evm/execution/block_coinbase.py

@@ -1,3 +1,4 @@
+from ...util.param import N_BYTES_ACCOUNT_ADDRESS
 from ..instruction import Instruction, Transition
 from ..table import BlockContextFieldTag
 from ..opcode import Opcode
@@ -6,18 +7,14 @@
 def coinbase(instruction: Instruction):
     opcode = instruction.opcode_lookup(True)
     instruction.constrain_equal(opcode, Opcode.COINBASE)
-    address = instruction.stack_push()
+
     # in real circuit also check address raw data is 160 bit length (20 bytes)
     # check block table for coinbase address
     instruction.constrain_equal(
-        address,
-        instruction.int_to_rlc(
-            instruction.block_context_lookup(BlockContextFieldTag.Coinbase),
-            # NOTE: We can replace this with N_BYTES_WORD if we reuse the 32
-            # byte RLC constraint in all places.  See:
-            # https://github.com/appliedzkp/zkevm-specs/issues/101
-            20,
-        ),
+        instruction.block_context_lookup(BlockContextFieldTag.Coinbase),
+        # NOTE: We can replace this with N_BYTES_WORD if we reuse the 32 byte RLC constraint in
+        # all places. See: https://github.com/appliedzkp/zkevm-specs/issues/101
+        instruction.rlc_to_fq_exact(instruction.stack_push(), N_BYTES_ACCOUNT_ADDRESS),
     )
 
     instruction.step_state_transition_in_same_context(

src/zkevm_specs/evm/execution/block_timestamp.py

@@ -6,11 +6,11 @@
 def timestamp(instruction: Instruction):
     opcode = instruction.opcode_lookup(True)
     instruction.constrain_equal(opcode, Opcode.TIMESTAMP)
-    timestamp = instruction.stack_push()
+
     # check block table for timestamp
     instruction.constrain_equal(
-        timestamp,
-        instruction.int_to_rlc(instruction.block_context_lookup(BlockContextFieldTag.Timestamp), 8),
+        instruction.block_context_lookup(BlockContextFieldTag.Timestamp),
+        instruction.rlc_to_fq_exact(instruction.stack_push(), 8),
     )
 
     instruction.step_state_transition_in_same_context(

src/zkevm_specs/evm/execution/calldatacopy.py

@@ -1,24 +1,24 @@
-from ...util import N_BYTES_MEMORY_ADDRESS, FQ
+from ...util import N_BYTES_MEMORY_ADDRESS, FQ, Expression
 from ..execution_state import ExecutionState
 from ..instruction import Instruction, Transition
-from ..table import RW, FixedTableTag, RWTableTag, CallContextFieldTag, TxContextFieldTag
+from ..table import RW, CallContextFieldTag, TxContextFieldTag
 
 
 def calldatacopy(instruction: Instruction):
     opcode = instruction.opcode_lookup(True)
 
-    memory_offset = instruction.stack_pop()
-    data_offset = instruction.stack_pop()
-    length = instruction.stack_pop()
+    memory_offset_word = instruction.stack_pop()
+    data_offset_word = instruction.stack_pop()
+    length_word = instruction.stack_pop()
 
     # convert rlc to FQ
-    memory_offset, length = instruction.memory_offset_and_length(memory_offset, length)
-    data_offset = instruction.rlc_to_fq_exact(data_offset, N_BYTES_MEMORY_ADDRESS)
+    memory_offset, length = instruction.memory_offset_and_length(memory_offset_word, length_word)
+    data_offset = instruction.rlc_to_fq_exact(data_offset_word, N_BYTES_MEMORY_ADDRESS)
 
     tx_id = instruction.call_context_lookup(CallContextFieldTag.TxId, RW.Read)
     if instruction.curr.is_root:
         call_data_length = instruction.tx_context_lookup(tx_id, TxContextFieldTag.CallDataLength)
-        call_data_offset = FQ.zero()
+        call_data_offset: Expression = FQ.zero()
     else:
         call_data_length = instruction.call_context_lookup(
             CallContextFieldTag.CallDataLength, RW.Read
@@ -34,12 +34,15 @@ def calldatacopy(instruction: Instruction):
 
     # When length != 0, constrain the state in the next execution state CopyToMemory
     if not instruction.is_zero(length):
+        assert instruction.next is not None
         instruction.constrain_equal(instruction.next.execution_state, ExecutionState.CopyToMemory)
         next_aux = instruction.next.aux_data
         instruction.constrain_equal(next_aux.src_addr, data_offset + call_data_offset)
         instruction.constrain_equal(next_aux.dst_addr, memory_offset)
-        instruction.constrain_equal(next_aux.src_addr_end, call_data_length + call_data_offset)
-        instruction.constrain_equal(next_aux.from_tx, instruction.curr.is_root)
+        instruction.constrain_equal(
+            next_aux.src_addr_end, call_data_length.expr() + call_data_offset
+        )
+        instruction.constrain_equal(next_aux.from_tx, FQ(instruction.curr.is_root))
         instruction.constrain_equal(next_aux.tx_id, tx_id)
 
     instruction.step_state_transition_in_same_context(