From 383812dffdf9d9499cfafab8879946a418019bd3 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Fri, 26 Feb 2021 14:08:27 -0800
Subject: [PATCH] Degree map

---
 src/constraint_polynomial.rs | 43 ++++++++++++++++++++++++++++++++----
 src/gates/output_graph.rs    | 12 ++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/src/constraint_polynomial.rs b/src/constraint_polynomial.rs
index adcedb30..40e3f2f6 100644
--- a/src/constraint_polynomial.rs
+++ b/src/constraint_polynomial.rs
@@ -26,7 +26,7 @@ pub(crate) struct EvaluationVars<'a, F: Field> {
 /// than content. This is useful when we want to use constraint polynomials as `HashMap` keys, but
 /// we want address-based hashing for performance reasons.
 #[derive(Clone)]
-pub struct ConstraintPolynomial<F: Field>(Rc<ConstraintPolynomialInner<F>>);
+pub struct ConstraintPolynomial<F: Field>(pub(crate) Rc<ConstraintPolynomialInner<F>>);
 
 impl<F: Field> ConstraintPolynomial<F> {
     pub fn constant(c: F) -> Self {
@@ -112,13 +112,48 @@ impl<F: Field> ConstraintPolynomial<F> {
     }
 
     pub fn cube(&self) -> Self {
-        self * self * self
+        self.exp(3)
     }
 
-    pub(crate) fn degree(&self) -> BigUint {
+    pub fn degree(&self) -> BigUint {
         self.0.degree()
     }
 
+    pub(crate) fn populate_degree_map(&self, degrees: &mut HashMap<Self, BigUint>) {
+        if degrees.contains_key(self) {
+            // Already visited this node in the polynomial graph.
+            return;
+        }
+
+        match self.0.as_ref() {
+            ConstraintPolynomialInner::Constant(_) =>
+                degrees.insert(self.clone(), BigUint::zero()),
+            ConstraintPolynomialInner::LocalConstant(_) =>
+                degrees.insert(self.clone(), BigUint::one()),
+            ConstraintPolynomialInner::NextConstant(_) =>
+                degrees.insert(self.clone(), BigUint::one()),
+            ConstraintPolynomialInner::LocalWireValue(_) =>
+                degrees.insert(self.clone(), BigUint::one()),
+            ConstraintPolynomialInner::NextWireValue(_) =>
+                degrees.insert(self.clone(), BigUint::one()),
+            ConstraintPolynomialInner::Sum { lhs, rhs } => {
+                lhs.populate_degree_map(degrees);
+                rhs.populate_degree_map(degrees);
+                degrees.insert(self.clone(), (&degrees[lhs]).max(&degrees[rhs]).clone())
+            }
+            ConstraintPolynomialInner::Product { lhs, rhs } => {
+                lhs.populate_degree_map(degrees);
+                rhs.populate_degree_map(degrees);
+                degrees.insert(self.clone(), &degrees[lhs] + &degrees[rhs])
+            }
+            ConstraintPolynomialInner::Exponentiation { base, exponent } => {
+                base.populate_degree_map(degrees);
+                degrees.insert(self.clone(),
+                               &degrees[base] * BigUint::from_usize(*exponent).unwrap())
+            }
+        };
+    }
+
     /// Returns the set of wires that this constraint would depend on if it were applied at a
     /// certain gate index.
     pub(crate) fn dependencies(&self, gate: usize) -> Vec<Wire> {
@@ -365,7 +400,7 @@ impl<F: Field> Product for ConstraintPolynomial<F> {
     }
 }
 
-enum ConstraintPolynomialInner<F: Field> {
+pub(crate) enum ConstraintPolynomialInner<F: Field> {
     Constant(F),
 
     LocalConstant(usize),
diff --git a/src/gates/output_graph.rs b/src/gates/output_graph.rs
index 4bbee83f..c4f7d30a 100644
--- a/src/gates/output_graph.rs
+++ b/src/gates/output_graph.rs
@@ -2,6 +2,8 @@ use std::iter;
 
 use crate::constraint_polynomial::{ConstraintPolynomial};
 use crate::field::field::Field;
+use std::collections::HashMap;
+use num::BigUint;
 
 /// Represents a set of deterministic gate outputs, expressed as polynomials over witness
 /// values.
@@ -29,10 +31,20 @@ impl<F: Field> OutputGraph<F> {
     ///
     /// Note that this uses a simple greedy algorithm, so the result may not be optimal in terms of wire
     /// count.
+    // TODO: This doesn't yet work with large exponentiations, i.e. x^n where n > new_degree. Not an
+    // immediate problem since our gates don't use those.
     pub fn shrink_degree(&self, new_degree: usize) -> Self {
         todo!()
     }
 
+    fn degree_map(&self) -> HashMap<ConstraintPolynomial<F>, BigUint> {
+        let mut degrees = HashMap::new();
+        for (_loc, out) in &self.outputs {
+            out.populate_degree_map(&mut degrees);
+        }
+        degrees
+    }
+
     /// Allocate a new wire for the given target polynomial, and return a new output graph with
     /// references to the target polynomial replaced with references to that wire.
     fn allocate_wire(&self, target: ConstraintPolynomial<F>) -> Self {