From 0a85e9371425d63e4597d5906af9c6b364011030 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Wed, 2 Dec 2020 03:48:19 +0800
Subject: [PATCH] Add lookup to circuit and test

---
 src/plonk.rs         | 65 ++++++++++++++++++++++++++++++++++++++++----
 src/plonk/circuit.rs | 25 +++++++++++++++++
 2 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/src/plonk.rs b/src/plonk.rs
index 5df7d86..c6a8be1 100644
--- a/src/plonk.rs
+++ b/src/plonk.rs
@@ -121,7 +121,7 @@ fn test_proving() {
     use crate::poly::commitment::{Blind, Params};
     use crate::transcript::DummyHash;
     use crate::tweedle::{EqAffine, Fp, Fq};
-    use circuit::{Advice, Column, Fixed};
+    use circuit::{Advice, Any, Column, Fixed};
     use std::marker::PhantomData;
     const K: u32 = 5;
 
@@ -144,6 +144,8 @@ fn test_proving() {
         sc: Column<Fixed>,
         sm: Column<Fixed>,
         sp: Column<Fixed>,
+        sl: Column<Fixed>,
+        sl2: Column<Fixed>,
 
         perm: usize,
         perm2: usize,
@@ -160,10 +162,12 @@ fn test_proving() {
         fn public_input<F>(&mut self, f: F) -> Result<Variable, Error>
         where
             F: FnOnce() -> Result<FF, Error>;
+        fn lookup_table(&mut self, values: &[Vec<FF>]) -> Result<(), Error>;
     }
 
     struct MyCircuit<F: FieldExt> {
         a: Option<F>,
+        lookup_tables: Vec<Vec<F>>,
     }
 
     struct StandardPLONK<'a, F: FieldExt, CS: Assignment<F> + 'a> {
@@ -297,6 +301,18 @@ fn test_proving() {
 
             Ok(Variable(self.config.a, index))
         }
+        fn lookup_table(&mut self, values: &[Vec<FF>]) -> Result<(), Error> {
+            for (&value_0, &value_1) in values[0].iter().zip(values[1].iter()) {
+                let index = self.current_gate;
+
+                self.current_gate += 1;
+                self.cs
+                    .assign_fixed(self.config.sl, index, || Ok(value_0))?;
+                self.cs
+                    .assign_fixed(self.config.sl2, index, || Ok(value_1))?;
+            }
+            Ok(())
+        }
     }
 
     impl<F: FieldExt> Circuit<F> for MyCircuit<F> {
@@ -319,6 +335,30 @@ fn test_proving() {
             let sb = meta.fixed_column();
             let sc = meta.fixed_column();
             let sp = meta.fixed_column();
+            let sl = meta.fixed_column();
+            let sl2 = meta.fixed_column();
+
+            /*
+             *    A    B        ...  sl   sl2
+             * [
+             *   aux   0        ...  0    0
+             *   a     a        ...  0    0
+             *   a     a^2      ...  0    0
+             *   a     a        ...  0    0
+             *   a     a^2      ...  0    0
+             *   ...   ...      ...  ...  ...
+             *   ...   ...      ...  aux  0
+             *   ...   ...      ...  a    a
+             *   ...   ...      ...  a    a^2
+             *   ...   ...      ...  0    0
+             *
+             * ]
+             */
+            meta.lookup(&[Column::<Any>::from(a)], &[Column::<Any>::from(sl)]);
+            meta.lookup(
+                &[Column::<Any>::from(a), Column::<Any>::from(b)],
+                &[Column::<Any>::from(sl), Column::<Any>::from(sl2)],
+            );
 
             meta.create_gate(|meta| {
                 let d = meta.query_advice(d, 1);
@@ -355,6 +395,8 @@ fn test_proving() {
                 sc,
                 sm,
                 sp,
+                sl,
+                sl2,
                 perm,
                 perm2,
             }
@@ -391,22 +433,33 @@ fn test_proving() {
                 cs.copy(b1, c0)?;
             }
 
+            cs.lookup_table(&self.lookup_tables)?;
+
             Ok(())
         }
     }
 
-    let circuit: MyCircuit<Fp> = MyCircuit {
-        a: Some(Fp::rand()),
+    let a = Fp::rand();
+    let a_squared = a * &a;
+    let aux = Fp::one() + Fp::one();
+    let lookup_table = vec![aux, a, a, Fp::zero()];
+    let lookup_table_2 = vec![Fp::zero(), a, a_squared, Fp::zero()];
+
+    let empty_circuit: MyCircuit<Fp> = MyCircuit {
+        a: None,
+        lookup_tables: vec![lookup_table.clone(), lookup_table_2.clone()],
     };
 
-    let empty_circuit: MyCircuit<Fp> = MyCircuit { a: None };
+    let circuit: MyCircuit<Fp> = MyCircuit {
+        a: Some(a),
+        lookup_tables: vec![lookup_table, lookup_table_2],
+    };
 
     // Initialize the proving key
     let pk = keygen(&params, &empty_circuit).expect("keygen should not fail");
 
     let mut pubinputs = pk.get_vk().get_domain().empty_lagrange();
-    pubinputs[0] = Fp::one();
-    pubinputs[0] += Fp::one();
+    pubinputs[0] = aux;
     let pubinput = params
         .commit_lagrange(&pubinputs, Blind::default())
         .to_affine();
diff --git a/src/plonk/circuit.rs b/src/plonk/circuit.rs
index 9db8f42..3b1cd45 100644
--- a/src/plonk/circuit.rs
+++ b/src/plonk/circuit.rs
@@ -358,6 +358,31 @@ impl<F: Field> ConstraintSystem<F> {
         index
     }
 
+    /// Add a lookup argument for some input columns and table columns.
+    pub fn lookup(
+        &mut self,
+        input_columns: &[Column<Any>],
+        table_columns: &[Column<Any>],
+    ) -> usize {
+        let index = self.lookups.len();
+        if self.lookups.is_empty() {
+            let at = Rotation(-1);
+            let len = self.rotations.len();
+            self.rotations.entry(at).or_insert(PointIndex(len));
+        }
+
+        for input in input_columns {
+            self.query_any_index(*input, 0);
+        }
+        for table in table_columns {
+            self.query_any_index(*table, 0);
+        }
+        self.lookups
+            .push(lookup::Argument::new(input_columns, table_columns));
+
+        index
+    }
+
     fn query_fixed_index(&mut self, column: Column<Fixed>, at: i32) -> usize {
         let at = Rotation(at);
         {