Added comments

src/circuit_builder.rs

@@ -11,8 +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::{GatePrefixes, Tree};
+use crate::gates::gate_tree::Tree;
 use crate::gates::noop::NoopGate;
 use crate::generator::{CopyGenerator, WitnessGenerator};
 use crate::hash::hash_n_to_hash;

src/field/extension_field/mod.rs

@@ -1,6 +1,7 @@
-use crate::field::field::Field;
 use std::convert::TryInto;
 
+use crate::field::field::Field;
+
 pub mod algebra;
 pub mod quadratic;
 pub mod quartic;

src/gates/gate.rs

@@ -1,10 +1,13 @@
+use std::collections::HashMap;
 use std::fmt::{Debug, Error, Formatter};
 use std::hash::{Hash, Hasher};
+use std::iter::FromIterator;
 use std::sync::Arc;
 
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::{Extendable, FieldExtension};
+use crate::gates::gate_tree::Tree;
 use crate::generator::WitnessGenerator;
 use crate::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
 
@@ -125,3 +128,17 @@ pub struct GateInstance<F: Extendable<D>, const D: usize> {
     pub gate_type: GateRef<F, D>,
     pub constants: Vec<F>,
 }
+
+/// Map each gate to a boolean prefix used to construct the gate's selector polynomial.
+#[derive(Debug, 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.traversal()),
+        }
+    }
+}

src/gates/gate_tree.rs

@@ -1,19 +1,10 @@
-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)]
+/// A binary tree where leaves hold some type `T` and other nodes are empty.
-enum Node<T> {
-    Terminus(T),
-    Bifurcation,
-}
-
 #[derive(Debug, Clone)]
 pub enum Tree<T> {
-    Node(T),
+    Leaf(T),
     Bifurcation(Option<Box<Tree<T>>>, Option<Box<Tree<T>>>),
 }
 
@@ -24,18 +15,23 @@ impl<T> Default for Tree<T> {
 }
 
 impl<T: Clone> Tree<T> {
-    pub fn preorder_traversal(&self) -> Vec<(T, Vec<bool>)> {
+    /// Traverse a tree using a depth-first traversal and collect data and position for each leaf.
+    /// A leaf's position is represented by its left/right path, where `false` means left and `true` means right.
+    pub fn traversal(&self) -> Vec<(T, Vec<bool>)> {
         let mut res = Vec::new();
         let prefix = [];
         self.traverse(&prefix, &mut res);
         res
     }
 
+    /// Utility function to traverse the tree.
     fn traverse(&self, prefix: &[bool], current: &mut Vec<(T, Vec<bool>)>) {
         match &self {
-            Tree::Node(t) => {
+            // If node is a leaf, collect the data and position.
+            Tree::Leaf(t) => {
                 current.push((t.clone(), prefix.to_vec()));
             }
+            // Otherwise, traverse the left subtree and then the right subtree.
             Tree::Bifurcation(left, right) => {
                 if let Some(l) = left {
                     let mut left_prefix = prefix.to_vec();
@@ -52,27 +48,22 @@ impl<T: Clone> Tree<T> {
     }
 }
 
-#[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>> {
+    /// Construct a binary tree of gates using the following greedy algorithm:
+    /// We want a tree where the maximum `M` of
+    /// `F(gate) = gate.degree() + gate.num_constants() + tree.depth(gate)`
+    /// over all gates is minimized. Such a tree is constructed by iterating over possible values of `M`
+    /// (from 1 to 99, then we give up) and then looking for a tree with this value of `M`
+    /// using `Self::find_tree`. This latter function greedily adds gates at the depth where
+    /// `F(gate)=M` to ensure no space is wasted. We return the first tree found in this manner,
+    /// i.e., the one with minimal `M` value.
     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();
+                tree.shorten();
                 println!(
                     "Found tree with max degree {} in {}s.",
                     max_degree,
@@ -85,6 +76,7 @@ impl<F: Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
         panic!("Can't find a tree.")
     }
 
+    /// Greedily add gates wherever possible. Returns `None` if this fails.
     fn find_tree(gates: &[GateRef<F, D>], max_degree: usize) -> Option<Self> {
         let mut tree = Tree::default();
 
@@ -94,31 +86,39 @@ impl<F: Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
         Some(tree)
     }
 
+    /// Try to add a gate in the tree. Returns `None` if this fails.
     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)
     }
 
+    /// Try to add a gate in the tree at a specified depth. Returns `None` if this fails.
     fn try_add_gate_at_depth(&mut self, g: &GateRef<F, D>, depth: usize) -> Option<()> {
+        // If depth is 0, we have to insert the gate here.
         if depth == 0 {
             return if let Tree::Bifurcation(_, _) = self {
-                *self = Tree::Node(g.clone());
+                // Insert the gate as a new leaf.
+                *self = Tree::Leaf(g.clone());
                 Some(())
             } else {
+                // A leaf is already here.
                 None
             };
         }
 
-        if let Tree::Node(_) = self {
+        // A leaf is already here so we cannot go deeper.
+        if let Tree::Leaf(_) = self {
             return None;
         }
 
         if let Tree::Bifurcation(left, right) = self {
             if let Some(left) = left {
+                // Try to add the gate to the left if there's already a left subtree.
                 if left.try_add_gate_at_depth(g, depth - 1).is_some() {
                     return Some(());
                 }
             } else {
+                // Add a new left subtree and try to add the gate to it.
                 let mut new_left = Tree::default();
                 if new_left.try_add_gate_at_depth(g, depth - 1).is_some() {
                     *left = Some(Box::new(new_left));
@@ -126,10 +126,12 @@ impl<F: Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
                 }
             }
             if let Some(right) = right {
+                // Try to add the gate to the right if there's already a right subtree.
                 if right.try_add_gate_at_depth(g, depth - 1).is_some() {
                     return Some(());
                 }
             } else {
+                // Add a new right subtree and try to add the gate to it.
                 let mut new_right = Tree::default();
                 if new_right.try_add_gate_at_depth(g, depth - 1).is_some() {
                     *right = Some(Box::new(new_right));
@@ -141,23 +143,27 @@ impl<F: Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
         None
     }
 
-    fn prune(&mut self) {
+    /// `Self::find_tree` returns a tree where each gate has `F(gate)=M` (see `Self::from_gates` comment).
+    /// This can produce subtrees with more nodes than necessary. This function removes useless nodes,
+    /// i.e., nodes that have a left but no right subtree.
+    fn shorten(&mut self) {
         if let Tree::Bifurcation(left, right) = self {
             if let (Some(left), None) = (left, right) {
+                // If the node has a left but no right subtree, set the node to its (shortened) left subtree.
                 let mut new = *left.clone();
-                new.prune();
+                new.shorten();
                 *self = new;
             }
         }
         if let Tree::Bifurcation(left, right) = self {
             if let Some(left) = left {
-                left.prune();
+                // Shorten the left subtree if there is one.
+                left.shorten();
             }
             if let Some(right) = right {
-                right.prune();
+                // Shorten the right subtree if there is one.
+                right.shorten();
             }
         }
     }
 }
-
-struct GateAdded;

src/plonk_challenger.rs

@@ -1,3 +1,5 @@
+use std::convert::TryInto;
+
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::{Extendable, FieldExtension};
@@ -5,7 +7,6 @@ use crate::field::field::Field;
 use crate::hash::{permute, SPONGE_RATE, SPONGE_WIDTH};
 use crate::proof::{Hash, HashTarget, OpeningSet};
 use crate::target::Target;
-use std::convert::TryInto;
 
 /// Observes prover messages, and generates challenges by hashing the transcript.
 #[derive(Clone)]

src/target.rs

@@ -1,6 +1,7 @@
+use std::ops::Range;
+
 use crate::circuit_data::CircuitConfig;
 use crate::wire::Wire;
-use std::ops::Range;
 
 /// A location in the witness.
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]

src/wire.rs

@@ -1,6 +1,7 @@
-use crate::circuit_data::CircuitConfig;
 use std::ops::Range;
 
+use crate::circuit_data::CircuitConfig;
+
 /// Represents a wire in the circuit.
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct Wire {

src/witness.rs

@@ -1,11 +1,11 @@
 use std::collections::HashMap;
+use std::convert::TryInto;
 
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::{Extendable, FieldExtension};
 use crate::field::field::Field;
 use crate::target::Target;
 use crate::wire::Wire;
-use std::convert::TryInto;
 
 #[derive(Clone, Debug)]
 pub struct PartialWitness<F: Field> {