Merge pull request #931 from mir-protocol/fp381-opcodes

evm/src/cpu/kernel/aggregator.rs

@@ -36,6 +36,7 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/core/transfer.asm"),
         include_str!("asm/core/util.asm"),
         include_str!("asm/core/access_lists.asm"),
+        include_str!("asm/curve/bls381/util.asm"),
         include_str!("asm/curve/bn254/curve_arithmetic/constants.asm"),
         include_str!("asm/curve/bn254/curve_arithmetic/curve_add.asm"),
         include_str!("asm/curve/bn254/curve_arithmetic/curve_mul.asm"),

evm/src/cpu/kernel/asm/curve/bls381/util.asm

@@ -0,0 +1,53 @@
+%macro add_fp381
+    // stack:         x0, x1, y0, y1
+    PROVER_INPUT(sf::bls381_base::add_hi)
+    // stack:     z1, x0, x1, y0, y1
+    SWAP4
+    // stack:     y1, x0, x1, y0, z1
+    PROVER_INPUT(sf::bls381_base::add_lo)
+    // stack: z0, y1, x0, x1, y0, z1
+    SWAP4
+    // stack: y0, y1, x0, x1, z0, z1
+    %pop4
+    // stack:                 z0, z1
+%endmacro
+
+%macro mul_fp381
+    // stack:         x0, x1, y0, y1
+    PROVER_INPUT(sf::bls381_base::mul_hi)
+    // stack:     z1, x0, x1, y0, y1
+    SWAP4
+    // stack:     y1, x0, x1, y0, z1
+    PROVER_INPUT(sf::bls381_base::mul_lo)
+    // stack: z0, y1, x0, x1, y0, z1
+    SWAP4
+    // stack: y0, y1, x0, x1, z0, z1
+    %pop4
+    // stack:                 z0, z1
+%endmacro
+
+%macro sub_fp381
+    // stack:         x0, x1, y0, y1
+    PROVER_INPUT(sf::bls381_base::sub_hi)
+    // stack:     z1, x0, x1, y0, y1
+    SWAP4
+    // stack:     y1, x0, x1, y0, z1
+    PROVER_INPUT(sf::bls381_base::sub_lo)
+    // stack: z0, y1, x0, x1, y0, z1
+    SWAP4
+    // stack: y0, y1, x0, x1, z0, z1
+    %pop4
+    // stack:                 z0, z1
+%endmacro
+
+global test_add_fp381:
+    %add_fp381
+    %jump(0xdeadbeef)
+
+global test_mul_fp381:
+    %mul_fp381
+    %jump(0xdeadbeef)
+
+global test_sub_fp381:
+    %sub_fp381
+    %jump(0xdeadbeef)

evm/src/cpu/kernel/tests/bls381.rs

@@ -0,0 +1,42 @@
+use anyhow::Result;
+use ethereum_types::U512;
+use rand::Rng;
+
+use crate::cpu::kernel::interpreter::{
+    run_interpreter_with_memory, InterpreterMemoryInitialization,
+};
+use crate::extension_tower::{Stack, BLS381};
+use crate::memory::segments::Segment::KernelGeneral;
+
+fn run_and_return_bls(label: String, x: BLS381, y: BLS381) -> BLS381 {
+    let mut stack = x.on_stack();
+    stack.extend(y.on_stack());
+    let setup = InterpreterMemoryInitialization {
+        label,
+        stack,
+        segment: KernelGeneral,
+        memory: vec![],
+    };
+    let interpreter = run_interpreter_with_memory(setup).unwrap();
+    let output = interpreter.stack();
+    BLS381 {
+        val: U512::from(output[1]) + (U512::from(output[0]) << 256),
+    }
+}
+
+#[test]
+fn test_bls_ops() -> Result<()> {
+    let mut rng = rand::thread_rng();
+    let x: BLS381 = rng.gen::<BLS381>();
+    let y: BLS381 = rng.gen::<BLS381>();
+
+    let output_add = run_and_return_bls("test_add_fp381".to_string(), x, y);
+    let output_mul = run_and_return_bls("test_mul_fp381".to_string(), x, y);
+    let output_sub = run_and_return_bls("test_sub_fp381".to_string(), x, y);
+
+    assert_eq!(output_add, x + y);
+    assert_eq!(output_mul, x * y);
+    assert_eq!(output_sub, x - y);
+
+    Ok(())
+}

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

@@ -1,6 +1,7 @@
 mod account_code;
 mod balance;
 mod bignum;
+mod bls381;
 mod bn254;
 mod core;
 mod ecc;

evm/src/extension_tower.rs

@@ -136,6 +136,14 @@ impl BLS381 {
             val: U512::from(val),
         }
     }
+
+    pub fn lo(self) -> U256 {
+        U256(self.val.0[..4].try_into().unwrap())
+    }
+
+    pub fn hi(self) -> U256 {
+        U256(self.val.0[4..].try_into().unwrap())
+    }
 }
 
 impl Distribution<BLS381> for Standard {
@@ -201,10 +209,10 @@ impl Mul for BLS381 {
 
     fn mul(self, other: Self) -> Self {
         // x1, y1 are at most ((q-1) // 2^256) < 2^125
-        let x0 = U512(self.val.0[..4].try_into().unwrap());
-        let x1 = U512(self.val.0[4..].try_into().unwrap());
-        let y0 = U512(other.val.0[..4].try_into().unwrap());
-        let y1 = U512(other.val.0[4..].try_into().unwrap());
+        let x0 = U512::from(self.lo());
+        let x1 = U512::from(self.hi());
+        let y0 = U512::from(other.lo());
+        let y1 = U512::from(other.hi());
 
         let z00 = BLS381 {
             val: x0.saturating_mul(y0) % BLS_BASE,
@@ -1196,6 +1204,12 @@ pub trait Stack {
     fn on_stack(self) -> Vec<U256>;
 }
 
+impl Stack for BLS381 {
+    fn on_stack(self) -> Vec<U256> {
+        vec![self.lo(), self.hi()]
+    }
+}
+
 impl Stack for Fp6<BN254> {
     fn on_stack(self) -> Vec<U256> {
         let f: [U256; 6] = unsafe { transmute(self) };

evm/src/generation/prover_input.rs

@@ -2,12 +2,12 @@ use std::mem::transmute;
 use std::str::FromStr;
 
 use anyhow::{bail, Error};
-use ethereum_types::{BigEndianHash, H256, U256};
+use ethereum_types::{BigEndianHash, H256, U256, U512};
 use plonky2::field::types::Field;
 
-use crate::extension_tower::{FieldExt, Fp12, BN254};
+use crate::extension_tower::{FieldExt, Fp12, BLS381, BN254};
 use crate::generation::prover_input::EvmField::{
-    Bn254Base, Bn254Scalar, Secp256k1Base, Secp256k1Scalar,
+    Bls381Base, Bls381Scalar, Bn254Base, Bn254Scalar, Secp256k1Base, Secp256k1Scalar,
 };
 use crate::generation::prover_input::FieldOp::{Inverse, Sqrt};
 use crate::generation::state::GenerationState;
@@ -30,6 +30,7 @@ impl<F: Field> GenerationState<F> {
         match input_fn.0[0].as_str() {
             "end_of_txns" => self.run_end_of_txns(),
             "ff" => self.run_ff(input_fn),
+            "sf" => self.run_sf(input_fn),
             "ffe" => self.run_ffe(input_fn),
             "mpt" => self.run_mpt(),
             "rlp" => self.run_rlp(),
@@ -56,6 +57,26 @@ impl<F: Field> GenerationState<F> {
         field.op(op, x)
     }
 
+    /// Special finite field operations.
+    fn run_sf(&self, input_fn: &ProverInputFn) -> U256 {
+        let field = EvmField::from_str(input_fn.0[1].as_str()).unwrap();
+        let inputs: [U256; 4] = match field {
+            Bls381Base => std::array::from_fn(|i| {
+                stack_peek(self, i).expect("Insufficient number of items on stack")
+            }),
+            _ => todo!(),
+        };
+        match input_fn.0[2].as_str() {
+            "add_lo" => field.add_lo(inputs),
+            "add_hi" => field.add_hi(inputs),
+            "mul_lo" => field.mul_lo(inputs),
+            "mul_hi" => field.mul_hi(inputs),
+            "sub_lo" => field.sub_lo(inputs),
+            "sub_hi" => field.sub_hi(inputs),
+            _ => todo!(),
+        }
+    }
+
     /// Finite field extension operations.
     fn run_ffe(&self, input_fn: &ProverInputFn) -> U256 {
         let field = EvmField::from_str(input_fn.0[1].as_str()).unwrap();
@@ -66,16 +87,12 @@ impl<F: Field> GenerationState<F> {
             .unwrap()
             .parse::<usize>()
             .unwrap();
-        let ptr = stack_peek(self, 11 - n).expect("Empty stack").as_usize();
+        let ptr = stack_peek(self, 11 - n)
+            .expect("Insufficient number of items on stack")
+            .as_usize();
 
         let f: [U256; 12] = match field {
-            Bn254Base => {
-                let mut f: [U256; 12] = [U256::zero(); 12];
-                for i in 0..12 {
-                    f[i] = kernel_peek(self, BnPairing, ptr + i);
-                }
-                f
-            }
+            Bn254Base => std::array::from_fn(|i| kernel_peek(self, BnPairing, ptr + i)),
             _ => todo!(),
         };
         field.field_extension_inverse(n, f)
@@ -126,6 +143,8 @@ impl<F: Field> GenerationState<F> {
 }
 
 enum EvmField {
+    Bls381Base,
+    Bls381Scalar,
     Bn254Base,
     Bn254Scalar,
     Secp256k1Base,
@@ -142,6 +161,8 @@ impl FromStr for EvmField {
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         Ok(match s {
+            "bls381_base" => Bls381Base,
+            "bls381_scalar" => Bls381Scalar,
             "bn254_base" => Bn254Base,
             "bn254_scalar" => Bn254Scalar,
             "secp256k1_base" => Secp256k1Base,
@@ -166,6 +187,8 @@ impl FromStr for FieldOp {
 impl EvmField {
     fn order(&self) -> U256 {
         match self {
+            EvmField::Bls381Base => todo!(),
+            EvmField::Bls381Scalar => todo!(),
             EvmField::Bn254Base => {
                 U256::from_str("0x30644e72e131a029b85045b68181585d97816a916871ca8d3c208c16d87cfd47")
                     .unwrap()
@@ -206,6 +229,54 @@ impl EvmField {
         modexp(x, q, n)
     }
 
+    fn add_lo(&self, inputs: [U256; 4]) -> U256 {
+        let [y1, x0, x1, y0] = inputs;
+        let x = U512::from(x0) + (U512::from(x1) << 256);
+        let y = U512::from(y0) + (U512::from(y1) << 256);
+        let z = BLS381 { val: x } + BLS381 { val: y };
+        z.lo()
+    }
+
+    fn add_hi(&self, inputs: [U256; 4]) -> U256 {
+        let [x0, x1, y0, y1] = inputs;
+        let x = U512::from(x0) + (U512::from(x1) << 256);
+        let y = U512::from(y0) + (U512::from(y1) << 256);
+        let z = BLS381 { val: x } + BLS381 { val: y };
+        z.hi()
+    }
+
+    fn mul_lo(&self, inputs: [U256; 4]) -> U256 {
+        let [y1, x0, x1, y0] = inputs;
+        let x = U512::from(x0) + (U512::from(x1) << 256);
+        let y = U512::from(y0) + (U512::from(y1) << 256);
+        let z = BLS381 { val: x } * BLS381 { val: y };
+        z.lo()
+    }
+
+    fn mul_hi(&self, inputs: [U256; 4]) -> U256 {
+        let [x0, x1, y0, y1] = inputs;
+        let x = U512::from(x0) + (U512::from(x1) << 256);
+        let y = U512::from(y0) + (U512::from(y1) << 256);
+        let z = BLS381 { val: x } * BLS381 { val: y };
+        z.hi()
+    }
+
+    fn sub_lo(&self, inputs: [U256; 4]) -> U256 {
+        let [y1, x0, x1, y0] = inputs;
+        let x = U512::from(x0) + (U512::from(x1) << 256);
+        let y = U512::from(y0) + (U512::from(y1) << 256);
+        let z = BLS381 { val: x } - BLS381 { val: y };
+        z.lo()
+    }
+
+    fn sub_hi(&self, inputs: [U256; 4]) -> U256 {
+        let [x0, x1, y0, y1] = inputs;
+        let x = U512::from(x0) + (U512::from(x1) << 256);
+        let y = U512::from(y0) + (U512::from(y1) << 256);
+        let z = BLS381 { val: x } - BLS381 { val: y };
+        z.hi()
+    }
+
     fn field_extension_inverse(&self, n: usize, f: [U256; 12]) -> U256 {
         let f: Fp12<BN254> = unsafe { transmute(f) };
         let f_inv: [U256; 12] = unsafe { transmute(f.inv()) };