Circuit for opcode SELFBALANCE (privacy-scaling-explorations#328)

* Implement bus mapping for SELFBALANCE

* Add SELFBALANCE evm circuit

* Replace caller with callee

* fmt

* Revert "Implement bus mapping for SELFBALANCE"

This reverts commit ecc000f.

* Use Cell to hold self_balance instead of Word

* Fix test failure from rebase

* Add bus mapping for SELFBALANCE

* Remove now unneeded test

* clippy

Co-authored-by: Mason Liang <[email protected]>
Co-authored-by: Mason Liang <[email protected]>
Co-authored-by: z2trillion <[email protected]>

bus-mapping/src/evm/opcodes.rs

@@ -19,6 +19,7 @@ mod mstore;
 mod pc;
 mod pop;
 mod push;
+mod selfbalance;
 mod sload;
 mod stackonlyop;
 mod stop;
@@ -41,6 +42,7 @@ use msize::Msize;
 use mstore::Mstore;
 use pc::Pc;
 use pop::Pop;
+use selfbalance::Selfbalance;
 use sload::Sload;
 use stackonlyop::StackOnlyOpcode;
 use stop::Stop;
@@ -123,7 +125,7 @@ fn fn_gen_associated_ops(opcode_id: &OpcodeId) -> FnGenAssociatedOps {
         // OpcodeId::DIFFICULTY => {},
         // OpcodeId::GASLIMIT => {},
         // OpcodeId::CHAINID => {},
-        // OpcodeId::SELFBALANCE => {},
+        OpcodeId::SELFBALANCE => Selfbalance::gen_associated_ops,
         // OpcodeId::BASEFEE => {},
         OpcodeId::POP => Pop::gen_associated_ops,
         OpcodeId::MLOAD => Mload::gen_associated_ops,

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

@@ -0,0 +1,132 @@
+use super::Opcode;
+use crate::circuit_input_builder::CircuitInputStateRef;
+use crate::operation::{AccountField, AccountOp, CallContextField, CallContextOp, RW};
+use crate::Error;
+use eth_types::{GethExecStep, ToWord};
+
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct Selfbalance;
+
+impl Opcode for Selfbalance {
+    fn gen_associated_ops(
+        state: &mut CircuitInputStateRef,
+        steps: &[GethExecStep],
+    ) -> Result<(), Error> {
+        let step = &steps[0];
+        let self_balance = steps[1].stack.last()?;
+        let callee_address = state.call().address;
+
+        // CallContext read of the callee_address
+        state.push_op(
+            RW::READ,
+            CallContextOp {
+                call_id: state.call().call_id,
+                field: CallContextField::CalleeAddress,
+                value: callee_address.to_word(),
+            },
+        );
+
+        // Account read for the balance of the callee_address
+        state.push_op(
+            RW::READ,
+            AccountOp {
+                address: callee_address,
+                field: AccountField::Balance,
+                value: self_balance,
+                value_prev: self_balance,
+            },
+        );
+
+        // Stack write of self_balance
+        state.push_stack_op(
+            RW::WRITE,
+            step.stack.last_filled().map(|a| a - 1),
+            self_balance,
+        );
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod selfbalance_tests {
+    use super::*;
+    use crate::circuit_input_builder::{ExecStep, TransactionContext};
+    use eth_types::{bytecode, evm_types::StackAddress, ToWord};
+    use pretty_assertions::assert_eq;
+
+    #[test]
+    fn selfbalance_opcode_impl() -> Result<(), Error> {
+        let code = bytecode! {
+            #[start]
+            SELFBALANCE
+            STOP
+        };
+
+        let mut geth_data = mock::new_single_tx_trace_code(&code)?;
+        geth_data.geth_trace.struct_logs =
+            geth_data.geth_trace.struct_logs[code.get_pos("start")..].to_vec();
+
+        // Get the execution steps from the external tracer
+        let block = crate::mock::BlockData::new_from_geth_data(geth_data);
+
+        let mut builder = block.new_circuit_input_builder();
+        builder.handle_tx(&block.eth_tx, &block.geth_trace).unwrap();
+
+        let mut test_builder = block.new_circuit_input_builder();
+        let mut tx = test_builder
+            .new_tx(&block.eth_tx, !block.geth_trace.failed)
+            .unwrap();
+        let mut tx_ctx = TransactionContext::new(&block.eth_tx, &block.geth_trace).unwrap();
+
+        // Generate step corresponding to SELFBALANCE
+        let mut step = ExecStep::new(
+            &block.geth_trace.struct_logs[0],
+            0,
+            test_builder.block_ctx.rwc,
+            0,
+        );
+        let mut state_ref = test_builder.state_ref(&mut tx, &mut tx_ctx, &mut step);
+
+        let callee_address = block.eth_tx.to.unwrap();
+        let self_balance = state_ref.sdb.get_account(&callee_address).1.balance;
+
+        // CallContext read for callee_address
+        state_ref.push_op(
+            RW::READ,
+            CallContextOp {
+                call_id: state_ref.call().call_id,
+                field: CallContextField::CalleeAddress,
+                value: callee_address.to_word(),
+            },
+        );
+
+        // Account read for balance of callee_address
+        state_ref.push_op(
+            RW::READ,
+            AccountOp {
+                address: callee_address,
+                field: AccountField::Balance,
+                value: self_balance,
+                value_prev: self_balance,
+            },
+        );
+
+        // Add the Stack write
+        state_ref.push_stack_op(RW::WRITE, StackAddress::from(1024 - 1), self_balance);
+
+        tx.steps_mut().push(step);
+        test_builder.block.txs_mut().push(tx);
+
+        // Compare first step bus mapping instance
+        assert_eq!(
+            builder.block.txs()[0].steps()[0].bus_mapping_instance,
+            test_builder.block.txs()[0].steps()[0].bus_mapping_instance,
+        );
+
+        // Compare containers
+        assert_eq!(builder.block.container, test_builder.block.container);
+
+        Ok(())
+    }
+}

bus-mapping/src/operation/container.rs

@@ -64,7 +64,7 @@ impl OperationContainer {
         }
     }
 
-    /// Inserts an [`Operation`] into the  container returning a lightwheight
+    /// Inserts an [`Operation`] into the  container returning a lightweight
     /// reference to it in the form of an [`OperationRef`] which points to the
     /// location of the inserted operation inside the corresponding container
     /// vector.

zkevm-circuits/src/evm_circuit/execution.rs

@@ -39,6 +39,7 @@ mod mul;
 mod pc;
 mod pop;
 mod push;
+mod selfbalance;
 mod signed_comparator;
 mod signextend;
 mod stop;
@@ -68,6 +69,7 @@ use mul::MulGadget;
 use pc::PcGadget;
 use pop::PopGadget;
 use push::PushGadget;
+use selfbalance::SelfbalanceGadget;
 use signed_comparator::SignedComparatorGadget;
 use signextend::SignextendGadget;
 use stop::StopGadget;
@@ -126,6 +128,7 @@ pub(crate) struct ExecutionConfig<F> {
     msize_gadget: MsizeGadget<F>,
     coinbase_gadget: CoinbaseGadget<F>,
     timestamp_gadget: TimestampGadget<F>,
+    selfbalance_gadget: SelfbalanceGadget<F>,
 }
 
 impl<F: FieldExt> ExecutionConfig<F> {
@@ -252,6 +255,7 @@ impl<F: FieldExt> ExecutionConfig<F> {
             pc_gadget: configure_gadget!(),
             pop_gadget: configure_gadget!(),
             push_gadget: configure_gadget!(),
+            selfbalance_gadget: configure_gadget!(),
             signed_comparator_gadget: configure_gadget!(),
             signextend_gadget: configure_gadget!(),
             stop_gadget: configure_gadget!(),
@@ -513,6 +517,7 @@ impl<F: FieldExt> ExecutionConfig<F> {
             ExecutionState::TIMESTAMP => {
                 assign_exec_step!(self.timestamp_gadget)
             }
+            ExecutionState::SELFBALANCE => assign_exec_step!(self.selfbalance_gadget),
             ExecutionState::CALLDATACOPY => {
                 assign_exec_step!(self.calldatacopy_gadget)
             }

zkevm-circuits/src/evm_circuit/execution/selfbalance.rs

@@ -0,0 +1,100 @@
+use crate::{
+    evm_circuit::{
+        execution::ExecutionGadget,
+        step::ExecutionState,
+        table::{AccountFieldTag, CallContextFieldTag},
+        util::{
+            common_gadget::SameContextGadget,
+            constraint_builder::{ConstraintBuilder, StepStateTransition, Transition::Delta},
+            Cell, Word,
+        },
+        witness::{Block, Call, ExecStep, Transaction},
+    },
+    util::Expr,
+};
+use eth_types::{ToLittleEndian, ToScalar};
+use halo2::{arithmetic::FieldExt, circuit::Region, plonk::Error};
+
+#[derive(Clone, Debug)]
+pub(crate) struct SelfbalanceGadget<F> {
+    same_context: SameContextGadget<F>,
+    callee_address: Cell<F>,
+    self_balance: Cell<F>,
+}
+
+impl<F: FieldExt> ExecutionGadget<F> for SelfbalanceGadget<F> {
+    const NAME: &'static str = "SELFBALANCE";
+
+    const EXECUTION_STATE: ExecutionState = ExecutionState::SELFBALANCE;
+
+    fn configure(cb: &mut ConstraintBuilder<F>) -> Self {
+        let callee_address = cb.call_context(None, CallContextFieldTag::CalleeAddress);
+
+        let self_balance = cb.query_cell();
+        cb.account_read(
+            callee_address.expr(),
+            AccountFieldTag::Balance,
+            self_balance.expr(),
+        );
+
+        cb.stack_push(self_balance.expr());
+
+        let opcode = cb.query_cell();
+        let step_state_transition = StepStateTransition {
+            rw_counter: Delta(3.expr()),
+            program_counter: Delta(1.expr()),
+            stack_pointer: Delta((-1).expr()),
+            ..Default::default()
+        };
+        let same_context = SameContextGadget::construct(cb, opcode, step_state_transition, None);
+
+        Self {
+            same_context,
+            self_balance,
+            callee_address,
+        }
+    }
+
+    fn assign_exec_step(
+        &self,
+        region: &mut Region<'_, F>,
+        offset: usize,
+        block: &Block<F>,
+        _: &Transaction,
+        call: &Call,
+        step: &ExecStep,
+    ) -> Result<(), Error> {
+        self.same_context.assign_exec_step(region, offset, step)?;
+
+        self.callee_address
+            .assign(region, offset, call.callee_address.to_scalar())?;
+
+        let self_balance = block.rws[step.rw_indices[2]].stack_value();
+        self.self_balance.assign(
+            region,
+            offset,
+            Some(Word::random_linear_combine(
+                self_balance.to_le_bytes(),
+                block.randomness,
+            )),
+        )?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::test_util::run_test_circuits;
+    use eth_types::bytecode;
+
+    #[test]
+    fn selfbalance_gadget_test() {
+        let bytecode = bytecode! {
+            #[start]
+            SELFBALANCE
+            STOP
+        };
+        assert_eq!(run_test_circuits(bytecode), Ok(()));
+    }
+}

zkevm-circuits/src/evm_circuit/witness.rs

@@ -1024,6 +1024,7 @@ impl From<&bus_mapping::circuit_input_builder::ExecStep> for ExecutionState {
             OpcodeId::COINBASE => ExecutionState::COINBASE,
             OpcodeId::TIMESTAMP => ExecutionState::TIMESTAMP,
             OpcodeId::GAS => ExecutionState::GAS,
+            OpcodeId::SELFBALANCE => ExecutionState::SELFBALANCE,
             _ => unimplemented!("unimplemented opcode {:?}", step.op),
         }
     }