Implementation ExecutionState::STOP (privacy-scaling-explorations#419)

* feat: implement STOP

* fix: add more comments and rename `CodeSource` to `CodeHash`

* feat: use `HashMap` to store `bytecodes` in `witness::Block`

bus-mapping/src/circuit_input_builder/input_state_ref.rs

@@ -370,6 +370,13 @@ impl<'a> CircuitInputStateRef<'a> {
             .ok_or(Error::CodeNotFound(code_hash))
     }
 
+    /// Reference to the caller's Call
+    pub fn caller(&self) -> Result<&Call, Error> {
+        self.tx_ctx
+            .caller_index()
+            .map(|caller_idx| &self.tx.calls()[caller_idx])
+    }
+
     /// Reference to the current Call
     pub fn call(&self) -> Result<&Call, Error> {
         self.tx_ctx
@@ -384,6 +391,11 @@ impl<'a> CircuitInputStateRef<'a> {
             .map(|call_idx| &mut self.tx.calls_mut()[call_idx])
     }
 
+    /// Reference to the current CallContext
+    pub fn caller_ctx(&self) -> Result<&CallContext, Error> {
+        self.tx_ctx.caller_ctx()
+    }
+
     /// Reference to the current CallContext
     pub fn call_ctx(&self) -> Result<&CallContext, Error> {
         self.tx_ctx.call_ctx()

bus-mapping/src/circuit_input_builder/transaction.rs

@@ -96,14 +96,24 @@ impl TransactionContext {
         &self.calls
     }
 
+    /// Return the index of the caller (the second last call in the call stack).
+    pub(crate) fn caller_index(&self) -> Result<usize, Error> {
+        self.caller_ctx().map(|call| call.index)
+    }
+
     /// Return the index of the current call (the last call in the call stack).
     pub(crate) fn call_index(&self) -> Result<usize, Error> {
+        self.call_ctx().map(|call| call.index)
+    }
+
+    pub(crate) fn caller_ctx(&self) -> Result<&CallContext, Error> {
         self.calls
-            .last()
+            .len()
+            .checked_sub(2)
+            .map(|idx| &self.calls[idx])
             .ok_or(Error::InvalidGethExecTrace(
                 "Call stack is empty but call is used",
             ))
-            .map(|call| call.index)
     }
 
     pub(crate) fn call_ctx(&self) -> Result<&CallContext, Error> {

bus-mapping/src/evm/opcodes.rs

@@ -34,6 +34,7 @@ mod mload;
 mod mstore;
 mod number;
 mod origin;
+mod r#return;
 mod selfbalance;
 mod sload;
 mod sstore;
@@ -56,6 +57,7 @@ use logs::Log;
 use mload::Mload;
 use mstore::Mstore;
 use origin::Origin;
+use r#return::Return;
 use selfbalance::Selfbalance;
 use sload::Sload;
 use sstore::Sstore;
@@ -200,13 +202,15 @@ fn fn_gen_associated_ops(opcode_id: &OpcodeId) -> FnGenAssociatedOps {
         // OpcodeId::CREATE => {},
         OpcodeId::CALL => Call::gen_associated_ops,
         // OpcodeId::CALLCODE => {},
-        // TODO: Handle RETURN by its own gen_associated_ops.
-        OpcodeId::RETURN => Stop::gen_associated_ops,
+        // OpcodeId::RETURN => {},
         // OpcodeId::DELEGATECALL => {},
         // OpcodeId::CREATE2 => {},
         // OpcodeId::STATICCALL => {},
-        // TODO: Handle REVERT by its own gen_associated_ops.
-        OpcodeId::REVERT => Stop::gen_associated_ops,
+        // OpcodeId::REVERT => {},
+        OpcodeId::REVERT | OpcodeId::RETURN => {
+            warn!("Using dummy gen_associated_ops for opcode {:?}", opcode_id);
+            Return::gen_associated_ops
+        }
         OpcodeId::SELFDESTRUCT => {
             warn!("Using dummy gen_selfdestruct_ops for opcode SELFDESTRUCT");
             dummy_gen_selfdestruct_ops
@@ -375,7 +379,7 @@ pub fn gen_begin_tx_ops(state: &mut CircuitInputStateRef) -> Result<ExecStep, Er
                 (CallContextField::LastCalleeReturnDataLength, 0.into()),
                 (CallContextField::IsRoot, 1.into()),
                 (CallContextField::IsCreate, 0.into()),
-                (CallContextField::CodeSource, code_hash.to_word()),
+                (CallContextField::CodeHash, code_hash.to_word()),
             ] {
                 state.push_op(
                     &mut exec_step,

bus-mapping/src/evm/opcodes/call.rs

@@ -180,7 +180,7 @@ impl Opcode for Call {
                     ),
                     (CallContextField::MemorySize, next_memory_word_size.into()),
                     (
-                        CallContextField::StateWriteCounter,
+                        CallContextField::ReversibleWriteCounter,
                         (exec_step.reversible_write_counter + 1).into(),
                     ),
                 ] {
@@ -220,7 +220,7 @@ impl Opcode for Call {
                     (CallContextField::LastCalleeReturnDataLength, 0.into()),
                     (CallContextField::IsRoot, 0.into()),
                     (CallContextField::IsCreate, 0.into()),
-                    (CallContextField::CodeSource, call.code_hash.to_word()),
+                    (CallContextField::CodeHash, call.code_hash.to_word()),
                 ] {
                     state.call_context_read(&mut exec_step, call.call_id, field, value);
                 }

bus-mapping/src/evm/opcodes/codesize.rs

@@ -18,8 +18,8 @@ impl Opcode for Codesize {
         let geth_step = &geth_steps[0];
         let mut exec_step = state.new_step(geth_step)?;
 
-        let code_source = state.call()?.code_hash;
-        let code = state.code(code_source)?;
+        let code_hash = state.call()?.code_hash;
+        let code = state.code(code_hash)?;
         let codesize = code.len();
 
         debug_assert_eq!(codesize, geth_steps[1].stack.last()?.as_usize());

bus-mapping/src/evm/opcodes/return.rs

@@ -0,0 +1,24 @@
+use super::Opcode;
+use crate::circuit_input_builder::{CircuitInputStateRef, ExecStep};
+use crate::Error;
+use eth_types::GethExecStep;
+
+/// Placeholder structure used to implement [`Opcode`] trait over it
+/// corresponding to the [`OpcodeId::RETURN`](crate::evm::OpcodeId::RETURN).
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct Return;
+
+impl Opcode for Return {
+    fn gen_associated_ops(
+        state: &mut CircuitInputStateRef,
+        geth_steps: &[GethExecStep],
+    ) -> Result<Vec<ExecStep>, Error> {
+        let geth_step = &geth_steps[0];
+        let exec_step = state.new_step(geth_step)?;
+
+        // TODO: Generate associated operations of RETURN
+
+        state.handle_return(geth_step)?;
+        Ok(vec![exec_step])
+    }
+}

bus-mapping/src/evm/opcodes/stop.rs

@@ -1,7 +1,10 @@
 use super::Opcode;
-use crate::circuit_input_builder::{CircuitInputStateRef, ExecStep};
-use crate::Error;
-use eth_types::GethExecStep;
+use crate::{
+    circuit_input_builder::{CircuitInputStateRef, ExecStep},
+    operation::{CallContextField, CallContextOp, RW},
+    Error,
+};
+use eth_types::{GethExecStep, ToWord};
 
 /// Placeholder structure used to implement [`Opcode`] trait over it
 /// corresponding to the [`OpcodeId::STOP`](crate::evm::OpcodeId::STOP)
@@ -18,8 +21,100 @@ impl Opcode for Stop {
         geth_steps: &[GethExecStep],
     ) -> Result<Vec<ExecStep>, Error> {
         let geth_step = &geth_steps[0];
-        let exec_step = state.new_step(geth_step)?;
+        let mut exec_step = state.new_step(geth_step)?;
+
+        let call = state.call()?.clone();
+
+        state.push_op(
+            &mut exec_step,
+            RW::READ,
+            CallContextOp {
+                call_id: call.call_id,
+                field: CallContextField::IsSuccess,
+                value: 1.into(),
+            },
+        );
+
+        if call.is_root {
+            state.push_op(
+                &mut exec_step,
+                RW::READ,
+                CallContextOp {
+                    call_id: call.call_id,
+                    field: CallContextField::IsPersistent,
+                    value: 1.into(),
+                },
+            );
+        } else {
+            // The following part corresponds to
+            // Instruction.step_state_transition_to_restored_context
+            // in python spec, and should be reusable among all expected halting opcodes or
+            // exceptions. TODO: Refactor it as a helper function.
+            let caller = state.caller()?.clone();
+            state.push_op(
+                &mut exec_step,
+                RW::READ,
+                CallContextOp {
+                    call_id: call.call_id,
+                    field: CallContextField::CallerId,
+                    value: caller.call_id.into(),
+                },
+            );
+
+            let geth_step_next = &geth_steps[1];
+            let caller_gas_left = geth_step_next.gas.0 - geth_step.gas.0;
+            for (field, value) in [
+                (CallContextField::IsRoot, (caller.is_root as u64).into()),
+                (
+                    CallContextField::IsCreate,
+                    (caller.is_create() as u64).into(),
+                ),
+                (CallContextField::CodeHash, caller.code_hash.to_word()),
+                (CallContextField::ProgramCounter, geth_step_next.pc.0.into()),
+                (
+                    CallContextField::StackPointer,
+                    geth_step_next.stack.stack_pointer().0.into(),
+                ),
+                (CallContextField::GasLeft, caller_gas_left.into()),
+                (
+                    CallContextField::MemorySize,
+                    geth_step_next.memory.word_size().into(),
+                ),
+                (
+                    CallContextField::ReversibleWriteCounter,
+                    state.caller_ctx()?.reversible_write_counter.into(),
+                ),
+            ] {
+                state.push_op(
+                    &mut exec_step,
+                    RW::READ,
+                    CallContextOp {
+                        call_id: caller.call_id,
+                        field,
+                        value,
+                    },
+                );
+            }
+
+            for (field, value) in [
+                (CallContextField::LastCalleeId, call.call_id.into()),
+                (CallContextField::LastCalleeReturnDataOffset, 0.into()),
+                (CallContextField::LastCalleeReturnDataLength, 0.into()),
+            ] {
+                state.push_op(
+                    &mut exec_step,
+                    RW::WRITE,
+                    CallContextOp {
+                        call_id: caller.call_id,
+                        field,
+                        value,
+                    },
+                );
+            }
+        }
+
         state.handle_return(geth_step)?;
+
         Ok(vec![exec_step])
     }
 }

bus-mapping/src/operation.rs

@@ -697,8 +697,8 @@ pub enum CallContextField {
     IsRoot,
     /// IsCreate
     IsCreate,
-    /// CodeSource
-    CodeSource,
+    /// CodeHash
+    CodeHash,
     /// ProgramCounter
     ProgramCounter,
     /// StackPointer
@@ -707,8 +707,8 @@ pub enum CallContextField {
     GasLeft,
     /// MemorySize
     MemorySize,
-    /// StateWriteCounter
-    StateWriteCounter,
+    /// ReversibleWriteCounter
+    ReversibleWriteCounter,
 }
 
 /// Represents an CallContext read/write operation.

zkevm-circuits/src/evm_circuit.rs

@@ -270,10 +270,10 @@ pub mod test {
             )
         }
 
-        fn load_bytecodes(
+        fn load_bytecodes<'a>(
             &self,
             layouter: &mut impl Layouter<F>,
-            bytecodes: &[Bytecode],
+            bytecodes: impl IntoIterator<Item = &'a Bytecode> + Clone,
             randomness: F,
         ) -> Result<(), Error> {
             layouter.assign_region(
@@ -290,7 +290,7 @@ pub mod test {
                     }
                     offset += 1;
 
-                    for bytecode in bytecodes.iter() {
+                    for bytecode in bytecodes.clone() {
                         for row in bytecode.table_assignments(randomness) {
                             for (column, value) in self.bytecode_table.iter().zip_eq(row) {
                                 region.assign_advice(
@@ -420,7 +420,11 @@ pub mod test {
             config.evm_circuit.load_byte_table(&mut layouter)?;
             config.load_txs(&mut layouter, &self.block.txs, self.block.randomness)?;
             config.load_rws(&mut layouter, &self.block.rws, self.block.randomness)?;
-            config.load_bytecodes(&mut layouter, &self.block.bytecodes, self.block.randomness)?;
+            config.load_bytecodes(
+                &mut layouter,
+                self.block.bytecodes.values(),
+                self.block.randomness,
+            )?;
             config.load_block(&mut layouter, &self.block.context, self.block.randomness)?;
             config
                 .evm_circuit
@@ -459,7 +463,7 @@ pub mod test {
         let k = k.max(log2_ceil(
             64 + block
                 .bytecodes
-                .iter()
+                .values()
                 .map(|bytecode| bytecode.bytes.len())
                 .sum::<usize>(),
         ));