feat/privacy-scaling-explorations#1186 Unify restore_ctx and handle_r…

…eturn in bus_mapping (privacy-scaling-explorations#1342)

### Description

- unify handle_restore_context and gen_restore_context_ops

### Issue Link

fix privacy-scaling-explorations#1186 

### Type of change

- [x] refactor (non-breaking change which adds functionality)

### Contents

- merge `gen_restore_context_ops` into `handle_restore_context`

bus-mapping/src/circuit_input_builder/input_state_ref.rs

@@ -901,9 +901,19 @@ impl<'a> CircuitInputStateRef<'a> {
         }
     }
 
-    /// Handle a return step caused by any opcode that causes a return to the
+    /// Handle a restore and a return step caused by any opcode that causes a return to the
     /// previous call context.
-    pub fn handle_return(&mut self, step: &GethExecStep) -> Result<(), Error> {
+    pub fn handle_return(
+        &mut self,
+        exec_step: &mut ExecStep,
+        geth_steps: &[GethExecStep],
+        need_restore: bool,
+    ) -> Result<(), Error> {
+        if need_restore {
+            self.handle_restore_context(exec_step, geth_steps)?;
+        }
+
+        let step = &geth_steps[0];
         // handle return_data
         let (return_data_offset, return_data_length) = {
             if !self.call()?.is_root {
@@ -986,25 +996,53 @@ impl<'a> CircuitInputStateRef<'a> {
     }
 
     /// Bus mapping for the RestoreContextGadget as used in RETURN.
-    // TODO: unify this with restore context bus mapping for STOP.
-    // TODO: unify this with the `handle return function above.`
     pub fn handle_restore_context(
         &mut self,
-        steps: &[GethExecStep],
         exec_step: &mut ExecStep,
+        steps: &[GethExecStep],
     ) -> Result<(), Error> {
         let call = self.call()?.clone();
+        let geth_step = steps
+            .get(0)
+            .ok_or(Error::InternalError("invalid index 0"))?;
+        let is_return_revert = geth_step.op == OpcodeId::REVERT || geth_step.op == OpcodeId::RETURN;
+
+        if !is_return_revert && !call.is_success {
+            // add call failure ops for exception cases
+            self.call_context_read(
+                exec_step,
+                call.call_id,
+                CallContextField::IsSuccess,
+                0u64.into(),
+            );
+
+            // Even call.rw_counter_end_of_reversion is zero for now, it will set in
+            // set_value_ops_call_context_rwc_eor later
+            // if call fails, no matter root or internal, read RwCounterEndOfReversion for
+            // circuit constraint.
+            self.call_context_read(
+                exec_step,
+                call.call_id,
+                CallContextField::RwCounterEndOfReversion,
+                call.rw_counter_end_of_reversion.into(),
+            );
+
+            if call.is_root {
+                return Ok(());
+            }
+        }
+
         let caller = self.caller()?.clone();
+        let geth_step_next = steps
+            .get(1)
+            .ok_or(Error::InternalError("invalid index 1"))?;
         self.call_context_read(
             exec_step,
             call.call_id,
             CallContextField::CallerId,
             caller.call_id.into(),
         );
 
-        let geth_step = &steps[0];
-        let geth_step_next = &steps[1];
-
         let [last_callee_return_data_offset, last_callee_return_data_length] = match geth_step.op {
             OpcodeId::STOP => [Word::zero(); 2],
             OpcodeId::REVERT | OpcodeId::RETURN => {
@@ -1018,7 +1056,7 @@ impl<'a> CircuitInputStateRef<'a> {
                     [offset, length]
                 }
             }
-            _ => unreachable!(),
+            _ => [Word::zero(), Word::zero()],
         };
 
         let curr_memory_word_size = (exec_step.memory_size as u64) / 32;
@@ -1041,7 +1079,11 @@ impl<'a> CircuitInputStateRef<'a> {
         };
         let gas_refund = geth_step.gas.0 - memory_expansion_gas_cost - code_deposit_cost;
 
-        let caller_gas_left = geth_step_next.gas.0 - gas_refund;
+        let caller_gas_left = if is_return_revert || call.is_success {
+            geth_step_next.gas.0 - gas_refund
+        } else {
+            geth_step_next.gas.0
+        };
 
         for (field, value) in [
             (CallContextField::IsRoot, (caller.is_root as u64).into()),
@@ -1306,89 +1348,4 @@ impl<'a> CircuitInputStateRef<'a> {
         }
         Ok(())
     }
-
-    /// gen bus mapping operations for context restore purpose
-    pub(crate) fn gen_restore_context_ops(
-        &mut self,
-        exec_step: &mut ExecStep,
-        geth_steps: &[GethExecStep],
-    ) -> Result<(), Error> {
-        let geth_step = &geth_steps[0];
-        let call = self.call()?.clone();
-        if !call.is_success {
-            // add call failure ops for exception cases
-            self.call_context_read(
-                exec_step,
-                call.call_id,
-                CallContextField::IsSuccess,
-                0u64.into(),
-            );
-
-            // Even call.rw_counter_end_of_reversion is zero for now, it will set in
-            // set_value_ops_call_context_rwc_eor later
-            // if call fails, no matter root or internal, read RwCounterEndOfReversion for
-            // circuit constraint.
-            self.call_context_read(
-                exec_step,
-                call.call_id,
-                CallContextField::RwCounterEndOfReversion,
-                call.rw_counter_end_of_reversion.into(),
-            );
-
-            if call.is_root {
-                return Ok(());
-            }
-        }
-
-        let caller = self.caller()?.clone();
-        self.call_context_read(
-            exec_step,
-            call.call_id,
-            CallContextField::CallerId,
-            caller.call_id.into(),
-        );
-
-        let geth_step_next = &geth_steps[1];
-        let caller_ctx = self.caller_ctx()?;
-        let caller_gas_left = if call.is_success {
-            geth_step_next.gas.0 - geth_step.gas.0
-        } else {
-            geth_step_next.gas.0
-        };
-
-        for (field, value) in [
-            (CallContextField::IsRoot, (caller.is_root as u64).into()),
-            (
-                CallContextField::IsCreate,
-                (caller.is_create() as u64).into(),
-            ),
-            (CallContextField::CodeHash, caller.code_hash.to_word()),
-            (CallContextField::ProgramCounter, geth_step_next.pc.0.into()),
-            (
-                CallContextField::StackPointer,
-                geth_step_next.stack.stack_pointer().0.into(),
-            ),
-            (CallContextField::GasLeft, caller_gas_left.into()),
-            (
-                CallContextField::MemorySize,
-                caller_ctx.memory.word_size().into(),
-            ),
-            (
-                CallContextField::ReversibleWriteCounter,
-                self.caller_ctx()?.reversible_write_counter.into(),
-            ),
-        ] {
-            self.call_context_read(exec_step, caller.call_id, field, value);
-        }
-
-        for (field, value) in [
-            (CallContextField::LastCalleeId, call.call_id.into()),
-            (CallContextField::LastCalleeReturnDataOffset, 0.into()),
-            (CallContextField::LastCalleeReturnDataLength, 0.into()),
-        ] {
-            self.call_context_write(exec_step, caller.call_id, field, value);
-        }
-
-        Ok(())
-    }
 }

bus-mapping/src/evm/opcodes.rs

@@ -331,17 +331,19 @@ pub fn gen_associated_ops(
         if let Some(fn_gen_error_ops) = fn_gen_error_state_associated_ops(&exec_error) {
             return fn_gen_error_ops(state, geth_steps);
         } else {
+            // For exceptions that fail to enter next call context, we need
+            // to restore call context of current caller
+            let mut need_restore = true;
+
             // For exceptions that already enter next call context, but fail immediately
             // (e.g. Depth, InsufficientBalance), we still need to parse the call.
             if geth_step.op.is_call_or_create() {
                 let call = state.parse_call(geth_step)?;
                 state.push_call(call);
-            // For exceptions that fail to enter next call context, we need
-            // to restore call context of current caller
-            } else {
-                state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
+                need_restore = false;
             }
-            state.handle_return(geth_step)?;
+
+            state.handle_return(&mut exec_step, geth_steps, need_restore)?;
             return Ok(vec![exec_step]);
         }
     }
@@ -701,6 +703,6 @@ fn dummy_gen_selfdestruct_ops(
         state.sdb.destruct_account(sender);
     }
 
-    state.handle_return(geth_step)?;
+    state.handle_return(&mut exec_step, geth_steps, false)?;
     Ok(vec![exec_step])
 }

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

@@ -229,7 +229,7 @@ impl<const N_ARGS: usize> Opcode for CallOpcode<N_ARGS> {
                 ] {
                     state.call_context_write(&mut exec_step, current_call.call_id, field, value);
                 }
-                state.handle_return(geth_step)?;
+                state.handle_return(&mut exec_step, geth_steps, false)?;
                 Ok(vec![exec_step])
             }
             // 3. Call to account with non-empty code.
@@ -306,7 +306,7 @@ impl<const N_ARGS: usize> Opcode for CallOpcode<N_ARGS> {
                 ] {
                     state.call_context_write(&mut exec_step, current_call.call_id, field, value);
                 }
-                state.handle_return(geth_step)?;
+                state.handle_return(&mut exec_step, geth_steps, false)?;
                 Ok(vec![exec_step])
             } //
         }

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

@@ -176,7 +176,7 @@ impl<const IS_CREATE2: bool> Opcode for DummyCreate<IS_CREATE2> {
 
         if call.code_hash == CodeDB::empty_code_hash() {
             // 1. Create with empty initcode.
-            state.handle_return(geth_step)?;
+            state.handle_return(&mut exec_step, geth_steps, false)?;
             Ok(vec![exec_step])
         } else {
             // 2. Create with non-empty initcode.

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

@@ -40,10 +40,9 @@ impl Opcode for InvalidJump {
                 condition,
             )?;
         }
-        // `IsSuccess` call context operation is added in gen_restore_context_ops
 
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
+        // `IsSuccess` call context operation is added in handle_return
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -90,8 +90,7 @@ impl Opcode for OOGCall {
             },
         );
 
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -30,9 +30,7 @@ impl Opcode for OOGExp {
             )?;
         }
 
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
-
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -43,9 +43,8 @@ impl Opcode for ErrorOOGLog {
             CallContextField::IsStatic,
             Word::from(state.call()?.is_static as u8),
         );
-        // common error handling
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
+
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -89,9 +89,7 @@ impl Opcode for OOGSloadSstore {
             );
         }
 
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
-
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -68,10 +68,8 @@ impl Opcode for ErrorReturnDataOutOfBound {
             return_data.len().into(),
         );
 
-        // `IsSuccess` call context operation is added in gen_restore_context_ops
-
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
+        // `IsSuccess` call context operation is added in handle_return
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -22,10 +22,7 @@ impl Opcode for ErrorSimple {
         let next_step = geth_steps.get(1);
         exec_step.error = state.get_step_err(geth_step, next_step).unwrap();
 
-        // handles all required steps
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
-
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -62,9 +62,8 @@ impl Opcode for ErrorWriteProtection {
             (current_call.is_static as u64).into(),
         );
 
-        // `IsSuccess` call context operation is added in gen_restore_context_ops
-        state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        state.handle_return(geth_step)?;
+        // `IsSuccess` call context operation is added in handle_return
+        state.handle_return(&mut exec_step, geth_steps, true)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -92,7 +92,7 @@ impl Opcode for ReturnRevert {
 
         // Case C in the specs.
         if !call.is_root {
-            state.handle_restore_context(steps, &mut exec_step)?;
+            state.handle_restore_context(&mut exec_step, steps)?;
         }
 
         // Case D in the specs.
@@ -132,7 +132,7 @@ impl Opcode for ReturnRevert {
             }
         }
 
-        state.handle_return(step)?;
+        state.handle_return(&mut exec_step, steps, false)?;
         Ok(vec![exec_step])
     }
 }

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

@@ -32,16 +32,7 @@ impl Opcode for Stop {
             1.into(),
         );
 
-        if !call.is_root {
-            // The following part corresponds to
-            // Instruction.step_state_transition_to_restored_context
-            // in python spec, and should be reusable among all expected halting opcodes or
-            // exceptions.
-            state.gen_restore_context_ops(&mut exec_step, geth_steps)?;
-        }
-
-        state.handle_return(geth_step)?;
-
+        state.handle_return(&mut exec_step, geth_steps, !call.is_root)?;
         Ok(vec![exec_step])
     }
 }