bus-mapping: fix return_data (privacy-scaling-explorations#874)

* bus-mapping: fix return_data * add zero check Co-authored-by: adria0.eth <[email protected]> Co-authored-by: Han <[email protected]>
poilqa · Nov 4, 2022 · b87a9d2 · b87a9d2
1 parent a7efd0d
commit b87a9d2
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 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
@@ -796,6 +796,29 @@ impl<'a> CircuitInputStateRef<'a> {
     /// Handle 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> {
+        // handle return_data
+        if !self.call()?.is_root {
+            match step.op {
+                OpcodeId::RETURN | OpcodeId::REVERT => {
+                    let offset = step.stack.nth_last(0)?.as_usize();
+                    let length = step.stack.nth_last(1)?.as_usize();
+                    // TODO: Try to get rid of clone.
+                    // At the moment it conflicts with `call_ctx` and `caller_ctx`.
+                    let callee_memory = self.call_ctx()?.memory.clone();
+                    let caller_ctx = self.caller_ctx_mut()?;
+                    caller_ctx.return_data.resize(length, 0);
+                    if length != 0 {
+                        caller_ctx.return_data[0..length]
+                            .copy_from_slice(&callee_memory.0[offset..offset + length]);
+                    }
+                }
+                _ => {
+                    let caller_ctx = self.caller_ctx_mut()?;
+                    caller_ctx.return_data.truncate(0);
+                }
+            }
+        }
+
         let call = self.call()?.clone();
         let call_ctx = self.call_ctx()?;
 

diff --git a/bus-mapping/src/evm/opcodes/return.rs b/bus-mapping/src/evm/opcodes/return.rs
@@ -94,9 +94,6 @@ impl Opcode for Return {
 
                 caller_ctx.memory.0[return_offset..return_offset + copy_length]
                     .copy_from_slice(&callee_memory.0[offset..offset + copy_length]);
-                caller_ctx.return_data.resize(length, 0);
-                caller_ctx.return_data[0..copy_length]
-                    .copy_from_slice(&callee_memory.0[offset..offset + copy_length]);
 
                 handle_copy(
                     state,