From 5acbb674ad1ad13e6baf28a2c19180f7c43a1550 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Tue, 22 Jun 2021 16:54:20 +0200
Subject: [PATCH] Add prefix to constant polys

---
 src/circuit_builder.rs | 28 +++++++++++++++++++---------
 src/gates/gate.rs      | 14 ++++++++++++++
 src/gates/gate_tree.rs |  7 +++----
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/src/circuit_builder.rs b/src/circuit_builder.rs
index 3a26497e..c91ef470 100644
--- a/src/circuit_builder.rs
+++ b/src/circuit_builder.rs
@@ -11,7 +11,8 @@ use crate::field::cosets::get_unique_coset_shifts;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
 use crate::gates::constant::ConstantGate;
-use crate::gates::gate::{GateInstance, GateRef};
+use crate::gates::gate::{GateInstance, GatePrefixes, GateRef};
+use crate::gates::gate_tree::Tree;
 use crate::gates::noop::NoopGate;
 use crate::generator::{CopyGenerator, WitnessGenerator};
 use crate::hash::hash_n_to_hash;
@@ -229,22 +230,28 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         }
     }
 
-    fn constant_polys(&self) -> Vec<PolynomialValues<F>> {
+    fn constant_polys(&self, prefixes: &GatePrefixes<F, D>) -> Vec<PolynomialValues<F>> {
         let num_constants = self
             .gate_instances
             .iter()
-            .map(|gate_inst| gate_inst.constants.len())
+            .map(|gate_inst| gate_inst.constants.len() + prefixes[&gate_inst.gate_type].len())
             .max()
             .unwrap();
         let constants_per_gate = self
             .gate_instances
             .iter()
             .map(|gate_inst| {
-                let mut padded_constants = gate_inst.constants.clone();
-                for _ in padded_constants.len()..num_constants {
-                    padded_constants.push(F::ZERO);
-                }
-                padded_constants
+                let mut prefixed_constants = Vec::new();
+                prefixed_constants.extend(prefixes[&gate_inst.gate_type].iter().map(|&b| {
+                    if b {
+                        F::ONE
+                    } else {
+                        F::ZERO
+                    }
+                }));
+                prefixed_constants.extend_from_slice(&gate_inst.constants);
+                prefixed_constants.resize(num_constants, F::ZERO);
+                prefixed_constants
             })
             .collect::<Vec<_>>();
 
@@ -288,7 +295,10 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let degree = self.gate_instances.len();
         info!("degree after blinding & padding: {}", degree);
 
-        let constant_vecs = self.constant_polys();
+        let gates = self.gates.iter().cloned().collect();
+        let gate_tree = Tree::from_gates(gates);
+
+        let constant_vecs = self.constant_polys(&gate_tree.into());
         let constants_commitment = ListPolynomialCommitment::new(
             constant_vecs.into_iter().map(|v| v.ifft()).collect(),
             self.config.fri_config.rate_bits,
diff --git a/src/gates/gate.rs b/src/gates/gate.rs
index 9219d5dc..cf939d9d 100644
--- a/src/gates/gate.rs
+++ b/src/gates/gate.rs
@@ -1,7 +1,9 @@
+use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::fmt::{Debug, Error, Formatter};
 use std::hash::{Hash, Hasher};
 use std::iter::FromIterator;
+use std::ops::Index;
 use std::sync::Arc;
 
 use crate::circuit_builder::CircuitBuilder;
@@ -142,3 +144,15 @@ impl<F: Extendable<D>, const D: usize> From<Tree<GateRef<F, D>>> for GatePrefixe
         }
     }
 }
+
+impl<F: Extendable<D>, T: Borrow<GateRef<F, D>>, const D: usize> Index<T> for GatePrefixes<F, D> {
+    type Output = Vec<bool>;
+
+    fn index(&self, index: T) -> &Self::Output {
+        &self.prefixes[index.borrow()]
+    }
+}
+
+// impl<F: Extendable<D>, const D: usize> GatePrefixes<F, D> {
+//     pub fn prefix_len()
+// }
diff --git a/src/gates/gate_tree.rs b/src/gates/gate_tree.rs
index 8081ab25..9195e30c 100644
--- a/src/gates/gate_tree.rs
+++ b/src/gates/gate_tree.rs
@@ -1,3 +1,5 @@
+use log::info;
+
 use crate::field::extension_field::Extendable;
 use crate::gates::gate::GateRef;
 
@@ -64,7 +66,7 @@ impl<F: Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
         for max_degree in 1..100 {
             if let Some(mut tree) = Self::find_tree(&gates, max_degree) {
                 tree.shorten();
-                println!(
+                info!(
                     "Found tree with max degree {} in {}s.",
                     max_degree,
                     timer.elapsed().as_secs_f32()
@@ -199,9 +201,6 @@ mod tests {
 
         let tree = Tree::from_gates(gates.clone());
         let mut gates_with_prefix = tree.traversal();
-        for (g, p) in &gates_with_prefix {
-            println!("{}: {:?}", g.0.id(), p);
-        }
 
         assert_eq!(
             gates_with_prefix.len(),