tstore in busmapping with test case

zemse · zemse · commit d28a696192c6 · 2024-04-04T11:55:27.000+05:30
diff --git a/bus-mapping/src/circuit_input_builder/input_state_ref.rs b/bus-mapping/src/circuit_input_builder/input_state_ref.rs
@@ -977,6 +977,10 @@ impl<'a> CircuitInputStateRef<'a> {
             OpEnum::Storage(op) => {
                 self.sdb.set_storage(&op.address, &op.key, &op.value);
             }
+            OpEnum::TransientStorage(op) => {
+                self.sdb
+                    .set_transient_storage(&op.address, &op.key, &op.value)
+            }
             OpEnum::TxAccessListAccount(op) => {
                 if !op.is_warm_prev && op.is_warm {
                     self.sdb.add_account_to_access_list(op.address);
diff --git a/bus-mapping/src/evm/opcodes.rs b/bus-mapping/src/evm/opcodes.rs
@@ -45,6 +45,7 @@ mod stackonlyop;
 mod stop;
 mod swap;
 mod tload;
+mod tstore;
 
 mod error_code_store;
 mod error_invalid_creation_code;
@@ -112,6 +113,7 @@ use stackonlyop::StackOnlyOpcode;
 use stop::Stop;
 use swap::Swap;
 use tload::Tload;
+use tstore::Tstore;
 
 #[cfg(any(feature = "test", test))]
 pub use crate::precompile::PrecompileCallArgs;
@@ -228,6 +230,7 @@ fn fn_gen_associated_ops(opcode_id: &OpcodeId) -> FnGenAssociatedOps {
         OpcodeId::GAS => StackOnlyOpcode::<0, 1>::gen_associated_ops,
         OpcodeId::JUMPDEST => Dummy::gen_associated_ops,
         OpcodeId::TLOAD => Tload::gen_associated_ops,
+        OpcodeId::TSTORE => Tstore::gen_associated_ops,
         OpcodeId::DUP1 => Dup::<1>::gen_associated_ops,
         OpcodeId::DUP2 => Dup::<2>::gen_associated_ops,
         OpcodeId::DUP3 => Dup::<3>::gen_associated_ops,
diff --git a/bus-mapping/src/evm/opcodes/tload.rs b/bus-mapping/src/evm/opcodes/tload.rs
@@ -18,7 +18,6 @@ impl Opcode for Tload {
         geth_steps: &[GethExecStep],
     ) -> Result<Vec<ExecStep>, Error> {
         let geth_step = &geth_steps[0];
-        let next_geth_step = &geth_steps[1];
         let mut exec_step = state.new_step(geth_step)?;
 
         let call_id = state.call()?.call_id;
@@ -59,11 +58,9 @@ impl Opcode for Tload {
         // Manage first stack read at latest stack position
         state.stack_read(&mut exec_step, stack_position, key)?;
 
-        // Storage read
-        let value = next_geth_step
-            .stack
-            .last()
-            .expect("No value in stack in TLOAD next step");
+        // Transient Storage read
+        let (_, value) = state.sdb.get_transient_storage(&contract_addr, &key);
+        let value = *value;
 
         state.push_op(
             &mut exec_step,
diff --git a/bus-mapping/src/evm/opcodes/tstore.rs b/bus-mapping/src/evm/opcodes/tstore.rs
@@ -0,0 +1,215 @@
+use super::Opcode;
+use crate::{
+    circuit_input_builder::{CircuitInputStateRef, ExecStep},
+    operation::{CallContextField, TransientStorageOp},
+    Error,
+};
+use eth_types::{GethExecStep, ToWord, Word};
+
+/// Placeholder structure used to implement [`Opcode`] trait over it
+/// corresponding to the [`OpcodeId::TSTORE`](crate::evm::OpcodeId::TSTORE)
+/// `OpcodeId`.
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct Tstore;
+
+impl Opcode for Tstore {
+    fn gen_associated_ops(
+        state: &mut CircuitInputStateRef,
+        geth_steps: &[GethExecStep],
+    ) -> Result<Vec<ExecStep>, Error> {
+        let geth_step = &geth_steps[0];
+        let mut exec_step = state.new_step(geth_step)?;
+
+        let contract_addr = state.call()?.address;
+
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::TxId,
+            Word::from(state.tx_ctx.id()),
+        )?;
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::IsStatic,
+            Word::from(state.call()?.is_static as u8),
+        )?;
+
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::RwCounterEndOfReversion,
+            Word::from(state.call()?.rw_counter_end_of_reversion),
+        )?;
+
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::IsPersistent,
+            Word::from(state.call()?.is_persistent as u8),
+        )?;
+
+        state.call_context_read(
+            &mut exec_step,
+            state.call()?.call_id,
+            CallContextField::CalleeAddress,
+            state.call()?.address.to_word(),
+        )?;
+
+        let key = geth_step.stack.nth_last(0)?;
+        let key_stack_position = geth_step.stack.nth_last_filled(0);
+        let value = geth_step.stack.nth_last(1)?;
+        let value_stack_position = geth_step.stack.nth_last_filled(1);
+
+        state.stack_read(&mut exec_step, key_stack_position, key)?;
+        state.stack_read(&mut exec_step, value_stack_position, value)?;
+
+        let (_, value_prev) = state.sdb.get_transient_storage(&contract_addr, &key);
+        let value_prev = *value_prev;
+
+        state.push_op_reversible(
+            &mut exec_step,
+            TransientStorageOp::new(
+                state.call()?.address,
+                key,
+                value,
+                value_prev,
+                state.tx_ctx.id(),
+            ),
+        )?;
+
+        Ok(vec![exec_step])
+    }
+}
+
+#[cfg(test)]
+mod tstore_tests {
+    use super::*;
+    use crate::{
+        circuit_input_builder::ExecState,
+        mock::BlockData,
+        operation::{CallContextOp, StackOp, RW},
+    };
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+        Word,
+    };
+    use mock::{test_ctx::helpers::tx_from_1_to_0, TestContext, MOCK_ACCOUNTS};
+    use pretty_assertions::assert_eq;
+
+    #[test]
+    fn tstore_opcode() {
+        let code = bytecode! {
+            // Write 0x6f to storage slot 0
+            PUSH1(0x6fu64)
+            PUSH1(0x00u64)
+            TSTORE
+            PUSH1(0x00u64)
+            TLOAD
+            STOP
+        };
+        let expected_prev_value = 0x00u64;
+
+        // Get the execution steps from the external tracer
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            |accs| {
+                accs[0]
+                    .address(MOCK_ACCOUNTS[0])
+                    .balance(Word::from(10u64.pow(19)))
+                    .code(code)
+                    .storage(vec![(0x00u64.into(), 0x6fu64.into())].into_iter());
+                accs[1]
+                    .address(MOCK_ACCOUNTS[1])
+                    .balance(Word::from(10u64.pow(19)));
+            },
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
+
+        let builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
+        let builder = builder
+            .handle_block(&block.eth_block, &block.geth_traces)
+            .unwrap();
+
+        let step = builder.block.txs()[0]
+            .steps()
+            .iter()
+            .rev() // find last tstore
+            .find(|step| step.exec_state == ExecState::Op(OpcodeId::TSTORE))
+            .unwrap();
+
+        assert_eq!(
+            [0, 1, 2, 3, 4]
+                .map(|idx| &builder.block.container.call_context
+                    [step.bus_mapping_instance[idx].as_usize()])
+                .map(|operation| (operation.rw(), operation.op())),
+            [
+                (
+                    RW::READ,
+                    &CallContextOp::new(1, CallContextField::TxId, Word::from(0x01)),
+                ),
+                (
+                    RW::READ,
+                    &CallContextOp::new(1, CallContextField::IsStatic, Word::from(0x00)),
+                ),
+                (
+                    RW::READ,
+                    &CallContextOp::new(
+                        1,
+                        CallContextField::RwCounterEndOfReversion,
+                        Word::from(0x00)
+                    ),
+                ),
+                (
+                    RW::READ,
+                    &CallContextOp::new(1, CallContextField::IsPersistent, Word::from(0x01)),
+                ),
+                (
+                    RW::READ,
+                    &CallContextOp::new(
+                        1,
+                        CallContextField::CalleeAddress,
+                        MOCK_ACCOUNTS[0].to_word(),
+                    ),
+                ),
+            ]
+        );
+
+        assert_eq!(
+            [5, 6]
+                .map(|idx| &builder.block.container.stack[step.bus_mapping_instance[idx].as_usize()])
+                .map(|operation| (operation.rw(), operation.op())),
+            [
+                (
+                    RW::READ,
+                    &StackOp::new(1, StackAddress::from(1022), Word::from(0x0u32))
+                ),
+                (
+                    RW::READ,
+                    &StackOp::new(1, StackAddress::from(1023), Word::from(0x6fu32))
+                ),
+            ]
+        );
+
+        let storage_op =
+            &builder.block.container.transient_storage[step.bus_mapping_instance[7].as_usize()];
+        assert_eq!(
+            (storage_op.rw(), storage_op.op()),
+            (
+                RW::WRITE,
+                &TransientStorageOp::new(
+                    MOCK_ACCOUNTS[0],
+                    Word::from(0x0u32),
+                    Word::from(0x6fu32),
+                    Word::from(expected_prev_value),
+                    1,
+                )
+            )
+        );
+    }
+}
diff --git a/bus-mapping/src/state_db.rs b/bus-mapping/src/state_db.rs
@@ -140,6 +140,8 @@ pub struct StateDB {
     // state before current transaction, to calculate gas cost for some opcodes like sstore.
     // So both dirty storage and committed storage are needed.
     dirty_storage: HashMap<(Address, Word), Word>,
+    // Transient storage, which is cleared after the transaction.
+    transient_storage: HashMap<(Address, Word), Word>,
     // Accounts that have been through `SELFDESTRUCT` under the situation that `is_persistent` is
     // `true`. These accounts will be reset once `commit_tx` is called.
     destructed_account: HashSet<Address>,
@@ -190,6 +192,17 @@ impl StateDB {
         }
     }
 
+    /// Get a reference to the transient storage value from [`Account`] at `addr`, at
+    /// `key`.  Returns false and a zero [`Word`] when the [`Account`] or `key`
+    /// wasn't found in the state.
+    /// Returns transient storage value, which is cleared after current tx
+    pub fn get_transient_storage(&self, addr: &Address, key: &Word) -> (bool, &Word) {
+        match self.transient_storage.get(&(*addr, *key)) {
+            Some(v) => (true, v),
+            None => (false, &VALUE_ZERO),
+        }
+    }
+
     /// Get a reference to the storage value from [`Account`] at `addr`, at
     /// `key`.  Returns false and a zero [`Word`] when the [`Account`] or `key`
     /// wasn't found in the state.
@@ -227,6 +240,12 @@ impl StateDB {
         self.dirty_storage.insert((*addr, *key), *value);
     }
 
+    /// Set transient storage value at `addr` and `key`.
+    /// Transient storage is cleared after transaction execution.
+    pub fn set_transient_storage(&mut self, addr: &Address, key: &Word, value: &Word) {
+        self.transient_storage.insert((*addr, *key), *value);
+    }
+
     /// Get nonce of account with `addr`.
     pub fn get_nonce(&self, addr: &Address) -> u64 {
         let (_, account) = self.get_account(addr);
