GMiMC, witness generation

src/circuit_builder.rs

@@ -3,7 +3,7 @@ use std::collections::HashSet;
 use crate::circuit_data::CircuitConfig;
 use crate::field::field::Field;
 use crate::gates::gate::{GateInstance, GateRef};
-use crate::generator::{CopyGenerator, WitnessGenerator2};
+use crate::generator::{CopyGenerator, WitnessGenerator};
 use crate::target::Target;
 use crate::gates::constant::ConstantGate2;
 use crate::wire::Wire;
@@ -12,7 +12,7 @@ pub struct CircuitBuilder2<F: Field> {
     config: CircuitConfig,
     gates: HashSet<GateRef<F>>,
     gate_instances: Vec<GateInstance<F>>,
-    generators: Vec<Box<dyn WitnessGenerator2<F>>>,
+    generators: Vec<Box<dyn WitnessGenerator<F>>>,
 }
 
 impl<F: Field> CircuitBuilder2<F> {
@@ -62,7 +62,7 @@ impl<F: Field> CircuitBuilder2<F> {
         assert!(y.is_routable(self.config));
     }
 
-    pub fn add_generator<G: WitnessGenerator2<F>>(&mut self, generator: G) {
+    pub fn add_generator<G: WitnessGenerator<F>>(&mut self, generator: G) {
         self.generators.push(Box::new(generator));
     }
 

src/gates/deterministic_gate.rs

@@ -4,7 +4,7 @@ use crate::circuit_data::CircuitConfig;
 use crate::constraint_polynomial::{ConstraintPolynomial, EvaluationVars};
 use crate::field::field::Field;
 use crate::gates::gate::Gate;
-use crate::generator::{SimpleGenerator, WitnessGenerator2};
+use crate::generator::{SimpleGenerator, WitnessGenerator};
 use crate::target::Target;
 use crate::wire::Wire;
 use crate::witness::PartialWitness;
@@ -35,7 +35,7 @@ pub trait DeterministicGate<F: Field>: 'static {
         &self,
         _config: CircuitConfig,
         _gate_index: usize,
-    ) -> Vec<Box<dyn WitnessGenerator2<F>>> {
+    ) -> Vec<Box<dyn WitnessGenerator<F>>> {
         Vec::new()
     }
 }
@@ -73,7 +73,7 @@ impl<F: Field, DG: DeterministicGate<F>> Gate<F> for DeterministicGateAdapter<F,
         gate_index: usize,
         local_constants: Vec<F>,
         next_constants: Vec<F>,
-    ) -> Vec<Box<dyn WitnessGenerator2<F>>> {
+    ) -> Vec<Box<dyn WitnessGenerator<F>>> {
         self.gate.outputs(config).outputs
             .into_iter()
             .map(|(location, out)| {
@@ -87,7 +87,7 @@ impl<F: Field, DG: DeterministicGate<F>> Gate<F> for DeterministicGateAdapter<F,
 
                 // We need the type system to treat this as a boxed `WitnessGenerator2<F>`, rather
                 // than a boxed `OutputGenerator<F>`.
-                let b: Box::<dyn WitnessGenerator2<F>> = Box::new(og);
+                let b: Box::<dyn WitnessGenerator<F>> = Box::new(og);
                 b
             })
             .chain(self.gate.additional_generators(config, gate_index))

src/gates/gate.rs

@@ -4,7 +4,7 @@ use std::rc::Rc;
 use crate::circuit_data::CircuitConfig;
 use crate::constraint_polynomial::ConstraintPolynomial;
 use crate::field::field::Field;
-use crate::generator::WitnessGenerator2;
+use crate::generator::WitnessGenerator;
 use num::ToPrimitive;
 
 /// A custom gate.
@@ -21,7 +21,7 @@ pub trait Gate<F: Field>: 'static {
         gate_index: usize,
         local_constants: Vec<F>,
         next_constants: Vec<F>,
-    ) -> Vec<Box<dyn WitnessGenerator2<F>>>;
+    ) -> Vec<Box<dyn WitnessGenerator<F>>>;
 
     /// The number of constants used by this gate.
     fn num_constants(&self, config: CircuitConfig) -> usize {

src/gates/gmimc.rs

@@ -1,19 +1,13 @@
-use std::convert::TryInto;
 use std::sync::Arc;
 
-use num::{BigUint, One};
+use num::{BigUint, FromPrimitive, One};
 
 use crate::circuit_data::CircuitConfig;
 use crate::constraint_polynomial::ConstraintPolynomial;
 use crate::field::field::Field;
 use crate::gates::deterministic_gate::{DeterministicGate, DeterministicGateAdapter};
-use crate::gates::gate::{Gate, GateRef};
+use crate::gates::gate::GateRef;
 use crate::gates::output_graph::{GateOutputLocation, OutputGraph};
-use crate::generator::{SimpleGenerator, WitnessGenerator2};
-use crate::gmimc::{gmimc_permute, gmimc_permute_array};
-use crate::target::Target;
-use crate::wire::Wire;
-use crate::witness::PartialWitness;
 
 /// Evaluates a full GMiMC permutation, and writes the output to the next gate's first `width`
 /// wires (which could be the input of another `GMiMCGate`).
@@ -126,14 +120,20 @@ impl<F: Field, const W: usize, const R: usize> DeterministicGate<F> for GMiMCGat
 
 #[cfg(test)]
 mod tests {
+    use std::convert::TryInto;
     use std::sync::Arc;
 
+    use num::ToPrimitive;
+
     use crate::circuit_data::CircuitConfig;
     use crate::field::crandall_field::CrandallField;
     use crate::field::field::Field;
     use crate::gates::deterministic_gate::DeterministicGate;
     use crate::gates::gmimc::GMiMCGate;
-    use crate::circuit_builder::CircuitBuilder2;
+    use crate::generator::generate_partial_witness;
+    use crate::gmimc::gmimc_permute_naive;
+    use crate::wire::Wire;
+    use crate::witness::PartialWitness;
 
     #[test]
     fn degree() {
@@ -147,10 +147,7 @@ mod tests {
             security_bits: 128,
         };
         let outs = gate.outputs(config);
-
         assert_eq!(outs.degree(), 3);
-        assert_eq!(outs.max_local_output_index(), Some(57));
-        panic!();
     }
 
     #[test]
@@ -159,17 +156,38 @@ mod tests {
         const W: usize = 12;
         const R: usize = 101;
         let constants = Arc::new([F::TWO; R]);
-        let gate = GMiMCGate::<F, W, R>::with_constants(constants);
+        type Gate = GMiMCGate::<F, W, R>;
+        let gate = Gate::with_constants(constants.clone());
 
         let config = CircuitConfig {
             num_wires: 200,
             num_routed_wires: 200,
             security_bits: 128,
         };
-        let mut builder = CircuitBuilder2::new(config);
-        builder.add_gate(gate, Vec::new());
-        // let circuit = builder.build();
 
-        // TODO: generate witness & compare output to normal GMiMC function.
+        let permutation_inputs = (0..W)
+            .map(F::from_canonical_usize)
+            .collect::<Vec<_>>();
+
+        let mut witness = PartialWitness::new();
+        witness.set_wire(Wire { gate: 0, input: Gate::WIRE_SWITCH }, F::ZERO);
+        for i in 0..W {
+            witness.set_wire(
+                Wire { gate: 0, input: Gate::wire_input(i) },
+                permutation_inputs[i]);
+        }
+
+        let generators = gate.0.generators(config, 0, vec![], vec![]);
+        generate_partial_witness(&mut witness, generators);
+
+        let expected_outputs: [F; W] = gmimc_permute_naive(
+            permutation_inputs.try_into().unwrap(),
+            constants);
+
+        for i in 0..W {
+            let out = witness.get_wire(
+                Wire { gate: 1, input: Gate::wire_output(i) });
+            assert_eq!(out, expected_outputs[i]);
+        }
     }
 }

src/generator.rs

@@ -1,9 +1,65 @@
+use std::borrow::Borrow;
+use std::collections::{HashMap, HashSet};
+
 use crate::field::field::Field;
 use crate::target::Target;
 use crate::witness::PartialWitness;
 
+pub(crate) fn generate_partial_witness<F: Field>(
+    witness: &mut PartialWitness<F>,
+    mut generators: Vec<Box<dyn WitnessGenerator<F>>>,
+) {
+    // Index generator indices by their watched targets.
+    let mut generator_indices_by_watches: HashMap<Target, Vec<usize>> = HashMap::new();
+    for (i, generator) in generators.iter().enumerate() {
+        for watch in generator.watch_list() {
+            generator_indices_by_watches
+                .entry(watch)
+                .or_insert_with(Vec::new)
+                .push(i);
+        }
+    }
+
+    // Build a list of "pending" generators which are queued to be run. Initially, all generators
+    // are queued.
+    let mut pending_generator_indices = HashSet::new();
+    for i in 0..generators.len() {
+        pending_generator_indices.insert(i);
+    }
+
+    // We also track a list of "expired" generators which have already returned false.
+    let mut expired_generator_indices = HashSet::new();
+
+    // Keep running generators until no generators are queued.
+    while !pending_generator_indices.is_empty() {
+        let mut next_pending_generator_indices = HashSet::new();
+
+        for &generator_idx in &pending_generator_indices {
+            let (result, finished) = generators[generator_idx].run(&witness);
+            if finished {
+                expired_generator_indices.insert(generator_idx);
+            }
+
+            // Enqueue unfinished generators that were watching one of the newly populated targets.
+            for watch in result.target_values.keys() {
+                if let Some(watching_generator_indices) = generator_indices_by_watches.get(watch) {
+                    for watching_generator_idx in watching_generator_indices {
+                        if !expired_generator_indices.contains(watching_generator_idx) {
+                            next_pending_generator_indices.insert(*watching_generator_idx);
+                        }
+                    }
+                }
+            }
+
+            witness.extend(result);
+        }
+
+        pending_generator_indices = next_pending_generator_indices;
+    }
+}
+
 /// A generator participates in the generation of the witness.
-pub trait WitnessGenerator2<F: Field>: 'static {
+pub trait WitnessGenerator<F: Field>: 'static {
     /// Targets to be "watched" by this generator. Whenever a target in the watch list is populated,
     /// the generator will be queued to run.
     fn watch_list(&self) -> Vec<Target>;
@@ -22,7 +78,7 @@ pub trait SimpleGenerator<F: Field>: 'static {
     fn run_once(&mut self, witness: &PartialWitness<F>) -> PartialWitness<F>;
 }
 
-impl<F: Field, SG: SimpleGenerator<F>> WitnessGenerator2<F> for SG {
+impl<F: Field, SG: SimpleGenerator<F>> WitnessGenerator<F> for SG {
     fn watch_list(&self) -> Vec<Target> {
         self.dependencies()
     }

src/main.rs

@@ -30,8 +30,8 @@ mod verifier;
 mod wire;
 mod witness;
 
-// 12 wire polys, 3 Z polys, 4 parts of quotient poly.
-const PROVER_POLYS: usize = 64 + 3 + (9 + 1); // TODO: Check
+// 112 wire polys, 3 Z polys, 4 parts of quotient poly.
+const PROVER_POLYS: usize = 113 + 3 + 4;
 
 fn main() {
     let overall_start = Instant::now();
@@ -56,7 +56,7 @@ fn bench_gmimc<F: Field>() {
     }
 
     const THREADS: usize = 12;
-    const LDE_BITS: i32 = 4;
+    const LDE_BITS: i32 = 3;
     const W: usize = 13;
     let hashes_per_poly = 1 << (13 + LDE_BITS);
     let threads = (0..THREADS).map(|_i| {

src/witness.rs

@@ -4,9 +4,9 @@ use crate::field::field::Field;
 use crate::target::Target;
 use crate::wire::Wire;
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct PartialWitness<F: Field> {
-    target_values: HashMap<Target, F>,
+    pub(crate) target_values: HashMap<Target, F>,
 }
 
 impl<F: Field> PartialWitness<F> {
@@ -53,4 +53,10 @@ impl<F: Field> PartialWitness<F> {
     pub fn set_wire(&mut self, wire: Wire, value: F) {
         self.set_target(Target::Wire(wire), value)
     }
+
+    pub fn extend(&mut self, other: PartialWitness<F>) {
+        for (target, value) in other.target_values {
+            self.set_target(target, value);
+        }
+    }
 }