Circuit for opcode BALANCE

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

@@ -1,8 +1,9 @@
 use crate::circuit_input_builder::{CircuitInputStateRef, ExecStep};
 use crate::evm::Opcode;
-use crate::operation::{TxAccessListAccountOp, RW};
+use crate::operation::{AccountField, CallContextField, TxAccessListAccountOp, RW};
+use crate::state_db::Account;
 use crate::Error;
-use eth_types::{GethExecStep, ToAddress, ToWord};
+use eth_types::{GethExecStep, ToAddress, ToWord, Word, U256};
 
 #[derive(Debug, Copy, Clone)]
 pub(crate) struct Balance;
@@ -12,16 +13,39 @@ impl Opcode for Balance {
         state: &mut CircuitInputStateRef,
         geth_steps: &[GethExecStep],
     ) -> Result<Vec<ExecStep>, Error> {
-        // TODO: finish this, only access list part is done
         let geth_step = &geth_steps[0];
         let mut exec_step = state.new_step(geth_step)?;
 
-        let address = geth_steps[0].stack.last()?.to_address();
+        // Read account address from stack.
+        let address = geth_step.stack.last()?.to_address();
         state.stack_read(
             &mut exec_step,
             geth_step.stack.last_filled(),
             address.to_word(),
         )?;
+
+        // Read transaction ID, rw_counter_end_of_reversion, and is_persistent
+        // from call context.
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::TxId,
+            U256::from(state.tx_ctx.id()),
+        );
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::RwCounterEndOfReversion,
+            U256::from(state.call()?.rw_counter_end_of_reversion as u64),
+        );
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::IsPersistent,
+            U256::from(state.call()?.is_persistent as u64),
+        );
+
+        // Update transaction access list for account address.
         let is_warm = state.sdb.check_account_in_access_list(&address);
         state.push_op_reversible(
             &mut exec_step,
@@ -34,6 +58,27 @@ impl Opcode for Balance {
             },
         )?;
 
+        // Read account balance.
+        let (exists, &Account { balance, .. }) = state.sdb.get_account(&address);
+        if exists {
+            state.account_read(
+                &mut exec_step,
+                address,
+                AccountField::Balance,
+                balance,
+                balance,
+            )?;
+        } else {
+            state.account_read(
+                &mut exec_step,
+                address,
+                AccountField::NonExisting,
+                Word::zero(),
+                Word::zero(),
+            )?;
+        };
+
+        // Write the BALANCE result to stack.
         state.stack_write(
             &mut exec_step,
             geth_steps[1].stack.nth_last_filled(0),
@@ -47,86 +92,77 @@ impl Opcode for Balance {
 #[cfg(test)]
 mod balance_tests {
     use super::*;
-    use crate::{
-        circuit_input_builder::ExecState,
-        mock::BlockData,
-        operation::{StackOp, RW},
-    };
-    use eth_types::{
-        address, bytecode,
-        evm_types::{OpcodeId, StackAddress},
-        geth_types::GethData,
-        Word,
-    };
-    use mock::eth;
-    use mock::test_ctx::{helpers::*, TestContext};
+    use crate::circuit_input_builder::ExecState;
+    use crate::mock::BlockData;
+    use crate::operation::{AccountOp, CallContextOp, StackOp};
+    use eth_types::evm_types::{OpcodeId, StackAddress};
+    use eth_types::geth_types::GethData;
+    use eth_types::{address, bytecode, Bytecode, Word, U256};
+    use mock::TestContext;
     use pretty_assertions::assert_eq;
 
-    // If the given account doesn't exist, it will push 0 onto the stack instead.
     #[test]
-    fn test_balance_of_non_exists_address() {
-        let address = address!("0x0000000000001111111111111111111111111111");
-
-        let code = bytecode! {
-            PUSH32(address.to_word())
-            BALANCE
-            STOP
-        };
-
-        // Get the execution steps from the external tracer
-        let block: GethData = TestContext::<2, 1>::new(
-            None,
-            account_0_code_account_1_no_code(code),
-            tx_from_1_to_0,
-            |block, _tx| block.number(0xcafeu64),
-        )
-        .unwrap()
-        .into();
-
-        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
-        builder
-            .handle_block(&block.eth_block, &block.geth_traces)
-            .unwrap();
-        let call_id = builder.block.txs()[0].calls()[0].call_id;
-
-        let step = builder.block.txs()[0]
-            .steps()
-            .iter()
-            .find(|step| step.exec_state == ExecState::Op(OpcodeId::BALANCE))
-            .unwrap();
-
-        let operation = &builder.block.container.stack[step.bus_mapping_instance[1].as_usize()];
+    fn test_balance_of_non_existing_address() {
+        test_ok(false, false);
+    }
 
-        assert_eq!(
-            { (operation.rw(), operation.op()) },
-            (
-                RW::WRITE,
-                &StackOp::new(call_id, StackAddress::from(1023), Word::zero())
-            )
-        );
+    #[test]
+    fn test_balance_of_cold_address() {
+        test_ok(true, false);
     }
 
     #[test]
-    fn test_balance_of_exists_address() {
-        let (addr_a, addr_b) = (mock::MOCK_ACCOUNTS[0], mock::MOCK_ACCOUNTS[1]);
+    fn test_balance_of_warm_address() {
+        test_ok(true, true);
+    }
+
+    fn test_ok(exists: bool, is_warm: bool) {
+        let address = address!("0xaabbccddee000000000000000000000000000000");
 
-        let code = bytecode! {
-            PUSH32(addr_a.to_word())
+        // Pop balance first for warm account.
+        let mut code = Bytecode::default();
+        if is_warm {
+            code.append(&bytecode! {
+                PUSH20(address.to_word())
+                BALANCE
+                POP
+            });
+        }
+        code.append(&bytecode! {
+            PUSH20(address.to_word())
             BALANCE
             STOP
+        });
+
+        let balance = if exists {
+            Word::from(800u64)
+        } else {
+            Word::zero()
         };
 
-        // Get the execution steps from the external tracer
-        let block: GethData = TestContext::<2, 1>::new(
+        // Get the execution steps from the external tracer.
+        let block: GethData = TestContext::<3, 1>::new(
             None,
             |accs| {
-                accs[0].address(addr_a).balance(eth(10)).code(code);
-                accs[1].address(addr_b).balance(eth(10));
+                accs[0]
+                    .address(address!("0x0000000000000000000000000000000000000010"))
+                    .balance(Word::from(1u64 << 20))
+                    .code(code.clone());
+                if exists {
+                    accs[1].address(address).balance(balance);
+                } else {
+                    accs[1]
+                        .address(address!("0x0000000000000000000000000000000000000020"))
+                        .balance(Word::from(1u64 << 20));
+                }
+                accs[2]
+                    .address(address!("0x0000000000000000000000000000000000cafe01"))
+                    .balance(Word::from(1u64 << 20));
             },
             |mut txs, accs| {
-                txs[0].from(accs[1].address).to(accs[0].address);
+                txs[0].to(accs[0].address).from(accs[2].address);
             },
-            |block, _tx| block,
+            |block, _tx| block.number(0xcafeu64),
         )
         .unwrap()
         .into();
@@ -136,37 +172,105 @@ mod balance_tests {
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();
 
-        let step = builder.block.txs()[0]
+        // Check if account address is in access list as a result of bus mapping.
+        assert!(builder.sdb.add_account_to_access_list(address));
+
+        let tx_id = 1;
+        let transaction = &builder.block.txs()[tx_id - 1];
+        let call_id = transaction.calls()[0].call_id;
+
+        let indices = transaction
             .steps()
             .iter()
-            .find(|step| step.exec_state == ExecState::Op(OpcodeId::BALANCE))
-            .unwrap();
+            .filter(|step| step.exec_state == ExecState::Op(OpcodeId::BALANCE))
+            .last()
+            .unwrap()
+            .bus_mapping_instance
+            .clone();
 
-        let call_id = builder.block.txs()[0].calls()[0].call_id;
-        let balance_a = builder.sdb.get_account(&addr_a).1.balance;
+        let container = builder.block.container;
 
-        assert_eq!(addr_a, block.eth_block.transactions[0].to.unwrap());
+        let operation = &container.stack[indices[0].as_usize()];
+        assert_eq!(operation.rw(), RW::READ);
         assert_eq!(
-            {
-                let operation =
-                    &builder.block.container.stack[step.bus_mapping_instance[0].as_usize()];
-                (operation.rw(), operation.op())
-            },
-            (
-                RW::READ,
-                &StackOp::new(call_id, StackAddress::from(1023), addr_a.to_word())
-            )
+            operation.op(),
+            &StackOp {
+                call_id,
+                address: StackAddress::from(1023u32),
+                value: address.to_word()
+            }
         );
+
+        let operation = &container.call_context[indices[1].as_usize()];
+        assert_eq!(operation.rw(), RW::READ);
         assert_eq!(
-            {
-                let operation =
-                    &builder.block.container.stack[step.bus_mapping_instance[1].as_usize()];
-                (operation.rw(), operation.op())
-            },
-            (
-                RW::WRITE,
-                &StackOp::new(call_id, StackAddress::from(1023), balance_a)
-            )
+            operation.op(),
+            &CallContextOp {
+                call_id,
+                field: CallContextField::TxId,
+                value: tx_id.into()
+            }
+        );
+
+        let operation = &container.call_context[indices[2].as_usize()];
+        assert_eq!(operation.rw(), RW::READ);
+        assert_eq!(
+            operation.op(),
+            &CallContextOp {
+                call_id,
+                field: CallContextField::RwCounterEndOfReversion,
+                value: U256::zero()
+            }
+        );
+
+        let operation = &container.call_context[indices[3].as_usize()];
+        assert_eq!(operation.rw(), RW::READ);
+        assert_eq!(
+            operation.op(),
+            &CallContextOp {
+                call_id,
+                field: CallContextField::IsPersistent,
+                value: U256::one()
+            }
+        );
+
+        let operation = &container.tx_access_list_account[indices[4].as_usize()];
+        assert_eq!(operation.rw(), RW::WRITE);
+        assert_eq!(
+            operation.op(),
+            &TxAccessListAccountOp {
+                tx_id,
+                address,
+                is_warm: true,
+                is_warm_prev: is_warm
+            }
+        );
+
+        let operation = &container.account[indices[5].as_usize()];
+        assert_eq!(operation.rw(), RW::READ);
+        assert_eq!(
+            operation.op(),
+            &AccountOp {
+                address,
+                field: if exists {
+                    AccountField::Balance
+                } else {
+                    AccountField::NonExisting
+                },
+                value: balance,
+                value_prev: balance,
+            }
+        );
+
+        let operation = &container.stack[indices[6].as_usize()];
+        assert_eq!(operation.rw(), RW::WRITE);
+        assert_eq!(
+            operation.op(),
+            &StackOp {
+                call_id,
+                address: 1023u32.into(),
+                value: balance,
+            }
         );
     }
 }