feat: serialize SymbolicConstraints as DAG (#14)

jonathanpwang · web-flow · commit a97c6664ac29 · 2024-12-21T12:28:58.000-08:00
* chore: clarify comment

* feat: serialize `SymbolicConstraints` as DAG

Including interactions. Doing this for now to keep interfaces the same.

In the upcoming frontend-backend split, we will make a breaking
interface change to switch to using the DAG as the primary interface and
make `SymbolicExpression` and `SymbolicConstraints` private (or remove
entirely).

* chore: bump version to v0.1.2
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,5 +1,5 @@
 [workspace.package]
-version = "0.1.0-alpha"
+version = "0.1.2-alpha"
 edition = "2021"
 rust-version = "1.82"
 authors = ["OpenVM contributors"]
diff --git a/crates/stark-backend/src/air_builders/symbolic/dag.rs b/crates/stark-backend/src/air_builders/symbolic/dag.rs
@@ -4,8 +4,12 @@ use p3_field::Field;
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 
-use crate::air_builders::symbolic::{
-    symbolic_expression::SymbolicExpression, symbolic_variable::SymbolicVariable,
+use super::SymbolicConstraints;
+use crate::{
+    air_builders::symbolic::{
+        symbolic_expression::SymbolicExpression, symbolic_variable::SymbolicVariable,
+    },
+    interaction::{Interaction, SymbolicInteraction},
 };
 
 /// A node in symbolic expression DAG.
@@ -50,18 +54,62 @@ pub struct SymbolicExpressionDag<F> {
     pub(crate) constraint_idx: Vec<usize>,
 }
 
-pub(crate) fn build_symbolic_expr_dag<F: Field>(
-    exprs: &[SymbolicExpression<F>],
-) -> SymbolicExpressionDag<F> {
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(bound = "F: Field")]
+pub struct SymbolicConstraintsDag<F> {
+    /// DAG with all symbolic expressions as nodes.
+    /// A subset of the nodes represents all constraints that will be
+    /// included in the quotient polynomial via DEEP-ALI.
+    pub constraints: SymbolicExpressionDag<F>,
+    /// List of all interactions, where expressions in the interactions
+    /// are referenced by node idx as `usize`.
+    ///
+    /// This is used by the prover for after challenge trace generation,
+    /// and some partial information may be used by the verifier.
+    ///
+    /// **However**, any contributions to the quotient polynomial from
+    /// logup are already included in `constraints` and do not need to
+    /// be separately calculated from `interactions`.
+    pub interactions: Vec<Interaction<usize>>,
+}
+
+pub(crate) fn build_symbolic_constraints_dag<F: Field>(
+    constraints: &[SymbolicExpression<F>],
+    interactions: &[SymbolicInteraction<F>],
+) -> SymbolicConstraintsDag<F> {
     let mut expr_to_idx = FxHashMap::default();
     let mut nodes = Vec::new();
-    let constraint_idx = exprs
+    let constraint_idx = constraints
         .iter()
         .map(|expr| topological_sort_symbolic_expr(expr, &mut expr_to_idx, &mut nodes))
         .collect();
-    SymbolicExpressionDag {
+    let interactions: Vec<Interaction<usize>> = interactions
+        .iter()
+        .map(|interaction| {
+            let fields: Vec<usize> = interaction
+                .fields
+                .iter()
+                .map(|field_expr| {
+                    topological_sort_symbolic_expr(field_expr, &mut expr_to_idx, &mut nodes)
+                })
+                .collect();
+            let count =
+                topological_sort_symbolic_expr(&interaction.count, &mut expr_to_idx, &mut nodes);
+            Interaction {
+                fields,
+                count,
+                bus_index: interaction.bus_index,
+                interaction_type: interaction.interaction_type,
+            }
+        })
+        .collect();
+    let constraints = SymbolicExpressionDag {
         nodes,
         constraint_idx,
+    };
+    SymbolicConstraintsDag {
+        constraints,
+        interactions,
     }
 }
 
@@ -139,8 +187,9 @@ fn topological_sort_symbolic_expr<'a, F: Field>(
 }
 
 impl<F: Field> SymbolicExpressionDag<F> {
-    /// Returns symbolic expressions for each constraint
-    pub fn to_symbolic_expressions(&self) -> Vec<SymbolicExpression<F>> {
+    /// Convert each node to a [`SymbolicExpression<F>`] reference and return
+    /// the full list.
+    fn to_symbolic_expressions(&self) -> Vec<Arc<SymbolicExpression<F>>> {
         let mut exprs: Vec<Arc<SymbolicExpression<_>>> = Vec::with_capacity(self.nodes.len());
         for node in &self.nodes {
             let expr = match *node {
@@ -186,10 +235,48 @@ impl<F: Field> SymbolicExpressionDag<F> {
             };
             exprs.push(Arc::new(expr));
         }
-        self.constraint_idx
-            .iter()
-            .map(|&idx| exprs[idx].as_ref().clone())
-            .collect()
+        exprs
+    }
+}
+
+// TEMPORARY conversions until we switch main interfaces to use SymbolicConstraintsDag
+impl<F: Field> From<SymbolicConstraintsDag<F>> for SymbolicConstraints<F> {
+    fn from(dag: SymbolicConstraintsDag<F>) -> Self {
+        let exprs = dag.constraints.to_symbolic_expressions();
+        let constraints = dag
+            .constraints
+            .constraint_idx
+            .into_iter()
+            .map(|idx| exprs[idx].as_ref().clone())
+            .collect::<Vec<_>>();
+        let interactions = dag
+            .interactions
+            .into_iter()
+            .map(|interaction| {
+                let fields = interaction
+                    .fields
+                    .into_iter()
+                    .map(|idx| exprs[idx].as_ref().clone())
+                    .collect();
+                let count = exprs[interaction.count].as_ref().clone();
+                Interaction {
+                    fields,
+                    count,
+                    bus_index: interaction.bus_index,
+                    interaction_type: interaction.interaction_type,
+                }
+            })
+            .collect::<Vec<_>>();
+        SymbolicConstraints {
+            constraints,
+            interactions,
+        }
+    }
+}
+
+impl<F: Field> From<SymbolicConstraints<F>> for SymbolicConstraintsDag<F> {
+    fn from(sc: SymbolicConstraints<F>) -> Self {
+        build_symbolic_constraints_dag(&sc.constraints, &sc.interactions)
     }
 }
 
@@ -198,17 +285,20 @@ mod tests {
     use p3_baby_bear::BabyBear;
     use p3_field::AbstractField;
 
-    use crate::air_builders::symbolic::{
-        dag::{build_symbolic_expr_dag, SymbolicExpressionDag, SymbolicExpressionNode},
-        symbolic_expression::SymbolicExpression,
-        symbolic_variable::{Entry, SymbolicVariable},
-        SymbolicConstraints,
+    use crate::{
+        air_builders::symbolic::{
+            dag::{build_symbolic_constraints_dag, SymbolicExpressionDag, SymbolicExpressionNode},
+            symbolic_expression::SymbolicExpression,
+            symbolic_variable::{Entry, SymbolicVariable},
+            SymbolicConstraints,
+        },
+        interaction::{Interaction, InteractionType},
     };
 
     type F = BabyBear;
 
     #[test]
-    fn test_symbolic_expressions_dag() {
+    fn test_symbolic_constraints_dag() {
         let expr = SymbolicExpression::Constant(F::ONE)
             * SymbolicVariable::new(
                 Entry::Main {
@@ -217,16 +307,22 @@ mod tests {
                 },
                 3,
             );
-        let exprs = vec![
+        let constraints = vec![
             SymbolicExpression::IsFirstRow * SymbolicExpression::IsLastRow
                 + SymbolicExpression::Constant(F::ONE)
                 + SymbolicExpression::IsFirstRow * SymbolicExpression::IsLastRow
                 + expr.clone(),
             expr.clone() * expr.clone(),
         ];
-        let expr_list = build_symbolic_expr_dag(&exprs);
+        let interactions = vec![Interaction {
+            bus_index: 0,
+            fields: vec![expr.clone(), SymbolicExpression::Constant(F::TWO)],
+            count: SymbolicExpression::Constant(F::ONE),
+            interaction_type: InteractionType::Send,
+        }];
+        let dag = build_symbolic_constraints_dag(&constraints, &interactions);
         assert_eq!(
-            expr_list,
+            dag.constraints,
             SymbolicExpressionDag::<F> {
                 nodes: vec![
                     SymbolicExpressionNode::IsFirstRow,
@@ -274,17 +370,28 @@ mod tests {
                         left_idx: 8,
                         right_idx: 8,
                         degree_multiple: 2
-                    }
+                    },
+                    SymbolicExpressionNode::Constant(F::TWO),
                 ],
                 constraint_idx: vec![9, 10],
             }
         );
+        assert_eq!(
+            dag.interactions,
+            vec![Interaction {
+                bus_index: 0,
+                fields: vec![8, 11],
+                count: 3,
+                interaction_type: InteractionType::Send,
+            }]
+        );
+
         let sc = SymbolicConstraints {
-            constraints: exprs,
-            interactions: vec![],
+            constraints,
+            interactions,
         };
         let ser_str = serde_json::to_string(&sc).unwrap();
         let new_sc: SymbolicConstraints<_> = serde_json::from_str(&ser_str).unwrap();
-        assert_eq!(sc.constraints, new_sc.constraints);
+        assert_eq!(sc, new_sc);
     }
 }
diff --git a/crates/stark-backend/src/air_builders/symbolic/mod.rs b/crates/stark-backend/src/air_builders/symbolic/mod.rs
@@ -7,7 +7,7 @@ use p3_air::{
 use p3_field::Field;
 use p3_matrix::{dense::RowMajorMatrix, Matrix};
 use p3_util::log2_ceil_usize;
-use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use serde::{Deserialize, Serialize};
 use tracing::instrument;
 
 use self::{
@@ -16,7 +16,6 @@ use self::{
 };
 use super::PartitionedAirBuilder;
 use crate::{
-    air_builders::symbolic::dag::{build_symbolic_expr_dag, SymbolicExpressionDag},
     interaction::{
         rap::InteractionPhaseAirBuilder, Interaction, InteractionBuilder, InteractionType,
         RapPhaseSeqKind, SymbolicInteraction,
@@ -31,13 +30,23 @@ pub mod symbolic_variable;
 
 /// Symbolic constraints for a single AIR with interactions.
 /// The constraints contain the constraints on the logup partial sums.
-#[derive(Clone, Debug, Serialize, Deserialize)]
-#[serde(bound = "F: Field")]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(
+    bound = "F: Field",
+    from = "dag::SymbolicConstraintsDag<F>",
+    into = "dag::SymbolicConstraintsDag<F>"
+)]
 pub struct SymbolicConstraints<F> {
     /// All constraints of the RAP, including the constraints on the logup partial sums.
     pub constraints: Vec<SymbolicExpression<F>>,
-    /// Only for debug purposes. `constraints` also contains the constraints on the logup partial sums.
-    pub interactions: Vec<Interaction<SymbolicExpression<F>>>,
+    /// Symbolic representation of chip interactions. This is used by
+    /// the prover for after challenge trace generation, and some partial
+    /// information may be used by the verifier.
+    ///
+    /// **However**, any contributions to the quotient polynomial from
+    /// logup are already included in `constraints` and do not need to
+    /// be separately calculated from `interactions`.
+    pub interactions: Vec<SymbolicInteraction<F>>,
 }
 
 impl<F: Field> SymbolicConstraints<F> {
@@ -467,27 +476,3 @@ fn gen_main_trace<F: Field>(
         .collect_vec();
     RowMajorMatrix::new(mat_values, width)
 }
-
-#[allow(dead_code)]
-fn serialize_symbolic_exprs<F: Field, S>(
-    data: &[SymbolicExpression<F>],
-    serializer: S,
-) -> Result<S::Ok, S::Error>
-where
-    S: Serializer,
-{
-    // Convert the number to a hex string before serializing
-    let dag = build_symbolic_expr_dag(data);
-    dag.serialize(serializer)
-}
-
-#[allow(dead_code)]
-fn deserialize_symbolic_exprs<'de, F: Field, D>(
-    deserializer: D,
-) -> Result<Vec<SymbolicExpression<F>>, D::Error>
-where
-    D: Deserializer<'de>,
-{
-    let dag = SymbolicExpressionDag::deserialize(deserializer)?;
-    Ok(dag.to_symbolic_expressions())
-}
diff --git a/crates/stark-backend/src/air_builders/verifier.rs b/crates/stark-backend/src/air_builders/verifier.rs
@@ -8,7 +8,7 @@ use p3_matrix::Matrix;
 
 use super::{
     symbolic::{
-        dag::{build_symbolic_expr_dag, SymbolicExpressionNode},
+        dag::{build_symbolic_constraints_dag, SymbolicExpressionNode},
         symbolic_expression::{SymbolicEvaluator, SymbolicExpression},
         symbolic_variable::{Entry, SymbolicVariable},
     },
@@ -52,7 +52,7 @@ where
     PubVar: Into<Expr> + Copy + Send + Sync,
 {
     pub fn eval_constraints(&mut self, constraints: &[SymbolicExpression<F>]) {
-        let dag = build_symbolic_expr_dag(constraints);
+        let dag = build_symbolic_constraints_dag(constraints, &[]).constraints;
         // node_idx -> evaluation
         // We do a simple serial evaluation in topological order.
         // This can be parallelized if necessary.
diff --git a/crates/stark-backend/src/interaction/mod.rs b/crates/stark-backend/src/interaction/mod.rs
@@ -18,13 +18,13 @@ pub mod stark_log_up;
 pub mod trace;
 mod utils;
 
-#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub enum InteractionType {
     Send,
     Receive,
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct Interaction<Expr> {
     pub fields: Vec<Expr>,
     pub count: Expr,
