Working gate tree generation

src/circuit_builder.rs

@@ -12,6 +12,7 @@ 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_tree::{GatePrefixes, Tree};
 use crate::gates::noop::NoopGate;
 use crate::generator::{CopyGenerator, WitnessGenerator};
 use crate::hash::hash_n_to_hash;
@@ -279,6 +280,12 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
     /// Builds a "full circuit", with both prover and verifier data.
     pub fn build(mut self) -> CircuitData<F, D> {
+        let gates = self.gates.iter().cloned().collect();
+        let tree = Tree::from_gates(gates);
+        let prefixes = GatePrefixes::from(tree);
+        for (g, p) in &prefixes.prefixes {
+            println!("{}: {:?}", g.0.id(), p);
+        }
         let start = Instant::now();
         info!(
             "degree before blinding & padding: {}",

src/gates/gate.rs

@@ -1,3 +1,4 @@
+use std::fmt::{Debug, Error, Formatter};
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
@@ -113,6 +114,12 @@ impl<F: Extendable<D>, const D: usize> Hash for GateRef<F, D> {
 
 impl<F: Extendable<D>, const D: usize> Eq for GateRef<F, D> {}
 
+impl<F: Extendable<D>, const D: usize> Debug for GateRef<F, D> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
+        write!(f, "{}", self.0.id())
+    }
+}
+
 /// A gate along with any constants used to configure it.
 pub struct GateInstance<F: Extendable<D>, const D: usize> {
     pub gate_type: GateRef<F, D>,

src/gates/gate_tree.rs

@@ -0,0 +1,156 @@
+use std::collections::HashMap;
+use std::iter::FromIterator;
+
+use crate::field::extension_field::Extendable;
+use crate::field::field::Field;
+use crate::gates::gate::GateRef;
+
+#[derive(Debug, Clone)]
+enum Node<T> {
+    Terminus(T),
+    Bifurcation,
+}
+
+#[derive(Debug, Clone)]
+pub struct Tree<T> {
+    node: Node<T>,
+    left: Option<Box<Tree<T>>>,
+    right: Option<Box<Tree<T>>>,
+}
+
+impl<T> Default for Tree<T> {
+    fn default() -> Self {
+        Self {
+            node: Node::Bifurcation,
+            left: None,
+            right: None,
+        }
+    }
+}
+
+impl<T: Clone> Tree<T> {
+    pub fn preorder_traversal(&self) -> Vec<(T, Vec<bool>)> {
+        let mut res = Vec::new();
+        let prefix = [];
+        self.traverse(&prefix, &mut res);
+        res
+    }
+
+    fn traverse(&self, prefix: &[bool], current: &mut Vec<(T, Vec<bool>)>) {
+        if let Node::Terminus(t) = &self.node {
+            current.push((t.clone(), prefix.to_vec()));
+        } else {
+            if let Some(l) = &self.left {
+                let mut left_prefix = prefix.to_vec();
+                left_prefix.push(false);
+                l.traverse(&left_prefix, current);
+            }
+            if let Some(r) = &self.right {
+                let mut right_prefix = prefix.to_vec();
+                right_prefix.push(true);
+                r.traverse(&right_prefix, current);
+            }
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct GatePrefixes<F: Extendable<D>, const D: usize> {
+    pub prefixes: HashMap<GateRef<F, D>, Vec<bool>>,
+}
+
+impl<F: Extendable<D>, const D: usize> From<Tree<GateRef<F, D>>> for GatePrefixes<F, D> {
+    fn from(tree: Tree<GateRef<F, D>>) -> Self {
+        GatePrefixes {
+            prefixes: HashMap::from_iter(tree.preorder_traversal()),
+        }
+    }
+}
+
+impl<F: Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
+    pub fn from_gates(mut gates: Vec<GateRef<F, D>>) -> Self {
+        let timer = std::time::Instant::now();
+        gates.sort_unstable_by_key(|g| -((g.0.degree() + g.0.num_constants()) as isize));
+
+        for max_degree in 1..100 {
+            if let Some(mut tree) = Self::find_tree(&gates, max_degree) {
+                tree.prune();
+                println!(
+                    "Found tree with max degree {} in {}s.",
+                    max_degree,
+                    timer.elapsed().as_secs_f32()
+                );
+                return tree;
+            }
+        }
+
+        panic!("Can't find a tree.")
+    }
+
+    fn find_tree(gates: &[GateRef<F, D>], max_degree: usize) -> Option<Self> {
+        let mut tree = Tree::default();
+
+        for g in gates {
+            tree.try_add_gate(g, max_degree)?;
+        }
+        Some(tree)
+    }
+
+    fn try_add_gate(&mut self, g: &GateRef<F, D>, max_degree: usize) -> Option<()> {
+        let depth = max_degree.checked_sub(g.0.num_constants() + g.0.degree())?;
+        self.try_add_gate_at_depth(g, depth).is_err().then(|| ())
+    }
+
+    fn try_add_gate_at_depth(&mut self, g: &GateRef<F, D>, depth: usize) -> Result<(), GateAdded> {
+        if depth == 0 {
+            return if let Node::Bifurcation = self.node {
+                self.node = Node::Terminus(g.clone());
+                Err(GateAdded)
+            } else {
+                Ok(())
+            };
+        }
+
+        if let Node::Terminus(_) = self.node {
+            return Ok(());
+        }
+
+        if let Some(left) = &mut self.left {
+            left.try_add_gate_at_depth(g, depth - 1)?;
+        } else {
+            let mut left = Tree::default();
+            if left.try_add_gate_at_depth(g, depth - 1).is_err() {
+                self.left = Some(Box::new(left));
+                return Err(GateAdded);
+            }
+        }
+        if let Some(right) = &mut self.right {
+            right.try_add_gate_at_depth(g, depth - 1)?;
+        } else {
+            let mut right = Tree::default();
+            if right.try_add_gate_at_depth(g, depth - 1).is_err() {
+                self.right = Some(Box::new(right));
+                return Err(GateAdded);
+            }
+        }
+
+        Ok(())
+    }
+
+    fn prune(&mut self) {
+        if let (Some(left), None) = (&self.left, &self.right) {
+            debug_assert!(matches!(self.node, Node::Bifurcation));
+            let mut new = *left.clone();
+            new.prune();
+            *self = new;
+        }
+        if let Some(left) = &mut self.left {
+            left.prune();
+        }
+        if let Some(right) = &mut self.right {
+            right.prune();
+        }
+    }
+}
+
+struct GateAdded;

src/gates/mod.rs

@@ -2,6 +2,7 @@ pub(crate) mod arithmetic;
 pub mod base_sum;
 pub mod constant;
 pub(crate) mod gate;
+pub mod gate_tree;
 pub mod gmimc;
 pub mod interpolation;
 pub mod mul_extension;