Feat: Builder Custom Gate Support

Cargo.toml

@@ -19,11 +19,11 @@ tiny-keccak = { version = "2.0", features = ["keccak"] }
 halo2curves = { git = "https://github.com/PolyhedraZK/halo2curves", default-features = false, features = [
     "bits",
 ] }
-arith = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev" }
-expander_config = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev", package = "config" }
-expander_circuit = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev", package = "circuit" }
-gkr = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev" }
-gf2 = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev" }
-mersenne31 = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev" }
-expander_transcript = { git = "https://github.com/PolyhedraZK/Expander", branch = "dev", package = "transcript" }
-
+arith = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb" }
+babybear = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb"}
+expander_config = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb", package = "config" }
+expander_circuit = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb", package = "circuit" }
+gkr = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb" }
+gf2 = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb" }
+mersenne31 = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb" }
+expander_transcript = { git = "https://github.com/NebraZKP/Expander", rev = "5470aaee59718f4d48977b63fac3baa046e7f5bb", package = "transcript" }

expander_compiler/Cargo.toml

@@ -14,6 +14,7 @@ halo2curves.workspace = true
 tiny-keccak.workspace = true
 expander_config.workspace = true
 expander_circuit.workspace = true
+babybear.workspace = true
 gkr.workspace = true
 arith.workspace = true
 gf2.workspace = true

expander_compiler/ec_go_lib/src/lib.rs

@@ -1,3 +1,5 @@
+#![allow(incomplete_features)]
+#![feature(generic_const_exprs)]
 use arith::FieldSerde;
 use expander_compiler::circuit::layered;
 use libc::{c_uchar, c_ulong, malloc};
@@ -49,6 +51,9 @@ fn compile_inner(ir_source: Vec<u8>, config_id: u64) -> Result<(Vec<u8>, Vec<u8>
         1 => compile_inner_with_config::<config::M31Config>(ir_source),
         2 => compile_inner_with_config::<config::BN254Config>(ir_source),
         3 => compile_inner_with_config::<config::GF2Config>(ir_source),
+        4 => compile_inner_with_config::<config::M31Gkr2Config>(ir_source),
+        5 => compile_inner_with_config::<config::BabyBearConfig>(ir_source),
+        6 => compile_inner_with_config::<config::BabyBearGkr2Config>(ir_source),
         _ => Err(format!("unknown config id: {}", config_id)),
     }
 }
@@ -156,6 +161,7 @@ fn prove_circuit_file_inner<C: expander_config::GKRConfig, CC: config::Config>(
 ) -> Vec<u8>
 where
     C::SimdCircuitField: arith::SimdField<Scalar = CC::CircuitField>,
+    [(); C::DEGREE_PLUS_ONE]:,
 {
     let config = expander_config::Config::<C>::new(
         expander_config::GKRScheme::Vanilla,
@@ -180,6 +186,7 @@ fn verify_circuit_file_inner<C: expander_config::GKRConfig, CC: config::Config>(
 ) -> u8
 where
     C::SimdCircuitField: arith::SimdField<Scalar = CC::CircuitField>,
+    [(); C::DEGREE_PLUS_ONE]:,
 {
     let config = expander_config::Config::<C>::new(
         expander_config::GKRScheme::Vanilla,
@@ -224,6 +231,18 @@ pub extern "C" fn prove_circuit_file(
             circuit_filename,
             witness,
         ),
+        4 => prove_circuit_file_inner::<expander_config::M31ExtConfigSha2, config::M31Gkr2Config>(
+            circuit_filename,
+            witness,
+        ),
+        5 => prove_circuit_file_inner::<
+            expander_config::BabyBearExt4ConfigSha2,
+            config::BabyBearConfig,
+        >(circuit_filename, witness),
+        6 => prove_circuit_file_inner::<
+            expander_config::BabyBearExt4ConfigSha2,
+            config::BabyBearGkr2Config,
+        >(circuit_filename, witness),
         _ => panic!("unknown config id: {}", config_id),
     };
     let proof_len = proof.len();
@@ -271,6 +290,19 @@ pub extern "C" fn verify_circuit_file(
             witness,
             proof,
         ),
+        4 => verify_circuit_file_inner::<expander_config::M31ExtConfigSha2, config::M31Gkr2Config>(
+            circuit_filename,
+            witness,
+            proof,
+        ),
+        5 => verify_circuit_file_inner::<
+            expander_config::BabyBearExt4ConfigSha2,
+            config::BabyBearConfig,
+        >(circuit_filename, witness, proof),
+        6 => verify_circuit_file_inner::<
+            expander_config::BabyBearExt4ConfigSha2,
+            config::BabyBearGkr2Config,
+        >(circuit_filename, witness, proof),
         _ => panic!("unknown config id: {}", config_id),
     }
 }

expander_compiler/src/circuit/config.rs

@@ -1,10 +1,12 @@
-use std::{fmt::Debug, hash::Hash};
-
 use crate::field::Field;
+use expander_config::GKRScheme;
+use std::{fmt::Debug, hash::Hash};
 
 pub trait Config: Default + Clone + Ord + Debug + Hash + Copy + 'static {
     type CircuitField: Field;
 
+    const SCHEME: GKRScheme = GKRScheme::Vanilla;
+
     const CONFIG_ID: usize;
 
     const COST_INPUT: usize = 1000;
@@ -25,6 +27,16 @@ impl Config for M31Config {
     const CONFIG_ID: usize = 1;
 }
 
+#[derive(Default, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)]
+pub struct M31Gkr2Config {}
+
+impl Config for M31Gkr2Config {
+    type CircuitField = crate::field::M31;
+
+    const CONFIG_ID: usize = 4;
+    const SCHEME: GKRScheme = GKRScheme::GkrSquare;
+}
+
 #[derive(Default, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)]
 pub struct BN254Config {}
 
@@ -48,3 +60,22 @@ impl Config for GF2Config {
 
     const ENABLE_RANDOM_COMBINATION: bool = false;
 }
+
+#[derive(Default, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)]
+pub struct BabyBearConfig {}
+
+impl Config for BabyBearConfig {
+    type CircuitField = crate::field::BabyBear;
+
+    const CONFIG_ID: usize = 5;
+}
+
+#[derive(Default, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)]
+pub struct BabyBearGkr2Config {}
+
+impl Config for BabyBearGkr2Config {
+    type CircuitField = crate::field::BabyBear;
+
+    const CONFIG_ID: usize = 6;
+    const SCHEME: GKRScheme = GKRScheme::GkrSquare;
+}

expander_compiler/src/field.rs

@@ -1,4 +1,5 @@
 pub use arith::{BN254Fr as BN254, Field as FieldArith, FieldForECC as FieldModulus};
+pub use babybear::BabyBear;
 pub use gf2::GF2;
 pub use mersenne31::M31;
 
@@ -7,6 +8,7 @@ use arith::{FieldForECC, FieldSerde, FieldSerdeError};
 
 pub trait Field: FieldArith + FieldForECC + FieldSerde {}
 
+impl Field for BabyBear {}
 impl Field for BN254 {}
 impl Field for GF2 {}
 impl Field for M31 {}

expander_compiler/src/frontend/api.rs

@@ -42,6 +42,8 @@ pub trait BasicAPI<C: Config> {
         y: impl ToVariableOrValue<C::CircuitField>,
     );
     fn get_random_value(&mut self) -> Variable;
+    /// Can be either power5 or power7
+    fn power_gate(&mut self, x: impl ToVariableOrValue<C::CircuitField>, power: usize) -> Variable;
 }
 
 pub trait UnconstrainedAPI<C: Config> {

expander_compiler/src/frontend/builder.rs

@@ -174,6 +174,28 @@ impl<C: Config> BasicAPI<C> for Builder<C> {
         self.new_var()
     }
 
+    fn power_gate(&mut self, x: impl ToVariableOrValue<C::CircuitField>, power: usize) -> Variable {
+        match power {
+            5 => {
+                let x = self.convert_to_variable(x);
+                self.instructions.push(SourceInstruction::CustomGate {
+                    gate_type: 12345,
+                    inputs: vec![x.id],
+                });
+                self.new_var()
+            }
+            7 => {
+                let x = self.convert_to_variable(x);
+                self.instructions.push(SourceInstruction::CustomGate {
+                    gate_type: 12347,
+                    inputs: vec![x.id],
+                });
+                self.new_var()
+            }
+            _ => panic!("Unsupported power gate"),
+        }
+    }
+
     fn div(
         &mut self,
         x: impl ToVariableOrValue<C::CircuitField>,
@@ -445,6 +467,14 @@ impl<C: Config> BasicAPI<C> for RootBuilder<C> {
     fn get_random_value(&mut self) -> Variable {
         self.last_builder().get_random_value()
     }
+
+    fn power_gate(
+        &mut self,
+        x: impl ToVariableOrValue<<C as Config>::CircuitField>,
+        power: usize,
+    ) -> Variable {
+        self.last_builder().power_gate(x, power)
+    }
 }
 
 impl<C: Config> RootBuilder<C> {

expander_compiler/src/hints/mod.rs

@@ -64,6 +64,7 @@ impl BuiltinHintIds {
     }
 }
 
+/// ??? Placeholder for the actual implementation
 fn stub_impl_general<F: Field>(hint_id: usize, inputs: &Vec<F>, num_outputs: usize) -> Vec<F> {
     let mut hasher = DefaultHasher::new();
     hint_id.hash(&mut hasher);
@@ -236,9 +237,21 @@ fn binop_hint_on_u256<F: Field, G: Fn(U256, U256) -> U256>(inputs: &[F], f: G) -
 }
 
 pub fn stub_impl<F: Field>(hint_id: usize, inputs: &Vec<F>, num_outputs: usize) -> Vec<F> {
+    if hint_id == 12345 {
+        let a2 = inputs[0] * inputs[0];
+        let a4 = a2 * a2;
+        return vec![a4 * inputs[0]];
+    } else if hint_id == 12346 {
+        return vec![inputs[0]];
+    } else if hint_id == 12347 {
+        let a2 = inputs[0].square();
+        let a4 = a2.square();
+        let a6 = a4 * a2;
+        return vec![a6 * inputs[0]];
+    }
     match BuiltinHintIds::from_usize(hint_id) {
         Some(hint_id) => impl_builtin_hint(hint_id, inputs, num_outputs),
-        None => stub_impl_general(hint_id, inputs, num_outputs),
+        None => panic!("invalid hint_id")
     }
 }
 

expander_compiler/src/layering/wire.rs

@@ -1,3 +1,4 @@
+use expander_config::GKRScheme;
 use std::collections::HashMap;
 
 use crate::{
@@ -20,6 +21,32 @@ struct LayoutQuery {
     var_pos: HashMap<usize, usize>,
 }
 
+/// Add value at `input_pos` to `output_pos` with coefficient `coef`.
+fn push_add_gate<C: Config>(
+    res: &mut Segment<C>,
+    input_pos: usize,
+    output_pos: usize,
+    coef: Coef<C>,
+) {
+    match C::SCHEME {
+        GKRScheme::Vanilla => {
+            res.gate_adds.push(GateAdd {
+                inputs: [input_pos],
+                output: output_pos,
+                coef,
+            });
+        }
+        GKRScheme::GkrSquare => {
+            res.gate_customs.push(GateCustom {
+                gate_type: 12346, // power1 gate
+                inputs: vec![input_pos],
+                output: output_pos,
+                coef,
+            });
+        }
+    }
+}
+
 impl LayoutQuery {
     // given a parent layer layout, this function query the layout of a sub circuit
     fn query<F, C: Config>(
@@ -278,11 +305,12 @@ impl<'a, C: Config> CompileContext<'a, C> {
             };
             // if it's not the first layer, just relay it
             if ic.min_layer[*x] != next_layer {
-                res.gate_adds.push(GateAdd {
-                    inputs: [aq.var_pos[x]],
-                    output: pos,
-                    coef: Coef::Constant(C::CircuitField::one()),
-                });
+                push_add_gate(
+                    &mut res,
+                    aq.var_pos[x],
+                    pos,
+                    Coef::Constant(C::CircuitField::one()),
+                );
                 continue;
             }
             if let Some(value) = ic.constant_like_variables.get(x) {
@@ -302,11 +330,12 @@ impl<'a, C: Config> CompileContext<'a, C> {
                             });
                         }
                         VarSpec::Linear(vid) => {
-                            res.gate_adds.push(GateAdd {
-                                inputs: [aq.var_pos[vid]],
-                                output: pos,
-                                coef: Coef::Constant(term.coef),
-                            });
+                            push_add_gate(
+                                &mut res,
+                                aq.var_pos[vid],
+                                pos,
+                                Coef::Constant(term.coef),
+                            );
                         }
                         VarSpec::Quad(vid0, vid1) => {
                             res.gate_muls.push(GateMul {
@@ -324,11 +353,7 @@ impl<'a, C: Config> CompileContext<'a, C> {
                             });
                         }
                         VarSpec::RandomLinear(vid) => {
-                            res.gate_adds.push(GateAdd {
-                                inputs: [aq.var_pos[vid]],
-                                output: pos,
-                                coef: Coef::Random,
-                            });
+                            push_add_gate(&mut res, aq.var_pos[vid], pos, Coef::Random);
                         }
                     }
                 }
@@ -345,11 +370,7 @@ impl<'a, C: Config> CompileContext<'a, C> {
                 } else {
                     Coef::Random
                 };
-                res.gate_adds.push(GateAdd {
-                    inputs: [aq.var_pos[v]],
-                    output: pos,
-                    coef,
-                });
+                push_add_gate(&mut res, aq.var_pos[v], pos, coef);
             }
             for i in cc.sub_circuit_ids.iter() {
                 let insn_id = ic.sub_circuit_insn_ids[*i];
@@ -373,11 +394,12 @@ impl<'a, C: Config> CompileContext<'a, C> {
                         placement.iter().position(|x| *x == vpid).unwrap()
                     }
                 };
-                res.gate_adds.push(GateAdd {
-                    inputs: [sub_cur_layout_all[&insn_id].offset + spos],
-                    output: pos,
-                    coef: Coef::Constant(C::CircuitField::one()),
-                });
+                push_add_gate(
+                    &mut res,
+                    sub_cur_layout_all[&insn_id].offset + spos,
+                    pos,
+                    Coef::Constant(C::CircuitField::one()),
+                );
             }
         }