Automatic signals generation (#152)

This is useful for:
 + Arbitrary boolean expressions
 + Multiplicative inverse elimination
 + Automatic witness generation in general
 + Conversion to lower degree polynomial identities

.github/workflows/rust.yml

@@ -49,7 +49,7 @@ jobs:
         uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: stable
+          toolchain: nightly-2023-04-24
           components: clippy
           override: true
       - name: Run Clippy

examples/fibonacci.rs

@@ -54,13 +54,15 @@ fn fibo_circuit<F: Field + From<u64> + Hash>() -> (Circuit<F>, Option<Assignment
 
             // in setup we define the constraints of the step
             ctx.setup(move |ctx| {
-                // regular constraints are for internal signals only
-                // constrain that a + b == c by calling `eq` function from constraint builder
-                ctx.constr(eq(a + b, c));
+                // regular constraints are for signals without rotation only
+
+                // `auto_eq` creates a constraint and also auto generates the witness of the left
+                // side.
+                ctx.auto_eq(c, a + b);
 
                 // transition constraints accepts forward signals as well
-                // constrain that b is equal to the next instance of a, by calling `next` on forward
-                // signal
+                // constrain that b is equal to the next instance of a by calling `eq` function from
+                // constraint builder and `next` on forward signal
                 ctx.transition(eq(b, a.next()));
                 // constrain that c is equal to the next instance of c, by calling `next` on forward
                 // signal
@@ -75,7 +77,8 @@ fn fibo_circuit<F: Field + From<u64> + Hash>() -> (Circuit<F>, Option<Assignment
                 // assign arbitrary input values from witness generation function to witnesses
                 ctx.assign(a, a_value.field());
                 ctx.assign(b, b_value.field());
-                ctx.assign(c, (a_value + b_value).field());
+
+                // c is auto generated by `auto_eq`
             })
         });
 

src/ast/mod.rs

@@ -204,6 +204,9 @@ pub struct StepType<F> {
     pub constraints: Vec<Constraint<F>>,
     pub transition_constraints: Vec<TransitionConstraint<F>>,
     pub lookups: Vec<Lookup<F>>,
+
+    pub auto_signals: HashMap<Queriable<F>, ASTExpr<F>>,
+
     pub annotations: HashMap<UUID, String>,
 }
 
@@ -227,9 +230,11 @@ impl<F> StepType<F> {
             constraints: Default::default(),
             transition_constraints: Default::default(),
             lookups: Default::default(),
+            auto_signals: Default::default(),
             annotations: Default::default(),
         }
     }
+
     pub fn uuid(&self) -> StepTypeUUID {
         self.id
     }

src/field.rs

@@ -43,4 +43,8 @@ pub trait Field:
     /// return zero if the element is zero. This is different from
     /// FF invert that returns None if the element is zero.
     fn mi(&self) -> Self;
+
+    /// Exponentiates `self` by `exp`, where `exp` is a little-endian order integer
+    /// exponent.
+    fn pow<S: AsRef<[u64]>>(&self, exp: S) -> Self;
 }

src/frontend/dsl/mod.rs

@@ -1,5 +1,5 @@
 use crate::{
-    ast::{query::Queriable, Circuit, ExposeOffset, StepType, StepTypeUUID},
+    ast::{query::Queriable, ASTExpr, Circuit, ExposeOffset, StepType, StepTypeUUID},
     field::Field,
     util::{uuid, UUID},
     wit_gen::{FixedGenContext, StepInstance, TraceContext},
@@ -11,7 +11,7 @@ use core::{fmt::Debug, hash::Hash};
 use std::marker::PhantomData;
 
 use self::{
-    cb::{Constraint, Typing},
+    cb::{eq, Constraint, Typing},
     lb::{LookupBuilder, LookupTable, LookupTableRegistry, LookupTableStore},
 };
 
@@ -310,6 +310,18 @@ impl<'a, F> StepTypeSetupContext<'a, F> {
     }
 }
 
+impl<'a, F: Eq + PartialEq + Hash + Debug + Clone> StepTypeSetupContext<'a, F> {
+    pub fn auto(&mut self, signal: Queriable<F>, expr: ASTExpr<F>) {
+        self.step_type.auto_signals.insert(signal, expr);
+    }
+
+    pub fn auto_eq(&mut self, signal: Queriable<F>, expr: ASTExpr<F>) {
+        self.auto(signal.clone(), expr.clone());
+
+        self.constr(eq(signal, expr));
+    }
+}
+
 impl<'a, F: Debug + Clone> StepTypeSetupContext<'a, F> {
     /// Adds a lookup to the step type.
     pub fn add_lookup<LB: LookupBuilder<F>>(&mut self, lookup_builder: LB) {

src/plonkish/backend/halo2.rs

@@ -30,6 +30,10 @@ impl<T: Field + From<u64>> ChiquitoField for T {
     fn mi(&self) -> Self {
         self.invert().unwrap_or(Self::ZERO)
     }
+
+    fn pow<S: AsRef<[u64]>>(&self, exp: S) -> Self {
+        Field::pow(self, exp)
+    }
 }
 
 #[allow(non_snake_case)]

src/plonkish/compiler/mod.rs

@@ -6,7 +6,7 @@ use crate::{
         Circuit, Column, Poly, PolyExpr, PolyLookup,
     },
     poly::Expr,
-    wit_gen::{FixedAssignment, TraceGenerator},
+    wit_gen::{AutoTraceGenerator, FixedAssignment, TraceGenerator},
 };
 use std::{hash::Hash, rc::Rc};
 
@@ -89,6 +89,7 @@ pub fn compile_phase1<
             unit.placement.clone(),
             unit.selector.clone(),
             TraceGenerator::new(Rc::clone(v), ast.num_steps),
+            AutoTraceGenerator::from(ast),
             unit.num_rows,
             unit.uuid,
         )

src/plonkish/ir/assignments.rs

@@ -1,10 +1,11 @@
 use std::{
     collections::HashMap,
     fmt,
+    hash::Hash,
     ops::{Deref, DerefMut},
 };
 
-use crate::field::Field;
+use crate::{field::Field, wit_gen::AutoTraceGenerator};
 
 use halo2_proofs::plonk::{Advice, Column as Halo2Column};
 
@@ -68,6 +69,7 @@ pub struct AssignmentGenerator<F, TraceArgs> {
     placement: Placement,
     selector: StepSelector<F>,
     trace_gen: TraceGenerator<F, TraceArgs>,
+    auto_trace_gen: AutoTraceGenerator<F>,
 
     num_rows: usize,
 
@@ -81,6 +83,7 @@ impl<F: Clone, TraceArgs> Clone for AssignmentGenerator<F, TraceArgs> {
             placement: self.placement.clone(),
             selector: self.selector.clone(),
             trace_gen: self.trace_gen.clone(),
+            auto_trace_gen: self.auto_trace_gen.clone(),
             num_rows: self.num_rows,
             ir_id: self.ir_id,
         }
@@ -94,18 +97,20 @@ impl<F: Clone, TraceArgs> Default for AssignmentGenerator<F, TraceArgs> {
             placement: Default::default(),
             selector: Default::default(),
             trace_gen: Default::default(),
+            auto_trace_gen: Default::default(),
             num_rows: Default::default(),
             ir_id: Default::default(),
         }
     }
 }
 
-impl<F: Field, TraceArgs> AssignmentGenerator<F, TraceArgs> {
+impl<F: Field + Hash, TraceArgs> AssignmentGenerator<F, TraceArgs> {
     pub fn new(
         columns: Vec<Column>,
         placement: Placement,
         selector: StepSelector<F>,
         trace_gen: TraceGenerator<F, TraceArgs>,
+        auto_trace_gen: AutoTraceGenerator<F>,
         num_rows: usize,
         ir_id: UUID,
     ) -> Self {
@@ -114,6 +119,7 @@ impl<F: Field, TraceArgs> AssignmentGenerator<F, TraceArgs> {
             placement,
             selector,
             trace_gen,
+            auto_trace_gen,
             num_rows,
             ir_id,
         }
@@ -136,6 +142,8 @@ impl<F: Field, TraceArgs> AssignmentGenerator<F, TraceArgs> {
         let mut offset: usize = 0;
         let mut assignments: Assignments<F> = Default::default();
 
+        let witness = self.auto_trace_gen.generate(witness);
+
         for step_instance in witness.step_instances.into_iter() {
             self.assign_step(&mut offset, &mut assignments, &step_instance);
         }

src/plonkish/ir/sc.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, rc::Rc};
+use std::{collections::HashMap, hash::Hash, rc::Rc};
 
 use crate::{field::Field, util::UUID, wit_gen::TraceWitness};
 
@@ -31,7 +31,7 @@ impl<F, MappingArgs> SuperCircuit<F, MappingArgs> {
     }
 }
 
-impl<F: Field, MappingArgs> SuperCircuit<F, MappingArgs> {
+impl<F: Field + Hash, MappingArgs> SuperCircuit<F, MappingArgs> {
     pub fn set_mapping<M: Fn(&mut MappingContext<F>, MappingArgs) + 'static>(
         &mut self,
         mapping: M,
@@ -60,7 +60,7 @@ impl<F> Default for MappingContext<F> {
     }
 }
 
-impl<F: Field> MappingContext<F> {
+impl<F: Field + Hash> MappingContext<F> {
     pub fn map<TraceArgs>(&mut self, gen: &AssignmentGenerator<F, TraceArgs>, args: TraceArgs) {
         self.assignments.insert(gen.uuid(), gen.generate(args));
     }
@@ -101,7 +101,7 @@ impl<F, MappingArgs> Default for MappingGenerator<F, MappingArgs> {
     }
 }
 
-impl<F: Field, MappingArgs> MappingGenerator<F, MappingArgs> {
+impl<F: Field + Hash, MappingArgs> MappingGenerator<F, MappingArgs> {
     pub fn new(mapping: Rc<Mapping<F, MappingArgs>>) -> Self {
         Self { mapping }
     }

src/poly.rs

@@ -1,5 +1,7 @@
 use std::{
+    collections::HashMap,
     fmt::Debug,
+    hash::Hash,
     ops::{Add, Mul, Neg, Sub},
 };
 
@@ -65,6 +67,28 @@ impl<F: Debug, V: Debug> Debug for Expr<F, V> {
     }
 }
 
+pub type VarAssignments<F, V> = HashMap<V, F>;
+
+impl<F: Field + Hash, V: Eq + PartialEq + Hash> Expr<F, V> {
+    pub fn eval(&self, assignments: &VarAssignments<F, V>) -> Option<F> {
+        match self {
+            Expr::Const(v) => Some(*v),
+            Expr::Sum(ses) => ses
+                .iter()
+                .fold(Some(F::ZERO), |acc, se| Some(acc? + se.eval(assignments)?)),
+            Expr::Mul(ses) => ses
+                .iter()
+                .fold(Some(F::ONE), |acc, se| Some(acc? * se.eval(assignments)?)),
+            Expr::Neg(se) => Some(F::ZERO - se.eval(assignments)?),
+            Expr::Pow(se, exp) => Some(se.eval(assignments)?.pow([*exp as u64])),
+            Expr::Query(q) => assignments.get(q).copied(),
+
+            // Not implemented, and not necessary for aexpr
+            Expr::Halo2Expr(_) => None,
+        }
+    }
+}
+
 impl<F: Clone, V: Clone> ToExpr<F, V> for Expr<F, V> {
     fn expr(&self) -> Expr<F, V> {
         self.clone()
@@ -186,3 +210,59 @@ impl<F, V> From<Expression<F>> for Expr<F, V> {
         Expr::Halo2Expr(value)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use halo2curves::bn256::Fr;
+
+    use crate::{field::Field, poly::VarAssignments};
+
+    use super::Expr;
+
+    #[test]
+    fn eval_const() {
+        use super::Expr::*;
+
+        let experiment: Expr<Fr, String> = Const(Fr::ONE);
+        let assignments: VarAssignments<Fr, String> = VarAssignments::default();
+
+        assert_eq!(experiment.eval(&assignments), Some(Fr::ONE))
+    }
+
+    #[test]
+    fn eval_var() {
+        use super::Expr::*;
+
+        let experiment: Expr<Fr, &str> = Query("a");
+        let mut assignments: VarAssignments<Fr, &str> = VarAssignments::default();
+        assignments.insert("a", Fr::ONE);
+
+        assert_eq!(experiment.eval(&assignments), Some(Fr::ONE))
+    }
+
+    #[test]
+    fn eval_expr() {
+        use super::Expr::*;
+
+        let experiment: Expr<Fr, &str> = (Query("a") * Query("b")) + Query("c") - Const(Fr::ONE);
+        let mut assignments: VarAssignments<Fr, &str> = VarAssignments::default();
+        assignments.insert("a", Fr::from(2));
+        assignments.insert("b", Fr::from(3));
+        assignments.insert("c", Fr::from(4));
+
+        assert_eq!(experiment.eval(&assignments), Some(Fr::from(9)))
+    }
+
+    #[test]
+    fn eval_expr_missing_var() {
+        use super::Expr::*;
+
+        let experiment: Expr<Fr, &str> = (Query("a") * Query("b")) + Query("c") - Const(Fr::ONE);
+        let mut assignments: VarAssignments<Fr, &str> = VarAssignments::default();
+        assignments.insert("a", Fr::from(2));
+        // REMOVE assignments.insert("b", Fr::from(3));
+        assignments.insert("c", Fr::from(4));
+
+        assert_eq!(experiment.eval(&assignments), None)
+    }
+}