Tree of scopes (#106)

* Tree of scopes

This is an extension of the context concept.

Earlier I was planning to store a simple stack of contexts, but I ended up storing the whole history, in a tree structure. This gives us more control over the output, i.e. we can print the gate count of a parent scope before those of its child scopes, which seems more user-friendly.

Sample gate count output:

    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] 27829 gates to root
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | 2373 gates to evaluate the vanishing polynomial at our challenge point, zeta.
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | | 1284 gates to evaluate gate constraints
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | 25312 gates to verify FRI proof
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | | 650 gates to verify 0'th FRI query
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | | | 96 gates to check FRI initial proof
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | | | 65 gates to compute x from its index
    [2021-07-19T18:09:24Z INFO  plonky2::circuit_builder] | | | 233 gates to combine initial oracles
    ...

Sample copy constraint failure:

    Error: Copy constraint 'root > verify FRI proof > verify 0'th FRI query > check FRI initial proof > verify 0'th initial Merkle proof > check Merkle root: 0-th hash element' between wire 12 of gate #2550 [...] and wire 0 of gate #0 [...] is not satisfied. Got values of 6861386743364621393 and 0 respectively.

* No min

* info -> debug

* Move to its own file

src/circuit_builder.rs

@@ -8,6 +8,7 @@ use crate::circuit_data::{
     CircuitConfig, CircuitData, CommonCircuitData, ProverCircuitData, ProverOnlyCircuitData,
     VerifierCircuitData, VerifierOnlyCircuitData,
 };
+use crate::context_tree::ContextTree;
 use crate::copy_constraint::CopyConstraint;
 use crate::field::cosets::get_unique_coset_shifts;
 use crate::field::extension_field::target::ExtensionTarget;
@@ -46,8 +47,8 @@ pub struct CircuitBuilder<F: Extendable<D>, const D: usize> {
 
     copy_constraints: Vec<CopyConstraint>,
 
-    /// A string used to give context to copy constraints.
-    context: String,
+    /// A tree of named scopes, used for debugging.
+    context_log: ContextTree,
 
     /// A vector of marked targets. The values assigned to these targets will be displayed by the prover.
     marked_targets: Vec<MarkedTargets<D>>,
@@ -68,7 +69,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             public_input_index: 0,
             virtual_target_index: 0,
             copy_constraints: Vec::new(),
-            context: String::new(),
+            context_log: ContextTree::new(),
             marked_targets: Vec::new(),
             generators: Vec::new(),
             constants_to_targets: HashMap::new(),
@@ -208,7 +209,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             "Tried to route a wire that isn't routable"
         );
         self.copy_constraints
-            .push(CopyConstraint::new((x, y), self.context.clone()));
+            .push(CopyConstraint::new((x, y), self.context_log.open_stack()));
     }
 
     /// Same as `assert_equal` for a named copy constraint.
@@ -223,7 +224,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         );
         self.copy_constraints.push(CopyConstraint::new(
             (x, y),
-            format!("{}: {}", self.context.clone(), name),
+            format!("{} > {}", self.context_log.open_stack(), name),
         ));
     }
 
@@ -305,8 +306,12 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         self.targets_to_constants.get(&target).cloned()
     }
 
-    pub fn set_context(&mut self, new_context: &str) {
-        self.context = new_context.to_string();
+    pub fn push_context(&mut self, ctx: &str) {
+        self.context_log.push(ctx, self.num_gates());
+    }
+
+    pub fn pop_context(&mut self) {
+        self.context_log.pop(self.num_gates());
     }
 
     pub fn add_marked(&mut self, targets: Markable<D>, name: &str) {
@@ -485,6 +490,12 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         wire_partition.get_sigma_polys(degree_log, k_is, subgroup)
     }
 
+    pub fn print_gate_counts(&self, min_delta: usize) {
+        self.context_log
+            .filter(self.num_gates(), min_delta)
+            .print(self.num_gates());
+    }
+
     /// Builds a "full circuit", with both prover and verifier data.
     pub fn build(mut self) -> CircuitData<F, D> {
         let quotient_degree_factor = 7; // TODO: add this as a parameter.

src/context_tree.rs

@@ -0,0 +1,124 @@
+use log::debug;
+
+/// The hierarchy of contexts, and the gate count contributed by each one. Useful for debugging.
+pub(crate) struct ContextTree {
+    /// The name of this scope.
+    name: String,
+    /// The gate count when this scope was created.
+    enter_gate_count: usize,
+    /// The gate count when this scope was destroyed, or None if it has not yet been destroyed.
+    exit_gate_count: Option<usize>,
+    /// Any child contexts.
+    children: Vec<ContextTree>,
+}
+
+impl ContextTree {
+    pub fn new() -> Self {
+        Self {
+            name: "root".to_string(),
+            enter_gate_count: 0,
+            exit_gate_count: None,
+            children: vec![],
+        }
+    }
+
+    /// Whether this context is still in scope.
+    fn is_open(&self) -> bool {
+        self.exit_gate_count.is_none()
+    }
+
+    /// A description of the stack of currently-open scopes.
+    pub fn open_stack(&self) -> String {
+        let mut stack = Vec::new();
+        self.open_stack_helper(&mut stack);
+        stack.join(" > ")
+    }
+
+    fn open_stack_helper(&self, stack: &mut Vec<String>) {
+        if self.is_open() {
+            stack.push(self.name.clone());
+            if let Some(last_child) = self.children.last() {
+                last_child.open_stack_helper(stack);
+            }
+        }
+    }
+
+    pub fn push(&mut self, ctx: &str, current_gate_count: usize) {
+        assert!(self.is_open());
+
+        if let Some(last_child) = self.children.last_mut() {
+            if last_child.is_open() {
+                last_child.push(ctx, current_gate_count);
+                return;
+            }
+        }
+
+        self.children.push(ContextTree {
+            name: ctx.to_string(),
+            enter_gate_count: current_gate_count,
+            exit_gate_count: None,
+            children: vec![],
+        })
+    }
+
+    /// Close the deepest open context from this tree.
+    pub fn pop(&mut self, current_gate_count: usize) {
+        assert!(self.is_open());
+
+        if let Some(last_child) = self.children.last_mut() {
+            if last_child.is_open() {
+                last_child.pop(current_gate_count);
+                return;
+            }
+        }
+
+        self.exit_gate_count = Some(current_gate_count);
+    }
+
+    fn gate_count_delta(&self, current_gate_count: usize) -> usize {
+        self.exit_gate_count.unwrap_or(current_gate_count) - self.enter_gate_count
+    }
+
+    /// Filter out children with a low gate count.
+    pub fn filter(&self, current_gate_count: usize, min_delta: usize) -> Self {
+        Self {
+            name: self.name.clone(),
+            enter_gate_count: self.enter_gate_count,
+            exit_gate_count: self.exit_gate_count,
+            children: self
+                .children
+                .iter()
+                .filter(|c| c.gate_count_delta(current_gate_count) >= min_delta)
+                .map(|c| c.filter(current_gate_count, min_delta))
+                .collect(),
+        }
+    }
+
+    pub fn print(&self, current_gate_count: usize) {
+        self.print_helper(current_gate_count, 0);
+    }
+
+    fn print_helper(&self, current_gate_count: usize, depth: usize) {
+        let prefix = "| ".repeat(depth);
+        debug!(
+            "{}{} gates to {}",
+            prefix,
+            self.gate_count_delta(current_gate_count),
+            self.name
+        );
+        for child in &self.children {
+            child.print_helper(current_gate_count, depth + 1);
+        }
+    }
+}
+
+/// Creates a named scope; useful for debugging.
+#[macro_export]
+macro_rules! context {
+    ($builder:expr, $ctx:expr, $exp:expr) => {{
+        $builder.push_context($ctx);
+        let res = $exp;
+        $builder.pop_context();
+        res
+    }};
+}

src/fri/recursive_verifier.rs

@@ -1,5 +1,6 @@
 use crate::circuit_builder::CircuitBuilder;
 use crate::circuit_data::CommonCircuitData;
+use crate::context;
 use crate::field::extension_field::target::{flatten_target, ExtensionTarget};
 use crate::field::extension_field::Extendable;
 use crate::field::field::Field;
@@ -86,19 +87,25 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         // Size of the LDE domain.
         let n = proof.final_poly.len() << (total_arities + config.rate_bits);
 
-        self.set_context("Recover the random betas used in the FRI reductions.");
-        let betas = proof
-            .commit_phase_merkle_roots
-            .iter()
-            .map(|root| {
-                challenger.observe_hash(root);
-                challenger.get_extension_challenge(self)
-            })
-            .collect::<Vec<_>>();
+        let betas = context!(
+            self,
+            "recover the random betas used in the FRI reductions.",
+            proof
+                .commit_phase_merkle_roots
+                .iter()
+                .map(|root| {
+                    challenger.observe_hash(root);
+                    challenger.get_extension_challenge(self)
+                })
+                .collect::<Vec<_>>()
+        );
         challenger.observe_extension_elements(&proof.final_poly.0);
 
-        self.set_context("Check PoW");
-        self.fri_verify_proof_of_work(proof, challenger, &config.fri_config);
+        context!(
+            self,
+            "check PoW",
+            self.fri_verify_proof_of_work(proof, challenger, &config.fri_config)
+        );
 
         // Check that parameters are coherent.
         debug_assert_eq!(
@@ -111,18 +118,22 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             "Number of reductions should be non-zero."
         );
 
-        for round_proof in &proof.query_round_proofs {
-            self.fri_verifier_query_round(
-                os,
-                zeta,
-                alpha,
-                initial_merkle_roots,
-                proof,
-                challenger,
-                n,
-                &betas,
-                round_proof,
-                common_data,
+        for (i, round_proof) in proof.query_round_proofs.iter().enumerate() {
+            context!(
+                self,
+                &format!("verify {}'th FRI query", i),
+                self.fri_verifier_query_round(
+                    os,
+                    zeta,
+                    alpha,
+                    initial_merkle_roots,
+                    proof,
+                    challenger,
+                    n,
+                    &betas,
+                    round_proof,
+                    common_data,
+                )
             );
         }
     }
@@ -139,8 +150,11 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             .zip(initial_merkle_roots)
             .enumerate()
         {
-            self.set_context(&format!("Verify {}-th initial Merkle proof.", i));
-            self.verify_merkle_proof(evals.clone(), x_index, root, merkle_proof);
+            context!(
+                self,
+                &format!("verify {}'th initial Merkle proof", i),
+                self.verify_merkle_proof(evals.clone(), x_index, root, merkle_proof)
+            );
         }
     }
 
@@ -245,41 +259,55 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         x_index = self.split_low_high(x_index, n_log, 64).0;
         let mut x_index_num_bits = n_log;
         let mut domain_size = n;
-        self.set_context("Check FRI initial proof.");
-        self.fri_verify_initial_proof(
-            x_index,
-            &round_proof.initial_trees_proof,
-            initial_merkle_roots,
+        context!(
+            self,
+            "check FRI initial proof",
+            self.fri_verify_initial_proof(
+                x_index,
+                &round_proof.initial_trees_proof,
+                initial_merkle_roots,
+            )
         );
         let mut old_x_index = self.zero();
-        // `subgroup_x` is `subgroup[x_index]`, i.e., the actual field element in the domain.
-        let g = self.constant(F::MULTIPLICATIVE_GROUP_GENERATOR);
-        let phi = self.constant(F::primitive_root_of_unity(n_log));
 
-        let reversed_x = self.reverse_limbs::<2>(x_index, n_log);
-        let phi = self.exp(phi, reversed_x, n_log);
-        let mut subgroup_x = self.mul(g, phi);
+        // `subgroup_x` is `subgroup[x_index]`, i.e., the actual field element in the domain.
+        let mut subgroup_x = context!(self, "compute x from its index", {
+            let g = self.constant(F::MULTIPLICATIVE_GROUP_GENERATOR);
+            let phi = self.constant(F::primitive_root_of_unity(n_log));
+
+            let reversed_x = self.reverse_limbs::<2>(x_index, n_log);
+            let phi = self.exp(phi, reversed_x, n_log);
+            self.mul(g, phi)
+        });
 
         for (i, &arity_bits) in config.reduction_arity_bits.iter().enumerate() {
             let next_domain_size = domain_size >> arity_bits;
             let e_x = if i == 0 {
-                self.fri_combine_initial(
-                    &round_proof.initial_trees_proof,
-                    alpha,
-                    os,
-                    zeta,
-                    subgroup_x,
-                    common_data,
+                context!(
+                    self,
+                    "combine initial oracles",
+                    self.fri_combine_initial(
+                        &round_proof.initial_trees_proof,
+                        alpha,
+                        os,
+                        zeta,
+                        subgroup_x,
+                        common_data,
+                    )
                 )
             } else {
                 let last_evals = &evaluations[i - 1];
                 // Infer P(y) from {P(x)}_{x^arity=y}.
-                self.compute_evaluation(
-                    subgroup_x,
-                    old_x_index,
-                    config.reduction_arity_bits[i - 1],
-                    last_evals,
-                    betas[i - 1],
+                context!(
+                    self,
+                    "infer evaluation using interpolation",
+                    self.compute_evaluation(
+                        subgroup_x,
+                        old_x_index,
+                        config.reduction_arity_bits[i - 1],
+                        last_evals,
+                        betas[i - 1],
+                    )
                 )
             };
             let mut evals = round_proof.steps[i].evals.clone();
@@ -288,12 +316,15 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
                 self.split_low_high(x_index, arity_bits, x_index_num_bits);
             evals = self.insert(low_x_index, e_x, evals);
             evaluations.push(evals);
-            self.set_context("Verify FRI round Merkle proof.");
-            self.verify_merkle_proof(
-                flatten_target(&evaluations[i]),
-                high_x_index,
-                proof.commit_phase_merkle_roots[i],
-                &round_proof.steps[i].merkle_proof,
+            context!(
+                self,
+                "verify FRI round Merkle proof.",
+                self.verify_merkle_proof(
+                    flatten_target(&evaluations[i]),
+                    high_x_index,
+                    proof.commit_phase_merkle_roots[i],
+                    &round_proof.steps[i].merkle_proof,
+                )
             );
 
             if i > 0 {
@@ -310,12 +341,16 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
         let last_evals = evaluations.last().unwrap();
         let final_arity_bits = *config.reduction_arity_bits.last().unwrap();
-        let purported_eval = self.compute_evaluation(
-            subgroup_x,
-            old_x_index,
-            final_arity_bits,
-            last_evals,
-            *betas.last().unwrap(),
+        let purported_eval = context!(
+            self,
+            "infer final evaluation using interpolation",
+            self.compute_evaluation(
+                subgroup_x,
+                old_x_index,
+                final_arity_bits,
+                last_evals,
+                *betas.last().unwrap(),
+            )
         );
         for _ in 0..final_arity_bits {
             subgroup_x = self.square(subgroup_x);
@@ -323,7 +358,11 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
         // Final check of FRI. After all the reductions, we check that the final polynomial is equal
         // to the one sent by the prover.
-        let eval = proof.final_poly.eval_scalar(self, subgroup_x);
+        let eval = context!(
+            self,
+            "evaluate final polynomial",
+            proof.final_poly.eval_scalar(self, subgroup_x)
+        );
         self.assert_equal_extension(eval, purported_eval);
     }
 }

src/lib.rs

@@ -2,6 +2,7 @@
 
 pub mod circuit_builder;
 pub mod circuit_data;
+pub mod context_tree;
 pub mod copy_constraint;
 pub mod field;
 pub mod fri;

src/merkle_proofs.rs

@@ -130,7 +130,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let leaf_index_rev = self.reverse_limbs::<2>(leaf_index, height);
         self.assert_equal(acc_leaf_index, leaf_index_rev);
 
-        self.named_assert_hashes_equal(state, merkle_root, "Check Merkle root".into())
+        self.named_assert_hashes_equal(state, merkle_root, "check Merkle root".into())
     }
 
     pub(crate) fn assert_hashes_equal(&mut self, x: HashTarget, y: HashTarget) {

src/polynomial/commitment.rs

@@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::circuit_builder::CircuitBuilder;
 use crate::circuit_data::CommonCircuitData;
+use crate::context;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
 use crate::field::field::Field;
@@ -283,15 +284,19 @@ impl<const D: usize> OpeningProofTarget<D> {
 
         let alpha = challenger.get_extension_challenge(builder);
 
-        builder.verify_fri_proof(
-            log2_ceil(common_data.degree()),
-            &os,
-            zeta,
-            alpha,
-            merkle_roots,
-            &self.fri_proof,
-            challenger,
-            common_data,
+        context!(
+            builder,
+            "verify FRI proof",
+            builder.verify_fri_proof(
+                log2_ceil(common_data.degree()),
+                &os,
+                zeta,
+                alpha,
+                merkle_roots,
+                &self.fri_proof,
+                challenger,
+                common_data,
+            )
         );
     }
 }

src/recursive_verifier.rs

@@ -1,5 +1,6 @@
 use crate::circuit_builder::CircuitBuilder;
 use crate::circuit_data::{CircuitConfig, CommonCircuitData, VerifierCircuitTarget};
+use crate::context;
 use crate::field::extension_field::Extendable;
 use crate::plonk_challenger::RecursiveChallenger;
 use crate::proof::{HashTarget, ProofTarget};
@@ -27,20 +28,25 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
         let mut challenger = RecursiveChallenger::new(self);
 
-        self.set_context("Challenger observes proof and generates challenges.");
-        let digest =
-            HashTarget::from_vec(self.constants(&inner_common_data.circuit_digest.elements));
-        challenger.observe_hash(&digest);
+        let (betas, gammas, alphas, zeta) =
+            context!(self, "observe proof and generates challenges", {
+                let digest = HashTarget::from_vec(
+                    self.constants(&inner_common_data.circuit_digest.elements),
+                );
+                challenger.observe_hash(&digest);
 
-        challenger.observe_hash(&proof.wires_root);
-        let betas = challenger.get_n_challenges(self, num_challenges);
-        let gammas = challenger.get_n_challenges(self, num_challenges);
+                challenger.observe_hash(&proof.wires_root);
+                let betas = challenger.get_n_challenges(self, num_challenges);
+                let gammas = challenger.get_n_challenges(self, num_challenges);
 
-        challenger.observe_hash(&proof.plonk_zs_partial_products_root);
-        let alphas = challenger.get_n_challenges(self, num_challenges);
+                challenger.observe_hash(&proof.plonk_zs_partial_products_root);
+                let alphas = challenger.get_n_challenges(self, num_challenges);
 
-        challenger.observe_hash(&proof.quotient_polys_root);
-        let zeta = challenger.get_extension_challenge(self);
+                challenger.observe_hash(&proof.quotient_polys_root);
+                let zeta = challenger.get_extension_challenge(self);
+
+                (betas, gammas, alphas, zeta)
+            });
 
         let local_constants = &proof.openings.constants;
         let local_wires = &proof.openings.wires;
@@ -54,38 +60,42 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let partial_products = &proof.openings.partial_products;
 
         let zeta_pow_deg = self.exp_power_of_2(zeta, inner_common_data.degree_bits);
-        self.set_context("Evaluate the vanishing polynomial at our challenge point, zeta.");
-        let vanishing_polys_zeta = eval_vanishing_poly_recursively(
+        let vanishing_polys_zeta = context!(
             self,
-            inner_common_data,
-            zeta,
-            zeta_pow_deg,
-            vars,
-            local_zs,
-            next_zs,
-            partial_products,
-            s_sigmas,
-            &betas,
-            &gammas,
-            &alphas,
+            "evaluate the vanishing polynomial at our challenge point, zeta.",
+            eval_vanishing_poly_recursively(
+                self,
+                inner_common_data,
+                zeta,
+                zeta_pow_deg,
+                vars,
+                local_zs,
+                next_zs,
+                partial_products,
+                s_sigmas,
+                &betas,
+                &gammas,
+                &alphas,
+            )
         );
 
-        self.set_context("Check vanishing and quotient polynomials.");
-        let quotient_polys_zeta = &proof.openings.quotient_polys;
-        let mut scale = ReducingFactorTarget::new(zeta_pow_deg);
-        let z_h_zeta = self.sub_extension(zeta_pow_deg, one);
-        for (i, chunk) in quotient_polys_zeta
-            .chunks(inner_common_data.quotient_degree_factor)
-            .enumerate()
-        {
-            let recombined_quotient = scale.reduce(chunk, self);
-            let computed_vanishing_poly = self.mul_extension(z_h_zeta, recombined_quotient);
-            self.named_route_extension(
-                vanishing_polys_zeta[i],
-                computed_vanishing_poly,
-                format!("Vanishing polynomial == Z_H * quotient, challenge {}", i),
-            );
-        }
+        context!(self, "check vanishing and quotient polynomials.", {
+            let quotient_polys_zeta = &proof.openings.quotient_polys;
+            let mut scale = ReducingFactorTarget::new(zeta_pow_deg);
+            let z_h_zeta = self.sub_extension(zeta_pow_deg, one);
+            for (i, chunk) in quotient_polys_zeta
+                .chunks(inner_common_data.quotient_degree_factor)
+                .enumerate()
+            {
+                let recombined_quotient = scale.reduce(chunk, self);
+                let computed_vanishing_poly = self.mul_extension(z_h_zeta, recombined_quotient);
+                self.named_route_extension(
+                    vanishing_polys_zeta[i],
+                    computed_vanishing_poly,
+                    format!("Vanishing polynomial == Z_H * quotient, challenge {}", i),
+                );
+            }
+        });
 
         let merkle_roots = &[
             inner_verifier_data.constants_sigmas_root,
@@ -361,6 +371,7 @@ mod tests {
 
         builder.add_recursive_verifier(pt, &config, &inner_data, &cd);
 
+        builder.print_gate_counts(0);
         let data = builder.build();
         let recursive_proof = data.prove(pw)?;
 

src/vanishing_poly.rs

@@ -1,5 +1,6 @@
 use crate::circuit_builder::CircuitBuilder;
 use crate::circuit_data::CommonCircuitData;
+use crate::context;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
 use crate::field::field::Field;
@@ -266,11 +267,15 @@ pub(crate) fn eval_vanishing_poly_recursively<F: Extendable<D>, const D: usize>(
     let max_degree = common_data.quotient_degree_factor;
     let (num_prods, final_num_prod) = common_data.num_partial_products;
 
-    let constraint_terms = evaluate_gate_constraints_recursively(
+    let constraint_terms = context!(
         builder,
-        &common_data.gates,
-        common_data.num_gate_constraints,
-        vars,
+        "evaluate gate constraints",
+        evaluate_gate_constraints_recursively(
+            builder,
+            &common_data.gates,
+            common_data.num_gate_constraints,
+            vars,
+        )
     );
 
     // The L_1(x) (Z(x) - 1) vanishing terms.