Merge pull request #1102 from mir-protocol/modexp-memory-context-change

Modexp memory context change

evm/src/cpu/kernel/asm/bignum/add.asm

@@ -14,10 +14,10 @@ global add_bignum:
 add_loop:
     // stack: carry, i, a_cur_loc, b_cur_loc, retdest
     DUP4
-    %mload_kernel_general
+    %mload_current_general
     // stack: b[cur], carry, i, a_cur_loc, b_cur_loc, retdest
     DUP4
-    %mload_kernel_general
+    %mload_current_general
     // stack: a[cur], b[cur], carry, i, a_cur_loc, b_cur_loc, retdest
     ADD
     ADD
@@ -32,7 +32,7 @@ add_loop:
     // stack: c[cur] = (a[cur] + b[cur] + carry) % 2^128, carry_new = (a[cur] + b[cur] + carry) // 2^128, i, a_cur_loc, b_cur_loc, retdest
     DUP4
     // stack: a_cur_loc, c[cur], carry_new, i, a_cur_loc, b_cur_loc, retdest
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: carry_new, i, a_cur_loc, b_cur_loc, retdest
     SWAP2
     %increment

evm/src/cpu/kernel/asm/bignum/addmul.asm

@@ -14,7 +14,7 @@ addmul_loop:
     // stack: carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
     DUP4
     // stack: b_cur_loc, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: b[cur], carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
     DUP6
     // stack: val, b[cur], carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
@@ -31,7 +31,7 @@ addmul_loop:
     // stack: prod_lo = val * b[cur] % 2^128, prod_hi = (val * b[cur]) // 2^128, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
     DUP5
     // stack: a_cur_loc, prod_lo, prod_hi, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: a[cur], prod_lo, prod_hi, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
     DUP1
     // stack: a[cur], a[cur], prod_lo, prod_hi, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
@@ -75,7 +75,7 @@ addmul_loop:
     // stack: to_write, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
     DUP4
     // stack: a_cur_loc, to_write, carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: carry_limb, i, a_cur_loc, b_cur_loc, val, retdest
     SWAP1
     // stack: i, carry_limb, a_cur_loc, b_cur_loc, val, retdest

evm/src/cpu/kernel/asm/bignum/cmp.asm

@@ -32,9 +32,9 @@ ge_loop:
     DUP3
     DUP3
     // stack: a_i_loc, b_i_loc, i, a_i_loc, b_i_loc, retdest
-    %mload_kernel_general
+    %mload_current_general
     SWAP1
-    %mload_kernel_general
+    %mload_current_general
     SWAP1
     // stack: a[i], b[i], i, a_i_loc, b_i_loc, retdest
     %stack (vals: 2) -> (vals, vals)

evm/src/cpu/kernel/asm/bignum/isone.asm

@@ -10,7 +10,7 @@ global isone_bignum:
     // stack: len, start_loc, retdest
     DUP2
     // stack: start_loc, len, start_loc, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: start_val, len, start_loc, retdest
     %eq_const(1)
     %jumpi(starts_with_one)

evm/src/cpu/kernel/asm/bignum/iszero.asm

@@ -17,7 +17,7 @@ iszero_loop:
     // stack: cur_loc, end_loc, retdest
     DUP1
     // stack: cur_loc, cur_loc, end_loc, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: cur_val, cur_loc, end_loc, retdest
     %jumpi(neqzero)
     // stack: cur_loc, end_loc, retdest

evm/src/cpu/kernel/asm/bignum/modexp.asm

@@ -65,22 +65,22 @@ e_zero_return:
     DUP3
     DUP8
     // stack: s1, b_loc, len, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
-    %memcpy_kernel_general
+    %memcpy_current_general
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
 
     // We store the accumulated output value x_i in output_loc, starting with x_0=1.
     PUSH 1
     DUP6
     // stack: out_loc, 1, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5,  retdest
-    %mstore_kernel_general
+    %mstore_current_general
 
 modexp_loop:
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
 
     // y := e % 2
     DUP3
     // stack: e_loc, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: e_first, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
     %mod_const(2)
     // stack: y = e_first % 2 = e % 2, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
@@ -108,15 +108,15 @@ modexp_mul_return:
     DUP11
     DUP7
     // stack: out_loc, s5, len, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
-    %memcpy_kernel_general
+    %memcpy_current_general
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
 
     // Zero out scratch_2..scratch_5.
     DUP1
     %mul_const(8)
     DUP8
     // stack: s2, 8 * len, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
-    %clear_kernel_general
+    %clear_current_general
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
 
 modexp_y_0:
@@ -142,15 +142,15 @@ modexp_square_return:
     DUP11
     DUP8
     // stack: s1, s5, len, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
-    %memcpy_kernel_general
+    %memcpy_current_general
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
 
     // Zero out scratch_2..scratch_5.
     DUP1
     %mul_const(8)
     DUP8
     // stack: s2, 8 * len, len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
-    %clear_kernel_general
+    %clear_current_general
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
 
     // e //= 2 (with shr_bignum)
@@ -186,7 +186,7 @@ b_and_e_zero:
     // stack: len, b_loc, e_loc, m_loc, out_loc, s1, s2, s3, s4, s5, retdest
     PUSH 1
     DUP6
-    %mstore_kernel_general
+    %mstore_current_general
     %pop10
     // stack: retdest
     JUMP

evm/src/cpu/kernel/asm/bignum/modmul.asm

@@ -31,7 +31,7 @@ modmul_remainder_loop:
     DUP3
     ADD
     // stack: out_loc[i], PI, i, len, a_loc, b_loc, m_loc, out_loc, s1, s2, s3, retdest
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: i, len, a_loc, b_loc, m_loc, out_loc, s1, s2, s3, retdest
     %increment
     DUP2
@@ -79,7 +79,7 @@ modmul_quotient_loop:
     DUP3
     ADD
     // stack: s1[i], PI, i, 2*len, len, a_loc, b_loc, m_loc, out_loc, s1, s2, s3, retdest
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: i, 2*len, len, a_loc, b_loc, m_loc, out_loc, s1, s2, s3, retdest
     %increment
     DUP2
@@ -136,9 +136,9 @@ modmul_check_loop:
     // stack: n, i, j, retdest
     %stack (l, idx: 2) -> (idx, l, idx)
     // stack: i, j, n, i, j, retdest
-    %mload_kernel_general
+    %mload_current_general
     SWAP1
-    %mload_kernel_general
+    %mload_current_general
     SWAP1
     // stack: mem[i], mem[j], n, i, j, retdest
     %assert_eq

evm/src/cpu/kernel/asm/bignum/mul.asm

@@ -20,7 +20,7 @@ mul_loop:
     // stack: mul_addmul_return, n, len, a_start_loc, bi, output_cur, retdest
     DUP5
     // stack: bi, mul_addmul_return, n, len, a_start_loc, bi, output_cur, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: b[i], mul_addmul_return, n, len, a_start_loc, bi, output_cur, retdest, b
     DUP5
     // stack: a_start_loc, b[i], mul_addmul_return, n, len, a_start_loc, bi, output_cur, retdest, b
@@ -37,7 +37,7 @@ mul_addmul_return:
     // stack: len, output_cur, carry_limb, n, len, a_start_loc, bi, output_cur, retdest
     ADD
     // stack: output_cur + len, carry_limb, n, len, a_start_loc, bi, output_cur, retdest
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: n, len, a_start_loc, bi, output_cur, retdest
     %decrement
     // stack: n-1, len, a_start_loc, bi, output_cur, retdest

evm/src/cpu/kernel/asm/bignum/shr.asm

@@ -22,7 +22,7 @@ shr_loop:
     // stack: i, carry, start_loc, retdest
     DUP1
     // stack: i, i, carry, start_loc, retdest
-    %mload_kernel_general
+    %mload_current_general
     // stack: a[i], i, carry, start_loc, retdest
     DUP1
     // stack: a[i], a[i], i, carry, start_loc, retdest
@@ -40,7 +40,7 @@ shr_loop:
     // stack: carry << 127 | a[i] >> 1, i, new_carry, start_loc, retdest
     DUP2
     // stack: i, carry << 127 | a[i] >> 1, i, new_carry, start_loc, retdest
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: i, new_carry, start_loc, retdest
     DUP1
     // stack: i, i, new_carry, start_loc, retdest

evm/src/cpu/kernel/asm/bignum/util.asm

@@ -1,13 +1,15 @@
-%macro memcpy_kernel_general
+%macro memcpy_current_general
     // stack: dst, src, len
-    %stack (dst, src, len) -> (0, @SEGMENT_KERNEL_GENERAL, dst, 0, @SEGMENT_KERNEL_GENERAL, src, len, %%after)
+    GET_CONTEXT
+    %stack (context, dst, src, len) -> (context, @SEGMENT_KERNEL_GENERAL, dst, context, @SEGMENT_KERNEL_GENERAL, src, len, %%after)
     %jump(memcpy)
 %%after:
 %endmacro
 
-%macro clear_kernel_general
+%macro clear_current_general
     // stack: dst, len
-    %stack (dst, len) -> (0, @SEGMENT_KERNEL_GENERAL, dst, 0, len, %%after)
+    GET_CONTEXT
+    %stack (context, dst, len) -> (context, @SEGMENT_KERNEL_GENERAL, dst, 0, len, %%after)
     %jump(memset)
 %%after:
 %endmacro

evm/src/cpu/kernel/asm/core/precompiles/expmod.asm

@@ -61,7 +61,7 @@ store_limbs:
     // stack: offset, retdest, num_limbs, limb[num_limbs - 1], ..limb[0]
     %stack (offset, ret, num, limb) -> (offset, limb, offset, ret, num)
     // stack: offset, limb[num_limbs - 1], offset, retdest, num_limbs, limb[num_limbs - 2], ..limb[0]
-    %mstore_kernel_general
+    %mstore_current_general
     // stack: offset, retdest, num_limbs, limb[num_limbs - 2], ..limb[0]
     %increment
     SWAP2
@@ -230,7 +230,7 @@ l_E_prime_return:
     %stack (kexit_info, l: 4) -> (l, kexit_info)
     // stack: len, l_M, l_E, l_B, kexit_info
 
-    // Copy B to kernel general memory.
+    // Copy B to memory.
     // stack: len, l_M, l_E, l_B, kexit_info
     DUP1
     // stack: len, len, l_M, l_E, l_B, kexit_info
@@ -262,7 +262,7 @@ copy_b_len_zero:
     %pop3
 copy_b_end:
     
-    // Copy E to kernel general memory.
+    // Copy E to memory.
     // stack: len, l_M, l_E, l_B, kexit_info
     DUP1
     // stack: len, len, l_M, l_E, l_B, kexit_info
@@ -294,7 +294,7 @@ copy_e_len_zero:
     %pop3
 copy_e_end:
 
-    // Copy M to kernel general memory.
+    // Copy M to memory.
     // stack: len, l_M, l_E, l_B, kexit_info
     DUP1
     // stack: len, len, l_M, l_E, l_B, kexit_info
@@ -389,7 +389,7 @@ copy_m_end:
 expmod_contd:
     // stack: len, l_M, kexit_info
 
-    // Copy the result value from kernel general memory to the parent's return data.
+    // Copy the result value from memory to the parent's return data.
 
     // Store return data size: l_M (number of bytes).
     SWAP1
@@ -411,7 +411,7 @@ expmod_contd:
     DUP2
     ADD
     // stack: cur_address=out+l_M_128-1, end_address=out-1, l_M_128, l_M%16, kexit_info
-    DUP1 %mload_kernel_general
+    DUP1 %mload_current_general
     %stack (cur_limb, cur_address, end_address, l_M_128, l_M_mod16, kexit_info) ->
         (@SEGMENT_RETURNDATA, 0, cur_limb, l_M_mod16, cur_address, end_address, l_M_128, kexit_info)
     %mload_context_metadata(@CTX_METADATA_PARENT_CONTEXT)
@@ -425,7 +425,7 @@ expmod_store_loop:
     // stack: cur_address, offset, end_address, l_M_128, kexit_info
     DUP3 DUP2 EQ %jumpi(expmod_store_end)
     // stack: cur_address, offset, end_address, l_M_128, kexit_info
-    DUP1 %mload_kernel_general
+    DUP1 %mload_current_general
     %stack (cur_limb, cur_address, offset, end_address, l_M_128, kexit_info) ->
          (offset, cur_limb, cur_address, end_address, l_M_128, kexit_info)
     %stack (offset, cur_limb) -> (@SEGMENT_RETURNDATA, offset, cur_limb, 16)

evm/src/cpu/kernel/interpreter.rs

@@ -204,6 +204,13 @@ impl<'a> Interpreter<'a> {
             .collect()
     }
 
+    pub(crate) fn get_current_general_memory(&self) -> Vec<U256> {
+        self.generation_state.memory.contexts[self.context].segments
+            [Segment::KernelGeneral as usize]
+            .content
+            .clone()
+    }
+
     pub(crate) fn get_kernel_general_memory(&self) -> Vec<U256> {
         self.get_memory_segment(Segment::KernelGeneral)
     }
@@ -212,12 +219,10 @@ impl<'a> Interpreter<'a> {
         self.get_memory_segment_bytes(Segment::RlpRaw)
     }
 
-    pub(crate) fn set_memory_segment(&mut self, segment: Segment, memory: Vec<U256>) {
-        self.generation_state.memory.contexts[0].segments[segment as usize].content = memory;
-    }
-
-    pub(crate) fn set_kernel_general_memory(&mut self, memory: Vec<U256>) {
-        self.set_memory_segment(Segment::KernelGeneral, memory)
+    pub(crate) fn set_current_general_memory(&mut self, memory: Vec<U256>) {
+        self.generation_state.memory.contexts[self.context].segments
+            [Segment::KernelGeneral as usize]
+            .content = memory;
     }
 
     pub(crate) fn set_memory_segment_bytes(&mut self, segment: Segment, memory: Vec<u8>) {

evm/src/cpu/kernel/tests/bignum/mod.rs

@@ -100,10 +100,10 @@ fn run_test(fn_label: &str, memory: Vec<U256>, stack: Vec<U256>) -> Result<(Vec<
     initial_stack.reverse();
 
     let mut interpreter = Interpreter::new_with_kernel(fn_label, initial_stack);
-    interpreter.set_kernel_general_memory(memory);
+    interpreter.set_current_general_memory(memory);
     interpreter.run()?;
 
-    let new_memory = interpreter.get_kernel_general_memory();
+    let new_memory = interpreter.get_current_general_memory();
 
     Ok((new_memory, interpreter.stack().to_vec()))
 }