From da3d34a0d4ecb83fb3a430250710609c2ac25c90 Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Tue, 22 Jun 2021 14:31:46 +0200 Subject: [PATCH] Working gate tree generation --- src/circuit_builder.rs | 7 ++ src/gates/gate.rs | 7 ++ src/gates/gate_tree.rs | 156 +++++++++++++++++++++++++++++++++++++++++ src/gates/mod.rs | 1 + 4 files changed, 171 insertions(+) create mode 100644 src/gates/gate_tree.rs diff --git a/src/circuit_builder.rs b/src/circuit_builder.rs index 3a26497e..232c9c20 100644 --- a/src/circuit_builder.rs +++ b/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, const D: usize> CircuitBuilder { /// Builds a "full circuit", with both prover and verifier data. pub fn build(mut self) -> CircuitData { + 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: {}", diff --git a/src/gates/gate.rs b/src/gates/gate.rs index 1765191e..d27c8f8b 100644 --- a/src/gates/gate.rs +++ b/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, const D: usize> Hash for GateRef { impl, const D: usize> Eq for GateRef {} +impl, const D: usize> Debug for GateRef { + 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, const D: usize> { pub gate_type: GateRef, diff --git a/src/gates/gate_tree.rs b/src/gates/gate_tree.rs new file mode 100644 index 00000000..39b3b9c2 --- /dev/null +++ b/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 { + Terminus(T), + Bifurcation, +} + +#[derive(Debug, Clone)] +pub struct Tree { + node: Node, + left: Option>>, + right: Option>>, +} + +impl Default for Tree { + fn default() -> Self { + Self { + node: Node::Bifurcation, + left: None, + right: None, + } + } +} + +impl Tree { + pub fn preorder_traversal(&self) -> Vec<(T, Vec)> { + let mut res = Vec::new(); + let prefix = []; + self.traverse(&prefix, &mut res); + res + } + + fn traverse(&self, prefix: &[bool], current: &mut Vec<(T, Vec)>) { + 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, const D: usize> { + pub prefixes: HashMap, Vec>, +} + +impl, const D: usize> From>> for GatePrefixes { + fn from(tree: Tree>) -> Self { + GatePrefixes { + prefixes: HashMap::from_iter(tree.preorder_traversal()), + } + } +} + +impl, const D: usize> Tree> { + pub fn from_gates(mut gates: Vec>) -> 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], max_degree: usize) -> Option { + 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, 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, 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; diff --git a/src/gates/mod.rs b/src/gates/mod.rs index d013c7cd..bb8b178b 100644 --- a/src/gates/mod.rs +++ b/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;