kind of working; a few tests are failing

src/backends/kimchi.rs

@@ -1,12 +1,13 @@
 use std::fmt::Write;
+use std::sync::Arc;
 use std::{
     collections::{HashMap, HashSet},
     iter::once,
     ops::Neg,
 };
 
 use ark_ff::Field;
-use itertools::{chain, Itertools};
+use itertools::{chain, izip, Itertools};
 use kimchi::mina_curves::pasta::{Pallas, Vesta, VestaParameters};
 use kimchi::mina_poseidon::sponge::{DefaultFqSponge, DefaultFrSponge};
 use kimchi::poly_commitment::srs::SRS;
@@ -53,7 +54,7 @@ type SpongeParams = PlonkSpongeConstantsKimchi;
 type BaseSponge = DefaultFqSponge<VestaParameters, SpongeParams>;
 type ScalarSponge = DefaultFrSponge<kimchi::mina_curves::pasta::Fp, SpongeParams>;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Kimchi {
     /// This is used to give a distinct number to each variable during circuit generation.
     pub(crate) next_variable: usize,
@@ -282,7 +283,7 @@ impl Backend for Kimchi {
                 for col in 0..3 {
                     // create each variable
                     let var = self.new_internal_var(
-                        Value::Hint(Box::new(move |compiler, env| {
+                        Value::Hint(Arc::new(move |compiler, env| {
                             let x1 = compiler.compute_var(env, prev_0)?;
                             let x2 = compiler.compute_var(env, prev_1)?;
                             let x3 = compiler.compute_var(env, prev_2)?;
@@ -798,7 +799,7 @@ impl Backend for Kimchi {
 
                 // compute the result
                 let res = self.new_internal_var(
-                    Value::Hint(Box::new(move |compiler, env| {
+                    Value::Hint(Arc::new(move |compiler, env| {
                         let x1 = compiler.compute_var(env, x1)?;
                         let x2 = compiler.compute_var(env, x2)?;
                         if x1 == x2 {
@@ -959,7 +960,7 @@ impl Backend for Kimchi {
                 let else_clone = *else_;
 
                 let res = self.new_internal_var(
-                    Value::Hint(Box::new(move |compiler, env| {
+                    Value::Hint(Arc::new(move |compiler, env| {
                         let cond = compiler.compute_var(env, cond_cell)?;
                         let res_var = if cond.is_one() {
                             &then_clone
@@ -1408,4 +1409,98 @@ impl Backend for Kimchi {
         // return asm + indexes
         Ok((prover_index, verifier_index))
     }
+
+    fn size(&self) -> usize {
+        self.gates.len()
+    }
+
+    fn private_input_indices(&self) -> &Vec<usize> {
+        self.private_input_indices.as_ref()
+    }
+
+    fn public_output(&self) -> &Option<Var<Self::Field>> {
+        &self.public_output
+    }
+
+    fn public_input_size(&self) -> usize {
+        self.public_input_size
+    }
+
+    fn add_witness(&mut self, var_idx: usize, val: Value<Self>) -> Option<Value<Self>> {
+        self.witness_vars.insert(var_idx, val)
+    }
+
+    fn generate_witness(
+        &self,
+        mut env: crate::witness::WitnessEnv<Self::Field>,
+        witness: &mut Vec<[Self::Field; 15]>,
+    ) -> Result<(Vec<Self::Field>, Vec<Self::Field>)> {
+        let mut public_outputs_vars: Vec<(usize, CellVar)> = vec![];
+
+        for (row, (gate, row_of_vars, debug_info)) in
+            izip!(&self.gates, &self.rows_of_vars(), &self.debug_info).enumerate()
+        {
+            // create the witness row
+            let mut witness_row = [Self::Field::zero(); NUM_REGISTERS];
+
+            for (col, var) in row_of_vars.iter().enumerate() {
+                let val = if let Some(var) = var {
+                    // if it's a public output, defer it's computation
+                    if matches!(
+                        self.witness_vars[&var.index],
+                        Value::PublicOutput(_)
+                    ) {
+                        public_outputs_vars.push((row, *var));
+                        Self::Field::zero()
+                    } else {
+                        self.compute_var(&mut env, *var)?
+                    }
+                } else {
+                    Self::Field::zero()
+                };
+                witness_row[col] = val;
+            }
+
+            // check if the row makes sense
+            let is_not_public_input = row >= self.public_input_size;
+            if is_not_public_input {
+                #[allow(clippy::single_match)]
+                match gate.typ {
+                    // only check the generic gate
+                    crate::circuit_writer::GateKind::DoubleGeneric => {
+                        let c = |i| gate.coeffs.get(i).copied().unwrap_or_else(Self::Field::zero);
+                        let w = &witness_row;
+                        let sum1 =
+                            c(0) * w[0] + c(1) * w[1] + c(2) * w[2] + c(3) * w[0] * w[1] + c(4);
+                        let sum2 =
+                            c(5) * w[3] + c(6) * w[4] + c(7) * w[5] + c(8) * w[3] * w[4] + c(9);
+                        if sum1 != Self::Field::zero() || sum2 != Self::Field::zero() {
+                            return Err(Error::new(
+                                "runtime",
+                                ErrorKind::InvalidWitness(row),
+                                debug_info.span,
+                            ));
+                        }
+                    }
+                    // for all other gates, we trust the gadgets
+                    _ => (),
+                }
+            }
+
+            //
+            witness.push(witness_row);
+        }
+        let mut public_output = vec![];
+        for (row, var) in public_outputs_vars {
+            let val = self.compute_var(&mut env, var)?;
+            witness[row][0] = val;
+            public_output.push(val);
+        }
+        let mut full_public_inputs = Vec::with_capacity(self.public_input_size);
+        for witness_row in witness.iter().take(self.public_input_size) {
+            full_public_inputs.push(witness_row[0]);
+        }
+        Ok((public_output, full_public_inputs))
+    }
+
 }

src/backends/mod.rs

@@ -1,26 +1,30 @@
 use std::collections::HashMap;
 
+use ::kimchi::{
+    mina_curves::pasta::Vesta, poly_commitment::evaluation_proof::OpeningProof,
+    prover_index::ProverIndex, verifier_index::VerifierIndex,
+};
 use ark_ff::Field;
-use ::kimchi::{mina_curves::pasta::Vesta, poly_commitment::evaluation_proof::OpeningProof, prover_index::ProverIndex, verifier_index::VerifierIndex};
 
 use crate::{
     circuit_writer::{CircuitWriter, DebugInfo, GateKind, VarInfo},
     compiler::Sources,
     constants::{Span, NUM_REGISTERS},
-    error::Result,
+    error::{Error, ErrorKind, Result},
     helpers::PrettyField,
     var::{CellVar, ConstOrCell, Value, Var},
-    witness::{CompiledCircuit, Witness},
+    witness::{CompiledCircuit, Witness, WitnessEnv},
 };
 
+use num_traits::Zero;
+
 pub mod kimchi;
 pub mod r1cs;
 
 type Curve = Vesta;
 
-
 /// Each backend implements this trait
-pub trait Backend {
+pub trait Backend: Clone {
     /// The circuit field / scalar field that the circuit is written on.
     type Field: Field + PrettyField;
 
@@ -34,6 +38,58 @@ pub trait Backend {
 
     fn next_variable(&mut self) -> usize;
 
+    fn size(&self) -> usize;
+
+    fn compute_var(&self, env: &mut WitnessEnv<Self::Field>, var: CellVar) -> Result<Self::Field> {
+        // fetch cache first
+        // TODO: if self was &mut, then we could use a Value::Cached(Field) to store things instead of that
+        if let Some(res) = env.cached_values.get(&var) {
+            return Ok(*res);
+        }
+
+        match &self.witness_vars()[&var.index] {
+            Value::Hint(func) => {
+                let res = func(self, env)
+                    .expect("that function doesn't return a var (type checker error)");
+                env.cached_values.insert(var, res);
+                Ok(res)
+            }
+            Value::Constant(c) => Ok(*c),
+            Value::LinearCombination(lc, cst) => {
+                let mut res = *cst;
+                for (coeff, var) in lc {
+                    res += self.compute_var(env, *var)? * *coeff;
+                }
+                env.cached_values.insert(var, res); // cache
+                Ok(res)
+            }
+            Value::Mul(lhs, rhs) => {
+                let lhs = self.compute_var(env, *lhs)?;
+                let rhs = self.compute_var(env, *rhs)?;
+                let res = lhs * rhs;
+                env.cached_values.insert(var, res); // cache
+                Ok(res)
+            }
+            Value::Inverse(v) => {
+                let v = self.compute_var(env, *v)?;
+                let res = v.inverse().unwrap_or_else(Self::Field::zero);
+                env.cached_values.insert(var, res); // cache
+                Ok(res)
+            }
+            Value::External(name, idx) => Ok(env.get_external(name)[*idx]),
+            Value::PublicOutput(var) => {
+                let var = var.ok_or_else(|| {
+                    Error::new("runtime", ErrorKind::MissingReturn, Span::default())
+                })?;
+                self.compute_var(env, var)
+            }
+            Value::Scale(scalar, var) => {
+                let var = self.compute_var(env, *var)?;
+                Ok(*scalar * var)
+            }
+        }
+    }
+
     fn witness_vars(&self) -> &HashMap<usize, Value<Self>>
     where
         Self: Sized;
@@ -42,6 +98,12 @@ pub trait Backend {
     where
         Self: Sized;
 
+    fn private_input_indices(&self) -> &Vec<usize>;
+
+    fn public_output(&self) -> &Option<Var<Self::Field>>;
+
+    fn public_input_size(&self) -> usize;
+
     fn new_internal_var(&mut self, val: Value<Self>, span: Span) -> CellVar
     where
         Self: Sized;
@@ -52,6 +114,8 @@ pub trait Backend {
 
     fn add_private_inputs(&mut self, name: String, num: usize, span: Span) -> Var<Self::Field>;
 
+    fn add_witness(&mut self, var_idx: usize, val: Value<Self>) -> Option<Value<Self>>;
+
     fn add_constant(
         &mut self,
         label: Option<&'static str>,
@@ -101,8 +165,12 @@ pub trait Backend {
         span: Span,
     ) -> Var<Self::Field>;
 
-    fn equal(&mut self, lhs: &Var<Self::Field>, rhs: &Var<Self::Field>, span: Span)
-        -> Var<Self::Field>;
+    fn equal(
+        &mut self,
+        lhs: &Var<Self::Field>,
+        rhs: &Var<Self::Field>,
+        span: Span,
+    ) -> Var<Self::Field>;
 
     fn equal_cells(
         &mut self,
@@ -151,4 +219,10 @@ pub trait Backend {
         ProverIndex<Curve, OpeningProof<Curve>>,
         VerifierIndex<Curve, OpeningProof<Curve>>,
     )>;
+
+    fn generate_witness(
+        &self,
+        env: WitnessEnv<Self::Field>,
+        witness: &mut Vec<[Self::Field; 15]>,
+    ) -> Result<(Vec<Self::Field>, Vec<Self::Field>)>;
 }

src/circuit_writer/mod.rs

@@ -53,33 +53,33 @@ where
     // Note: I don't think we need this, but it acts as a nice redundant failsafe.
     pub(crate) finalized: bool,
 
-    /// This is used to give a distinct number to each variable during circuit generation.
-    pub(crate) next_variable: usize,
-
-    /// This is how you compute the value of each variable during witness generation.
-    /// It is created during circuit generation.
-    pub(crate) witness_vars: HashMap<usize, Value<B>>,
-
-    /// Size of the public input.
-    pub(crate) public_input_size: usize,
-
-    /// If a public output is set, this will be used to store its [Var].
-    /// The public output generation works as follows:
-    /// 1. This cvar is created and inserted in the circuit (gates) during compilation of the public input
-    ///    (as the public output is the end of the public input)
-    /// 2. When the `return` statement of the circuit is parsed,
-    ///    it will set this `public_output` variable again to the correct vars.
-    /// 3. During witness generation, the public output computation
-    ///    is delayed until the very end.
-    pub(crate) public_output: Option<Var<B::Field>>,
-
-    /// Indexes used by the private inputs
-    /// (this is useful to check that they appear in the circuit)
-    pub(crate) private_input_indices: Vec<usize>,
-
-    /// We cache the association between a constant and its _constrained_ variable,
-    /// this is to avoid creating a new constraint every time we need to hardcode the same constant.
-    pub(crate) cached_constants: HashMap<B::Field, CellVar>,
+    // /// This is used to give a distinct number to each variable during circuit generation.
+    // pub(crate) next_variable: usize,
+
+    // /// This is how you compute the value of each variable during witness generation.
+    // /// It is created during circuit generation.
+    // pub(crate) witness_vars: HashMap<usize, Value<B>>,
+
+    // /// Size of the public input.
+    // pub(crate) public_input_size: usize,
+
+    // /// If a public output is set, this will be used to store its [Var].
+    // /// The public output generation works as follows:
+    // /// 1. This cvar is created and inserted in the circuit (gates) during compilation of the public input
+    // ///    (as the public output is the end of the public input)
+    // /// 2. When the `return` statement of the circuit is parsed,
+    // ///    it will set this `public_output` variable again to the correct vars.
+    // /// 3. During witness generation, the public output computation
+    // ///    is delayed until the very end.
+    // pub(crate) public_output: Option<Var<B::Field>>,
+
+    // /// Indexes used by the private inputs
+    // /// (this is useful to check that they appear in the circuit)
+    // pub(crate) private_input_indices: Vec<usize>,
+
+    // /// We cache the association between a constant and its _constrained_ variable,
+    // /// this is to avoid creating a new constraint every time we need to hardcode the same constant.
+    // pub(crate) cached_constants: HashMap<B::Field, CellVar>,
 
     /// A vector of debug information that maps to each row of the created circuit.
     pub(crate) debug_info: Vec<DebugInfo>,
@@ -168,12 +168,6 @@ impl<B: Backend> CircuitWriter<B> {
             typed,
             backend,
             finalized: false,
-            next_variable: 0,
-            witness_vars: HashMap::new(),
-            public_input_size: 0,
-            public_output: None,
-            private_input_indices: vec![],
-            cached_constants: HashMap::new(),
             debug_info: vec![],
         }
     }
@@ -236,7 +230,7 @@ impl<B: Backend> CircuitWriter<B> {
                 }
                 circuit_writer.backend.add_public_inputs(name.value.clone(), len, name.span)
             } else {
-                circuit_writer.add_private_inputs(name.value.clone(), len, name.span)
+                circuit_writer.backend.add_private_inputs(name.value.clone(), len, name.span)
             };
 
             // constrain what needs to be constrained
@@ -286,9 +280,9 @@ impl<B: Backend> CircuitWriter<B> {
             });
         }
 
-        for var in 0..circuit_writer.next_variable {
+        for var in 0..circuit_writer.backend.next_variable() {
             if !written_vars.contains(&var) {
-                if circuit_writer.private_input_indices.contains(&var) {
+                if circuit_writer.backend.private_input_indices().contains(&var) {
                     // compute main sig
                     let (_main_sig, main_span) = {
                         let fn_info = circuit_writer.main_info()?.clone();