Change context to current context for BN precompiles (#1428)

* Change context to current for BN precompiles

* Rename segments

* rustfmt

evm/src/cpu/kernel/asm/curve/bn254/curve_arithmetic/curve_mul.asm

@@ -27,12 +27,12 @@ bn_mul_valid_point:
 bn_mul_after_glv:
     // stack: bneg, a, b, x, y, bn_msm, bn_mul_end, retdest
     // Store bneg at this (otherwise unused) location. Will be used later in the MSM.
-    %mstore_kernel(@SEGMENT_KERNEL_BN_TABLE_Q, @BN_BNEG_LOC)
+    %mstore_current(@SEGMENT_BN_TABLE_Q, @BN_BNEG_LOC)
     // stack: a, b, x, y, bn_msm, bn_mul_end, retdest
-    PUSH bn_mul_after_a SWAP1 PUSH @SEGMENT_KERNEL_BN_WNAF_A PUSH @BN_SCALAR %jump(wnaf)
+    PUSH bn_mul_after_a SWAP1 PUSH @SEGMENT_BN_WNAF_A PUSH @BN_SCALAR %jump(wnaf)
 bn_mul_after_a:
     // stack: b, x, y, bn_msm, bn_mul_end, retdest
-    PUSH bn_mul_after_b SWAP1 PUSH @SEGMENT_KERNEL_BN_WNAF_B PUSH @BN_SCALAR %jump(wnaf)
+    PUSH bn_mul_after_b SWAP1 PUSH @SEGMENT_BN_WNAF_B PUSH @BN_SCALAR %jump(wnaf)
 bn_mul_after_b:
     // stack: x, y, bn_msm, bn_mul_end, retdest
     %jump(bn_precompute_table)

evm/src/cpu/kernel/asm/curve/bn254/curve_arithmetic/msm.asm

@@ -42,31 +42,31 @@ bn_msm_loop_add_b_nonzero:
 
 %macro bn_mload_wnaf_a
     // stack: i
-    %mload_kernel(@SEGMENT_KERNEL_BN_WNAF_A)
+    %mload_current(@SEGMENT_BN_WNAF_A)
 %endmacro
 
 %macro bn_mload_wnaf_b
     // stack: i
-    %mload_kernel(@SEGMENT_KERNEL_BN_WNAF_B)
+    %mload_current(@SEGMENT_BN_WNAF_B)
 %endmacro
 
 %macro bn_mload_point_a
     // stack: w
     DUP1
-    %mload_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    %mload_current(@SEGMENT_BN_TABLE_Q)
     //stack: Gy, w
-    SWAP1 %decrement %mload_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    SWAP1 %decrement %mload_current(@SEGMENT_BN_TABLE_Q)
     //stack: Gx, Gy
 %endmacro
 
 %macro bn_mload_point_b
     // stack: w
     DUP1
-    %mload_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    %mload_current(@SEGMENT_BN_TABLE_Q)
-    PUSH @BN_BNEG_LOC %mload_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    PUSH @BN_BNEG_LOC %mload_current(@SEGMENT_BN_TABLE_Q)
     %stack (bneg, Gy, w) -> (@BN_BASE, Gy, bneg, bneg, Gy, w)
     SUB SWAP1 ISZERO MUL SWAP2 MUL ADD
-    SWAP1 %decrement %mload_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    SWAP1 %decrement %mload_current(@SEGMENT_BN_TABLE_Q)
     //stack: Gx, Gy
     PUSH @BN_GLV_BETA
     MULFP254

evm/src/cpu/kernel/asm/curve/bn254/curve_arithmetic/precomputation.asm

@@ -1,5 +1,5 @@
 // Precompute a table of multiples of the BN254 point `Q = (Qx, Qy)`.
-// Let `(Qxi, Qyi) = i * Q`, then store in the `SEGMENT_KERNEL_BN_TABLE_Q` segment of memory the values
+// Let `(Qxi, Qyi) = i * Q`, then store in the `SEGMENT_BN_TABLE_Q` segment of memory the values
 // `i-1 => Qxi`, `i => Qyi if i < 16 else -Qy(32-i)` for `i in range(1, 32, 2)`.
 global bn_precompute_table:
     // stack: Qx, Qy, retdest
@@ -12,14 +12,14 @@ bn_precompute_table_loop:
     // stack i, Qx2, Qy2, Qx, Qy, retdest
     PUSH 1 DUP2 SUB
     %stack (im, i, Qx2, Qy2, Qx, Qy, retdest) -> (i, Qy, im, Qx, i, Qx2, Qy2, Qx, Qy, retdest)
-    %mstore_kernel(@SEGMENT_KERNEL_BN_TABLE_Q) %mstore_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    %mstore_current(@SEGMENT_BN_TABLE_Q) %mstore_current(@SEGMENT_BN_TABLE_Q)
     // stack: i, Qx2, Qy2, Qx, Qy, retdest
     DUP1 PUSH 32 SUB PUSH 1 DUP2 SUB
     // stack: 31-i, 32-i, i, Qx2, Qy2, Qx, Qy, retdest
     DUP7 PUSH @BN_BASE SUB
     // TODO: Could maybe avoid storing Qx a second time here, not sure if it would be more efficient.
     %stack (Qyy, iii, ii, i, Qx2, Qy2, Qx, Qy, retdest) -> (iii, Qx, ii, Qyy, i, Qx2, Qy2, Qx, Qy, retdest)
-    %mstore_kernel(@SEGMENT_KERNEL_BN_TABLE_Q) %mstore_kernel(@SEGMENT_KERNEL_BN_TABLE_Q)
+    %mstore_current(@SEGMENT_BN_TABLE_Q) %mstore_current(@SEGMENT_BN_TABLE_Q)
     // stack: i, Qx2, Qy2, Qx, Qy, retdest
     PUSH 2 ADD
     // stack: i+2, Qx2, Qy2, Qx, Qy, retdest

evm/src/cpu/kernel/asm/curve/wnaf.asm

@@ -34,7 +34,8 @@ wnaf_loop_contd:
     DUP2 SWAP1 SUB
     %stack (n, m, segment, o, retdest) -> (129, o, m, o, segment, n, retdest)
     SUB
-    %stack (i, m, o, segment, n, retdest) -> (m, 0, segment, i, o, segment, n, retdest)
+    GET_CONTEXT
+    %stack (ctx, i, m, o, segment, n, retdest) -> (m, ctx, segment, i, o, segment, n, retdest)
     MSTORE_GENERAL
     // stack: o, segment, n, retdest
     DUP3 ISZERO %jumpi(wnaf_end)

evm/src/cpu/kernel/asm/memory/core.asm

@@ -107,6 +107,19 @@
     // stack: (empty)
 %endmacro
 
+%macro mstore_current(segment, offset)
+    // stack: value
+    PUSH $offset
+    // stack: offset, value
+    PUSH $segment
+    // stack: segment, offset, value
+    GET_CONTEXT
+    // stack: context, segment, offset, value
+    %stack(context, segment, offset, value) -> (value, context, segment, offset)
+    MSTORE_GENERAL
+    // stack: (empty)
+%endmacro
+
 // Load a single byte from user code.
 %macro mload_current_code
     // stack: offset

evm/src/cpu/kernel/constants/mod.rs

@@ -154,7 +154,7 @@ const EC_CONSTANTS: [(&str, [u8; 32]); 20] = [
     ),
     (
         "BN_BNEG_LOC",
-        // This just needs to be large enough to not interfere with anything else in SEGMENT_KERNEL_BN_TABLE_Q.
+        // This just needs to be large enough to not interfere with anything else in SEGMENT_BN_TABLE_Q.
         hex!("0000000000000000000000000000000000000000000000000000000000001337"),
     ),
     (

evm/src/memory/segments.rs

@@ -133,9 +133,9 @@ impl Segment {
             Segment::ShiftTable => "SEGMENT_SHIFT_TABLE",
             Segment::JumpdestBits => "SEGMENT_JUMPDEST_BITS",
             Segment::EcdsaTable => "SEGMENT_KERNEL_ECDSA_TABLE",
-            Segment::BnWnafA => "SEGMENT_KERNEL_BN_WNAF_A",
+            Segment::BnWnafA => "SEGMENT_BN_WNAF_A",
-            Segment::BnWnafB => "SEGMENT_KERNEL_BN_WNAF_B",
+            Segment::BnWnafB => "SEGMENT_BN_WNAF_B",
-            Segment::BnTableQ => "SEGMENT_KERNEL_BN_TABLE_Q",
+            Segment::BnTableQ => "SEGMENT_BN_TABLE_Q",
             Segment::BnPairing => "SEGMENT_KERNEL_BN_PAIRING",
             Segment::AccessedAddresses => "SEGMENT_ACCESSED_ADDRESSES",
             Segment::AccessedStorageKeys => "SEGMENT_ACCESSED_STORAGE_KEYS",