From 01050195879550fa68d9abd3d669ffc7b658d3f5 Mon Sep 17 00:00:00 2001
From: Mason Liang <mason@scroll.io>
Date: Mon, 29 Apr 2024 22:10:32 -0400
Subject: [PATCH] Handle reverts and write protection failure for TSTORE in bus
 mapping

---
 bus-mapping/src/circuit_input_builder/input_state_ref.rs | 9 +++++++++
 bus-mapping/src/evm/opcodes/error_write_protection.rs    | 1 +
 bus-mapping/src/operation/container.rs                   | 7 +++++--
 zkevm-circuits/src/evm_circuit/execution/tstore.rs       | 2 +-
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/bus-mapping/src/circuit_input_builder/input_state_ref.rs b/bus-mapping/src/circuit_input_builder/input_state_ref.rs
index 17cdddcade..44e6365356 100644
--- a/bus-mapping/src/circuit_input_builder/input_state_ref.rs
+++ b/bus-mapping/src/circuit_input_builder/input_state_ref.rs
@@ -935,6 +935,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() {
@@ -1436,6 +1444,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
diff --git a/bus-mapping/src/evm/opcodes/error_write_protection.rs b/bus-mapping/src/evm/opcodes/error_write_protection.rs
index 65c3cb9eee..b5571f5215 100644
--- a/bus-mapping/src/evm/opcodes/error_write_protection.rs
+++ b/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,
diff --git a/bus-mapping/src/operation/container.rs b/bus-mapping/src/operation/container.rs
index 0e3551af9b..e4064b5820 100644
--- a/bus-mapping/src/operation/container.rs
+++ b/bus-mapping/src/operation/container.rs
@@ -124,8 +124,11 @@ impl OperationContainer {
                 OperationRef::from((Target::Storage, self.storage.len() - 1))
             }
             OpEnum::TransientStorage(op) => {
-                self.transient_storage
-                    .push(Operation::new(rwc, rwc_inner_chunk, rw, op));
+                self.transient_storage.push(if reversible {
+                    Operation::new_reversible(rwc, rwc_inner_chunk, rw, op)
+                } else {
+                    Operation::new(rwc, rwc_inner_chunk, rw, op)
+                });
                 OperationRef::from((Target::TransientStorage, self.transient_storage.len() - 1))
             }
             OpEnum::TxAccessListAccount(op) => {
diff --git a/zkevm-circuits/src/evm_circuit/execution/tstore.rs b/zkevm-circuits/src/evm_circuit/execution/tstore.rs
index d29ba4c7f9..d0a89397c2 100644
--- a/zkevm-circuits/src/evm_circuit/execution/tstore.rs
+++ b/zkevm-circuits/src/evm_circuit/execution/tstore.rs
@@ -198,7 +198,7 @@ mod test {
                 accs[0]
                     .address(MOCK_ACCOUNTS[0])
                     .balance(Word::from(10u64.pow(19)))
-                    .code(bytecode)
+                    .code(bytecode);
                 accs[1]
                     .address(MOCK_ACCOUNTS[1])
                     .balance(Word::from(10u64.pow(19)));