Implement TLOAD and TSTORE

bus-mapping/src/circuit_input_builder.rs

@@ -696,6 +696,7 @@ impl<'a> CircuitInputBuilder {
         log::trace!("gen_end_tx_ops");
         let end_tx_steps =
             gen_associated_steps(&mut self.state_ref(&mut tx, &mut tx_ctx), ExecState::EndTx)?;
+        self.sdb.clear_transient_storage();
         tx.steps_mut().extend(end_tx_steps);
 
         debug_assert_eq!(

bus-mapping/src/circuit_input_builder/input_state_ref.rs

@@ -1179,6 +1179,14 @@ impl<'a> CircuitInputStateRef<'a> {
                     None
                 }
             }
+            OperationRef(Target::TransientStorage, idx) => {
+                let operation = &self.block.container.transient_storage[*idx];
+                if operation.rw().is_write() && operation.reversible() {
+                    Some(OpEnum::TransientStorage(operation.op().reverse()))
+                } else {
+                    None
+                }
+            }
             OperationRef(Target::TxAccessListAccount, idx) => {
                 let operation = &self.block.container.tx_access_list_account[*idx];
                 if operation.rw().is_write() && operation.reversible() {
@@ -1221,6 +1229,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);
@@ -1694,6 +1706,7 @@ impl<'a> CircuitInputStateRef<'a> {
                     OpcodeId::RETURNDATACOPY => Some(ExecError::ReturnDataOutOfBounds),
                     // Break write protection (CALL with value will be handled below)
                     OpcodeId::SSTORE
+                    | OpcodeId::TSTORE
                     | OpcodeId::CREATE
                     | OpcodeId::CREATE2
                     | OpcodeId::SELFDESTRUCT

bus-mapping/src/evm/opcodes.rs

@@ -59,6 +59,8 @@ mod sstore;
 mod stackonlyop;
 mod stop;
 mod swap;
+mod tload;
+mod tstore;
 
 mod error_codestore;
 mod error_contract_address_collision;
@@ -113,6 +115,7 @@ use extcodecopy::Extcodecopy;
 use extcodehash::Extcodehash;
 use extcodesize::Extcodesize;
 use gasprice::GasPrice;
+// use invalid_tx::InvalidTx;
 use logs::Log;
 use mload::Mload;
 use mstore::Mstore;
@@ -127,6 +130,8 @@ use sstore::Sstore;
 use stackonlyop::StackPopOnlyOpcode;
 use stop::Stop;
 use swap::Swap;
+use tload::Tload;
+use tstore::Tstore;
 
 pub use sstore::calc_expected_tx_refund;
 
@@ -242,6 +247,8 @@ fn fn_gen_associated_ops(opcode_id: &OpcodeId) -> FnGenAssociatedOps {
         OpcodeId::MSIZE => Msize::gen_associated_ops,
         OpcodeId::GAS => Gas::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,

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

@@ -31,6 +31,7 @@ impl Opcode for ErrorWriteProtection {
         // assert op code can only be following codes
         assert!([
             OpcodeId::SSTORE,
+            OpcodeId::TSTORE,
             OpcodeId::CREATE,
             OpcodeId::CREATE2,
             OpcodeId::CALL,

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

@@ -0,0 +1,140 @@
+use super::Opcode;
+use crate::{
+    circuit_input_builder::{CircuitInputStateRef, ExecStep},
+    operation::{CallContextField, TransientStorageOp, RW},
+    Error,
+};
+use eth_types::{GethExecStep, ToWord, Word};
+
+/// Placeholder structure used to implement [`Opcode`] trait over it
+/// corresponding to the [`OpcodeId::TLOAD`](crate::evm::OpcodeId::TLOAD)
+/// `OpcodeId`.
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct Tload;
+
+impl Opcode for Tload {
+    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 call_id = state.call()?.call_id;
+        let contract_addr = state.call()?.address;
+
+        state.call_context_read(
+            &mut exec_step,
+            call_id,
+            CallContextField::TxId,
+            Word::from(state.tx_ctx.id()),
+        )?;
+
+        state.call_context_read(
+            &mut exec_step,
+            call_id,
+            CallContextField::CalleeAddress,
+            contract_addr.to_word(),
+        )?;
+
+        // First stack read
+        let key = state.stack_pop(&mut exec_step)?;
+        #[cfg(feature = "enable-stack")]
+        assert_eq!(key, geth_step.stack.last()?);
+
+        // Transient Storage read
+        let (_, &value) = state.sdb.get_transient_storage(&contract_addr, &key);
+        #[cfg(feature = "enable-stack")]
+        assert_eq!(
+            value,
+            geth_steps[1].stack.last()?,
+            "inconsistent tload: step proof {value:?}, result {:?} in contract {contract_addr:?}, key {key:?}", geth_steps[1].stack.last()?,
+        );
+
+        state.push_op(
+            &mut exec_step,
+            RW::READ,
+            TransientStorageOp::new(contract_addr, key, value, value, state.tx_ctx.id()),
+        )?;
+
+        // First stack write
+        state.stack_push(&mut exec_step, value)?;
+
+        Ok(vec![exec_step])
+    }
+}
+
+#[cfg(test)]
+mod tload_tests {
+    use super::*;
+    use crate::{circuit_input_builder::ExecState, mock::BlockData, operation::StackOp};
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+    };
+    use mock::{
+        test_ctx::{helpers::*, TestContext},
+        MOCK_ACCOUNTS,
+    };
+    use pretty_assertions::assert_eq;
+
+    #[test]
+    fn tload_opcode() {
+        let code = bytecode! {
+            // Load transient storage slot 0
+            PUSH1(0x00u64)
+            TLOAD
+            STOP
+        };
+        let expected_loaded_value = 0;
+
+        // 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 step = builder.block.txs()[0]
+            .steps()
+            .iter()
+            .find(|step| step.exec_state == ExecState::Op(OpcodeId::TLOAD))
+            .unwrap();
+
+        println!("{:?}", step.bus_mapping_instance);
+
+        assert_eq!(
+            [&builder.block.container.stack[step.bus_mapping_instance[2].as_usize()]]
+                .map(|operation| (operation.rw(), operation.op())),
+            [(
+                RW::READ,
+                &StackOp::new(1, StackAddress::from(1023), Word::from(0x0u32))
+            )]
+        );
+
+        let transient_storage_op =
+            &builder.block.container.transient_storage[step.bus_mapping_instance[3].as_usize()];
+        assert_eq!(
+            (transient_storage_op.rw(), transient_storage_op.op()),
+            (
+                RW::READ,
+                &TransientStorageOp::new(
+                    MOCK_ACCOUNTS[0],
+                    Word::from(0x0u32),
+                    Word::from(expected_loaded_value),
+                    Word::from(expected_loaded_value),
+                    1,
+                )
+            )
+        );
+    }
+}