Merge pull request #26 from mir-protocol/recursive_gmimc

Recursive evaluation of GMiMCGate

src/gadgets/arithmetic.rs

@@ -7,11 +7,22 @@ use crate::wire::Wire;
 use crate::witness::PartialWitness;
 
 impl<F: Field> CircuitBuilder<F> {
+    /// Computes `-x`.
     pub fn neg(&mut self, x: Target) -> Target {
         let neg_one = self.neg_one();
         self.mul(x, neg_one)
     }
 
+    /// Computes `x^2`.
+    pub fn square(&mut self, x: Target) -> Target {
+        self.mul(x, x)
+    }
+
+    /// Computes `x^3`.
+    pub fn cube(&mut self, x: Target) -> Target {
+        self.mul_many(&[x, x, x])
+    }
+
     /// Computes `const_0 * multiplicand_0 * multiplicand_1 + const_1 * addend`.
     pub fn arithmetic(
         &mut self,
@@ -111,6 +122,17 @@ impl<F: Field> CircuitBuilder<F> {
         None
     }
 
+    /// Computes `x * y + z`.
+    pub fn mul_add(&mut self, x: Target, y: Target, z: Target) -> Target {
+        self.arithmetic(F::ONE, x, y, F::ONE, z)
+    }
+
+    /// Computes `x * y - z`.
+    pub fn mul_sub(&mut self, x: Target, y: Target, z: Target) -> Target {
+        self.arithmetic(F::ONE, x, y, F::NEG_ONE, z)
+    }
+
+    /// Computes `x + y`.
     pub fn add(&mut self, x: Target, y: Target) -> Target {
         let one = self.one();
         // x + y = 1 * x * 1 + 1 * y
@@ -125,12 +147,14 @@ impl<F: Field> CircuitBuilder<F> {
         sum
     }
 
+    /// Computes `x - y`.
     pub fn sub(&mut self, x: Target, y: Target) -> Target {
         let one = self.one();
         // x - y = 1 * x * 1 + (-1) * y
         self.arithmetic(F::ONE, x, one, F::NEG_ONE, y)
     }
 
+    /// Computes `x * y`.
     pub fn mul(&mut self, x: Target, y: Target) -> Target {
         // x * y = 1 * x * y + 0 * x
         self.arithmetic(F::ONE, x, y, F::ZERO, x)

src/gates/gmimc.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::field::Field;
 use crate::gates::gate::{Gate, GateRef};
+use crate::gates::gmimc_eval::GMiMCEvalGate;
 use crate::generator::{SimpleGenerator, WitnessGenerator};
 use crate::gmimc::gmimc_automatic_constants;
 use crate::target::Target;
@@ -66,7 +67,7 @@ impl<F: Field, const R: usize> Gate<F> for GMiMCGate<F, R> {
     }
 
     fn eval_unfiltered(&self, vars: EvaluationVars<F>) -> Vec<F> {
-        let mut constraints = Vec::with_capacity(W + R);
+        let mut constraints = Vec::with_capacity(self.num_constraints());
 
         // Assert that `swap` is binary.
         let swap = vars.local_wires[Self::WIRE_SWAP];
@@ -119,7 +120,94 @@ impl<F: Field, const R: usize> Gate<F> for GMiMCGate<F, R> {
         builder: &mut CircuitBuilder<F>,
         vars: EvaluationTargets,
     ) -> Vec<Target> {
-        unimplemented!()
+        let mut constraints = Vec::with_capacity(self.num_constraints());
+
+        // Assert that `swap` is binary. Usually we would assert that
+        //     swap(swap - 1) = 0
+        // but to make it work with a single ArithmeticGate, we will instead write it as
+        //     swap*swap - swap = 0
+        let swap = vars.local_wires[Self::WIRE_SWAP];
+        constraints.push(builder.mul_sub(swap, swap, swap));
+
+        let old_index_acc = vars.local_wires[Self::WIRE_INDEX_ACCUMULATOR_OLD];
+        let new_index_acc = vars.local_wires[Self::WIRE_INDEX_ACCUMULATOR_NEW];
+        // computed_new_index_acc = 2 * old_index_acc + swap
+        let two = builder.two();
+        let computed_new_index_acc = builder.mul_add(two, old_index_acc, swap);
+        constraints.push(builder.sub(computed_new_index_acc, new_index_acc));
+
+        let mut state = Vec::with_capacity(12);
+        for i in 0..4 {
+            let a = vars.local_wires[i];
+            let b = vars.local_wires[i + 4];
+            let delta = builder.sub(b, a);
+            state.push(builder.mul_add(swap, delta, a));
+        }
+        for i in 0..4 {
+            let a = vars.local_wires[i + 4];
+            let b = vars.local_wires[i];
+            let delta = builder.sub(b, a);
+            state.push(builder.mul_add(swap, delta, a));
+        }
+        for i in 8..12 {
+            state.push(vars.local_wires[i]);
+        }
+
+        // Value that is implicitly added to each element.
+        // See https://affine.group/2020/02/starkware-challenge
+        let mut addition_buffer = builder.zero();
+
+        for r in 0..R {
+            let active = r % W;
+            let gate = builder.add_gate(GMiMCEvalGate::get(), vec![self.constants[r]]);
+
+            let cubing_input = vars.local_wires[Self::wire_cubing_input(r)];
+            builder.route(
+                cubing_input,
+                Target::Wire(Wire {
+                    gate,
+                    input: GMiMCEvalGate::WIRE_CUBING_INPUT,
+                }),
+            );
+
+            builder.route(
+                addition_buffer,
+                Target::Wire(Wire {
+                    gate,
+                    input: GMiMCEvalGate::WIRE_ADDITION_BUFFER_OLD,
+                }),
+            );
+
+            builder.route(
+                state[active],
+                Target::Wire(Wire {
+                    gate,
+                    input: GMiMCEvalGate::WIRE_STATE_A_OLD,
+                }),
+            );
+
+            constraints.push(Target::Wire(Wire {
+                gate,
+                input: GMiMCEvalGate::WIRE_CONSTRAINT,
+            }));
+
+            addition_buffer = Target::Wire(Wire {
+                gate,
+                input: GMiMCEvalGate::WIRE_ADDITION_BUFFER_NEW,
+            });
+
+            state[active] = Target::Wire(Wire {
+                gate,
+                input: GMiMCEvalGate::WIRE_STATE_A_NEW,
+            });
+        }
+
+        for i in 0..W {
+            state[i] = builder.add(state[i], addition_buffer);
+            constraints.push(builder.sub(state[i], vars.local_wires[Self::wire_output(i)]));
+        }
+
+        constraints
     }
 
     fn generators(

src/gates/gmimc_eval.rs

@@ -4,10 +4,21 @@ use crate::gates::gate::{Gate, GateRef};
 use crate::generator::{SimpleGenerator, WitnessGenerator};
 use crate::target::Target;
 use crate::vars::{EvaluationTargets, EvaluationVars};
+use crate::wire::Wire;
 use crate::witness::PartialWitness;
 
 /// Performs some arithmetic involved in the evaluation of GMiMC's constraint polynomials for one
-/// round.
+/// round. In particular, this performs the following computations:
+///
+/// - `constraint := state_a_old + addition_buffer_old + C_r - cubing_input`
+/// - `f := cubing_input^3`
+/// - `addition_buffer_new := addition_buffer_old + f`
+/// - `state_a_new := state_a_old - f`
+///
+/// Here `state_a_{old,new}` represent the old and new states of the `a`th element of the GMiMC
+/// permutation. `addition_buffer_{old,new}` represents a value that is implicitly added to each
+/// element; see https://affine.group/2020/02/starkware-challenge. `C_r` represents the round
+/// constant for round `r`.
 #[derive(Debug)]
 pub struct GMiMCEvalGate;
 
@@ -15,6 +26,16 @@ impl GMiMCEvalGate {
     pub fn get<F: Field>() -> GateRef<F> {
         GateRef::new(GMiMCEvalGate)
     }
+
+    pub const CONST_C_R: usize = 0;
+
+    pub const WIRE_CONSTRAINT: usize = 0;
+    pub const WIRE_STATE_A_OLD: usize = 1;
+    pub const WIRE_STATE_A_NEW: usize = 2;
+    pub const WIRE_ADDITION_BUFFER_OLD: usize = 3;
+    pub const WIRE_ADDITION_BUFFER_NEW: usize = 4;
+    pub const WIRE_CUBING_INPUT: usize = 5;
+    const WIRE_F: usize = 6;
 }
 
 impl<F: Field> Gate<F> for GMiMCEvalGate {
@@ -23,7 +44,34 @@ impl<F: Field> Gate<F> for GMiMCEvalGate {
     }
 
     fn eval_unfiltered(&self, vars: EvaluationVars<F>) -> Vec<F> {
-        todo!()
+        let c_r = vars.local_constants[Self::CONST_C_R];
+        let constraint = vars.local_wires[Self::WIRE_CONSTRAINT];
+        let state_a_old = vars.local_wires[Self::WIRE_STATE_A_OLD];
+        let state_a_new = vars.local_wires[Self::WIRE_STATE_A_NEW];
+        let addition_buffer_old = vars.local_wires[Self::WIRE_ADDITION_BUFFER_OLD];
+        let addition_buffer_new = vars.local_wires[Self::WIRE_ADDITION_BUFFER_NEW];
+        let cubing_input = vars.local_wires[Self::WIRE_CUBING_INPUT];
+        let f = vars.local_wires[Self::WIRE_F];
+
+        let mut constraints = Vec::with_capacity(self.num_constraints());
+
+        // constraint := state_a_old + addition_buffer_old + C_r - cubing_input
+        let computed_constraint = state_a_old + addition_buffer_old + c_r - cubing_input;
+        constraints.push(constraint - computed_constraint);
+
+        // f := cubing_input^3
+        let computed_f = cubing_input.cube();
+        constraints.push(f - computed_f);
+
+        // addition_buffer_new := addition_buffer_old + f
+        let computed_addition_buffer_new = addition_buffer_old + f;
+        constraints.push(addition_buffer_new - computed_addition_buffer_new);
+
+        // state_a_new := state_a_old - f
+        let computed_state_a_new = state_a_old - f;
+        constraints.push(state_a_new - computed_state_a_new);
+
+        constraints
     }
 
     fn eval_unfiltered_recursively(
@@ -31,7 +79,35 @@ impl<F: Field> Gate<F> for GMiMCEvalGate {
         builder: &mut CircuitBuilder<F>,
         vars: EvaluationTargets,
     ) -> Vec<Target> {
-        unimplemented!()
+        let c_r = vars.local_constants[Self::CONST_C_R];
+        let constraint = vars.local_wires[Self::WIRE_CONSTRAINT];
+        let state_a_old = vars.local_wires[Self::WIRE_STATE_A_OLD];
+        let state_a_new = vars.local_wires[Self::WIRE_STATE_A_NEW];
+        let addition_buffer_old = vars.local_wires[Self::WIRE_ADDITION_BUFFER_OLD];
+        let addition_buffer_new = vars.local_wires[Self::WIRE_ADDITION_BUFFER_NEW];
+        let cubing_input = vars.local_wires[Self::WIRE_CUBING_INPUT];
+        let f = vars.local_wires[Self::WIRE_F];
+
+        let mut constraints = Vec::with_capacity(self.num_constraints());
+
+        // constraint := state_a_old + addition_buffer_old + C_r - cubing_input
+        let sum = builder.add_many(&[state_a_old, addition_buffer_old, c_r]);
+        let computed_constraint = builder.sub(sum, cubing_input);
+        constraints.push(builder.sub(constraint, computed_constraint));
+
+        // f := cubing_input^3
+        let computed_f = builder.cube(cubing_input);
+        constraints.push(builder.sub(f, computed_f));
+
+        // addition_buffer_new := addition_buffer_old + f
+        let computed_addition_buffer_new = builder.add(addition_buffer_old, f);
+        constraints.push(builder.sub(addition_buffer_new, computed_addition_buffer_new));
+
+        // state_a_new := state_a_old - f
+        let computed_state_a_new = builder.sub(state_a_old, f);
+        constraints.push(builder.sub(state_a_new, computed_state_a_new));
+
+        constraints
     }
 
     fn generators(
@@ -41,13 +117,13 @@ impl<F: Field> Gate<F> for GMiMCEvalGate {
     ) -> Vec<Box<dyn WitnessGenerator<F>>> {
         let gen = GMiMCEvalGenerator::<F> {
             gate_index,
-            constant: local_constants[0],
+            c_r: local_constants[Self::CONST_C_R],
         };
         vec![Box::new(gen)]
     }
 
     fn num_wires(&self) -> usize {
-        6
+        7
     }
 
     fn num_constants(&self) -> usize {
@@ -59,22 +135,87 @@ impl<F: Field> Gate<F> for GMiMCEvalGate {
     }
 
     fn num_constraints(&self) -> usize {
-        unimplemented!()
+        4
     }
 }
 
 #[derive(Debug)]
 struct GMiMCEvalGenerator<F: Field> {
     gate_index: usize,
-    constant: F,
+    c_r: F,
 }
 
 impl<F: Field> SimpleGenerator<F> for GMiMCEvalGenerator<F> {
     fn dependencies(&self) -> Vec<Target> {
-        todo!()
+        let gate = self.gate_index;
+        vec![
+            Target::Wire(Wire {
+                gate,
+                input: GMiMCEvalGate::WIRE_CUBING_INPUT,
+            }),
+            Target::Wire(Wire {
+                gate,
+                input: GMiMCEvalGate::WIRE_ADDITION_BUFFER_OLD,
+            }),
+            Target::Wire(Wire {
+                gate,
+                input: GMiMCEvalGate::WIRE_STATE_A_OLD,
+            }),
+        ]
     }
 
     fn run_once(&self, witness: &PartialWitness<F>) -> PartialWitness<F> {
-        todo!()
+        let gate = self.gate_index;
+        let wire_constraint = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_CONSTRAINT,
+        };
+        let wire_state_a_old = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_STATE_A_OLD,
+        };
+        let wire_state_a_new = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_STATE_A_NEW,
+        };
+        let wire_addition_buffer_old = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_ADDITION_BUFFER_OLD,
+        };
+        let wire_addition_buffer_new = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_ADDITION_BUFFER_NEW,
+        };
+        let wire_cubing_input = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_CUBING_INPUT,
+        };
+        let wire_f = Wire {
+            gate,
+            input: GMiMCEvalGate::WIRE_F,
+        };
+
+        let addition_buffer_old = witness.get_wire(wire_addition_buffer_old);
+        let state_a_old = witness.get_wire(wire_state_a_old);
+        let cubing_input = witness.get_wire(wire_cubing_input);
+
+        // constraint := state_a_old + addition_buffer_old + C_r - cubing_input
+        let constraint = state_a_old + addition_buffer_old + self.c_r - cubing_input;
+
+        // f := cubing_input^3
+        let f = cubing_input.cube();
+
+        // addition_buffer_new := addition_buffer_old + f
+        let addition_buffer_new = addition_buffer_old + f;
+
+        // state_a_new := state_a_old - f
+        let state_a_new = state_a_old - f;
+
+        let mut witness = PartialWitness::new();
+        witness.set_wire(wire_constraint, constraint);
+        witness.set_wire(wire_f, f);
+        witness.set_wire(wire_state_a_new, addition_buffer_new);
+        witness.set_wire(wire_addition_buffer_new, state_a_new);
+        witness
     }
 }