Prune child context in create and call faults (#747)

* Prune child context after create failure

* call_insufficient_balance fix

* Fix create failures and drain_all_but_one_64th_gas

* Revert "fix (#738)"

This reverts commit f8c5f9c.

* Remove unnecessary changes

* Reapply "fix (#738)"

This reverts commit 2d66aa7.

* Add warning when the last MemAfter is not empty

* Apply comment

* Apply comment

---------

Co-authored-by: Robin Salen <salenrobin@gmail.com>

evm_arithmetization/src/cpu/kernel/asm/core/call.asm

@@ -247,7 +247,9 @@ after_call_instruction_failed:
     %jump(after_call_instruction_contd)
 
 call_insufficient_balance:
-    %stack (new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size) ->
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
+    %prune_context
+    %stack (kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size) ->
         (callgas, kexit_info, 0)
     %shl_const(192) SWAP1 SUB
     // stack: kexit_info', 0

evm_arithmetization/src/cpu/kernel/asm/core/create.asm

@@ -123,8 +123,8 @@ run_constructor:
     // stack: new_ctx, address, kexit_info
 
     // All but 1/64 of the sender's remaining gas goes to the constructor.
-    SWAP2
-    // stack: kexit_info, address, new_ctx
+    %stack(new_ctx, address, kexit_info) -> (kexit_info, new_ctx, address, new_ctx)
+    // stack: kexit_info, new_ctx, address, new_ctx
     %drain_all_but_one_64th_gas
     %stack (kexit_info, drained_gas, address, new_ctx) -> (drained_gas, new_ctx, address, kexit_info)
     %set_new_ctx_gas_limit_no_check
@@ -211,22 +211,30 @@ nonce_overflow:
     %stack (sender, address, value, code_offset, code_len, kexit_info) -> (kexit_info, 0)
     EXIT_KERNEL
 
+// stack: new_ctx, address, kexit_info
 create_collision:
     %revert_checkpoint
     %mstore_context_metadata(@CTX_METADATA_RETURNDATA_SIZE, 0)
-    %stack (new_ctx, address, kexit_info) -> (kexit_info, 0)
+    // Collisions are checked when running the constructor and prior entering the new context
+    // (but after writing some values in the new context), contrary to the other checks here. 
+    // This is why we need to prune the new context.
+    %prune_context
+    %stack (address, kexit_info) -> (kexit_info, 0)
     EXIT_KERNEL
 
+// stack: new_ctx, leftover_gas, success, address, kexit_info
 create_first_byte_ef:
     %revert_checkpoint
     %stack (leftover_gas, success, address, kexit_info) -> (kexit_info, 0)
     EXIT_KERNEL
 
+// stack: code_size, new_ctx, leftover_gas, success, address, kexit_info
 create_code_too_large:
     %revert_checkpoint
     %stack (code_size, leftover_gas, success, address, kexit_info) -> (kexit_info, 0)
     EXIT_KERNEL
 
+// stack: code_size_cost, new_ctx, leftover_gas, success, address, kexit_info
 create_oog:
     %revert_checkpoint
     %mstore_context_metadata(@CTX_METADATA_RETURNDATA_SIZE, 0)

evm_arithmetization/src/cpu/kernel/asm/core/gas.asm

@@ -40,6 +40,31 @@ global sys_gas:
     // stack: kexit_info'
 %endmacro
 
+// Charge gas. Faults if we exceed the limit for the current context,
+// and prune context in case of an exception.
+%macro charge_gas_and_prune
+    // stack: gas, kexit_info, new_ctx, retdest
+    %shl_const(192)
+    ADD
+    // stack: kexit_info', new_ctx
+    %ctx_gas_limit
+    // stack: gas_limit, kexit_info', new_ctx
+    DUP2 %shr_const(192)
+    // stack: gas_used, gas_limit, kexit_info', new_ctx
+    GT
+    // stack: out_of_gas, kexit_info', new_ctx
+    %jumpi(fault_exception_and_prune)
+    // stack: kexit_info', new_ctx
+    SWAP1 POP
+%endmacro
+
+// Prunes previously created context before faulting.
+global fault_exception_and_prune:
+    // stack: kexit_info', new_ctx
+    SWAP1 %prune_context
+    // stack: kexit_info'
+    %jump(fault_exception)
+
 // Charge a constant amount of gas.
 %macro charge_gas_const(gas)
     // stack: kexit_info
@@ -88,13 +113,13 @@ global sys_gasprice:
 // Given the current kexit_info, drains all but one 64th of its remaining gas.
 // Returns how much gas was drained.
 %macro drain_all_but_one_64th_gas
-    // stack: kexit_info
+    // stack: kexit_info, new_ctx
     DUP1 %leftover_gas
-    // stack: leftover_gas, kexit_info
+    // stack: leftover_gas, kexit_info, new_ctx
     %all_but_one_64th
-    // stack: all_but_one_64th, kexit_info
-    %stack (all_but_one_64th, kexit_info) -> (all_but_one_64th, kexit_info, all_but_one_64th)
-    %charge_gas
+    // stack: all_but_one_64th, kexit_info, new_ctx
+    %stack (all_but_one_64th, kexit_info, new_ctx) -> (all_but_one_64th, kexit_info, new_ctx, all_but_one_64th)
+    %charge_gas_and_prune
     // stack: kexit_info, drained_gas
 %endmacro
 

evm_arithmetization/src/cpu/kernel/asm/core/util.asm

@@ -121,4 +121,39 @@
     // stack: addr
     PUSH @U256_MAX
     MSTORE_GENERAL
-%endmacro
+%endmacro
+
+// Adds stale_ctx to the list of stale contexts. You need to return to a previous, older context with
+// a SET_CONTEXT instruction. By assumption, stale_ctx is greater than the current context.
+global prune_context:
+    // stack: stale_ctx, retdest
+    GET_CONTEXT
+    // stack: curr_ctx, stale_ctx, retdest
+    // When we go to stale_ctx, we want its stack to contain curr_ctx so that we can immediately
+    // call SET_CONTEXT. For that, we need a stack length of 1, and store curr_ctx in Segment::Stack[0].
+    PUSH @SEGMENT_STACK
+    DUP3 ADD
+    // stack: stale_ctx_stack_addr, curr_ctx, stale_ctx, retdest
+    DUP2
+    // stack: curr_ctx, stale_ctx_stack_addr, curr_ctx, stale_ctx, retdest
+    MSTORE_GENERAL
+    // stack: curr_ctx, stale_ctx, retdest
+    PUSH @CTX_METADATA_STACK_SIZE
+    DUP3 ADD
+    // stack: stale_ctx_stack_size_addr, curr_ctx, stale_ctx, retdest
+    PUSH 1
+    MSTORE_GENERAL
+    // stack: curr_ctx, stale_ctx, retdest
+    POP
+    SET_CONTEXT
+    // We're now in stale_ctx, with stack: curr_ctx, retdest
+    %set_and_prune_ctx
+    // We're now in curr_ctx, with stack: retdest
+    JUMP
+
+%macro prune_context
+    // stack: stale_ctx
+    %stack (stale_ctx) -> (stale_ctx, %%after)
+    %jump(prune_context)
+%%after:
+%endmacro

evm_arithmetization/src/cpu/kernel/asm/memory/syscalls.asm

@@ -391,30 +391,3 @@ mcopy_empty:
     %jump(memcpy_bytes)
 %endmacro
 
-// Adds stale_ctx to the list of stale contexts. You need to return to a previous, older context with
-// a SET_CONTEXT instruction. By assumption, stale_ctx is greater than the current context.
-%macro prune_context
-    // stack: stale_ctx
-    GET_CONTEXT
-    // stack: curr_ctx, stale_ctx
-    // When we go to stale_ctx, we want its stack to contain curr_ctx so that we can immediately
-    // call SET_CONTEXT. For that, we need a stack length of 1, and store curr_ctx in Segment::Stack[0].
-    PUSH @SEGMENT_STACK
-    DUP3 ADD
-    // stack: stale_ctx_stack_addr, curr_ctx, stale_ctx
-    DUP2
-    // stack: curr_ctx, stale_ctx_stack_addr, curr_ctx, stale_ctx
-    MSTORE_GENERAL
-    // stack: curr_ctx, stale_ctx
-    PUSH @CTX_METADATA_STACK_SIZE
-    DUP3 ADD
-    // stack: stale_ctx_stack_size_addr, curr_ctx, stale_ctx
-    PUSH 1
-    MSTORE_GENERAL
-    // stack: curr_ctx, stale_ctx
-    POP
-    SET_CONTEXT
-    // We're now in stale_ctx, with stack: curr_ctx
-    %set_and_prune_ctx
-    // We're now in curr_ctx, with an empty stack.
-%endmacro

evm_arithmetization/src/generation/mod.rs

@@ -620,6 +620,11 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         )
     );
 
+    let is_last_segment =
+        segment_data.registers_after.program_counter == KERNEL.global_labels["halt"];
+    if is_last_segment && final_len != 0 {
+        log::warn!("This is the last segment, but MemoryAfter is still not empty!");
+    }
     if final_len == 0 && OPTIONAL_TABLE_INDICES.contains(&MemAfter) {
         log::debug!("MemAfter table not in use");
         table_in_use[*MemAfter] = false;