Transient Storage Support (#1797)

#1761

- [x] Geth utils update
- [x] TLOAD TSTORE in bus mapping
- [x] TLOAD TSTORE in zkevm circuits 

<!--
### Description

[_PR description_]

### Issue Link

[_link issue here_]

### Type of change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [ ] Refactor (no updates to logic)

### Contents

- [_item_]

### Rationale

[_design decisions and extended information_]

### How Has This Been Tested?

[_explanation_]

<hr>

## How to fill a PR description 

Please give a concise description of your PR.

The target readers could be future developers, reviewers, and auditors.
By reading your description, they should easily understand the changes
proposed in this pull request.

MUST: Reference the issue to resolve

### Single responsibility

Is RECOMMENDED to create single responsibility commits, but not
mandatory.

Anyway, you MUST enumerate the changes in a unitary way, e.g.

```
This PR contains:
- Cleanup of xxxx, yyyy
- Changed xxxx to yyyy in order to bla bla
- Added xxxx function to ...
- Refactored ....
```

### Design choices

RECOMMENDED to:
- What types of design choices did you face?
- What decisions you have made?
- Any valuable information that could help reviewers to think critically
-->

---------

Co-authored-by: Miha Stopar <miha.stopar@xlab.si>
Co-authored-by: Chih Cheng Liang <chihchengliang@gmail.com>

bus-mapping/src/circuit_input_builder.rs

@@ -460,6 +460,7 @@ impl<'a, C: CircuitsParams> CircuitInputBuilder<C> {
             // Generate EndTx step
             let end_tx_step =
                 gen_associated_steps(&mut self.state_ref(&mut tx, &mut tx_ctx), ExecState::EndTx)?;
+            self.sdb.clear_transient_storage();
             tx.steps_mut().push(end_tx_step.clone());
             (end_tx_step, last_call)
         } else if self.feature_config.invalid_tx {

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);

bus-mapping/src/evm/opcodes.rs

@@ -44,6 +44,8 @@ mod sstore;
 mod stackonlyop;
 mod stop;
 mod swap;
+mod tload;
+mod tstore;
 
 mod error_code_store;
 mod error_invalid_creation_code;
@@ -65,7 +67,6 @@ mod precompiles;
 #[cfg(test)]
 mod memory_expansion_test;
 
-use self::{invalid_tx::InvalidTx, sha3::Sha3};
 use address::Address;
 use balance::Balance;
 use begin_end_tx::BeginEndTx;
@@ -96,6 +97,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;
@@ -104,13 +106,16 @@ use return_revert::ReturnRevert;
 use returndatacopy::Returndatacopy;
 use returndatasize::Returndatasize;
 use selfbalance::Selfbalance;
+use sha3::Sha3;
 use sload::Sload;
 use sstore::Sstore;
 use stackonlyop::StackOnlyOpcode;
 use stop::Stop;
 use swap::Swap;
+use tload::Tload;
+use tstore::Tstore;
 
-#[cfg(feature = "test")]
+#[cfg(any(feature = "test", test))]
 pub use crate::precompile::PrecompileCallArgs;
 
 /// Generic opcode trait which defines the logic of the
@@ -224,6 +229,8 @@ fn fn_gen_associated_ops(opcode_id: &OpcodeId) -> FnGenAssociatedOps {
         OpcodeId::MSIZE => StackOnlyOpcode::<0, 1>::gen_associated_ops,
         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,

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

@@ -0,0 +1,136 @@
+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 = geth_step.stack.last()?;
+        let stack_position = geth_step.stack.last_filled();
+
+        // Manage first stack read at latest stack position
+        state.stack_read(&mut exec_step, stack_position, key)?;
+
+        // Transient Storage read
+        let (_, &value) = state.sdb.get_transient_storage(&contract_addr, &key);
+
+        state.push_op(
+            &mut exec_step,
+            RW::READ,
+            TransientStorageOp::new(contract_addr, key, value, value, state.tx_ctx.id()),
+        )?;
+
+        // First stack write
+        state.stack_write(&mut exec_step, stack_position, 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 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()
+            .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,
+                )
+            )
+        );
+    }
+}