fix scaled sum degree

scroll-tech · Oct 2, 2024 · 91ad507 · 91ad507
1 parent bea2f4f
commit 91ad507
Showing 1 changed file with 50 additions and 1 deletion.
diff --git a/ceno_zkvm/src/expression.rs b/ceno_zkvm/src/expression.rs
@@ -52,7 +52,9 @@ impl<E: ExtensionField> Expression<E> {
             Expression::Constant(_) => 0,
             Expression::Sum(a_expr, b_expr) => max(a_expr.degree(), b_expr.degree()),
             Expression::Product(a_expr, b_expr) => a_expr.degree() + b_expr.degree(),
-            Expression::ScaledSum(_, _, _) => 1,
+            Expression::ScaledSum(x_expr, a_expr, b_expr) => {
+                max(x_expr.degree() + a_expr.degree(), b_expr.degree())
+            }
             Expression::Challenge(_, _, _, _) => 0,
         }
     }
@@ -794,4 +796,51 @@ mod tests {
         assert_eq!(s, "WitIn(0)");
         assert_eq!(wtns_acc, vec![0]);
     }
+
+    #[test]
+    fn test_degree_1() {
+        let expr: Expression<GoldilocksExt2> = Expression::WitIn(1);
+        assert_eq!(expr.degree(), 1);
+    }
+
+    #[test]
+    fn test_degree_2_product() {
+        let expr: Expression<GoldilocksExt2> = Expression::Product(
+            Box::new(Expression::WitIn(1)),
+            Box::new(Expression::WitIn(2)),
+        );
+        assert_eq!(expr.degree(), 2);
+    }
+
+    #[test]
+    fn test_degree_3_sum() {
+        let expr: Expression<GoldilocksExt2> = Expression::Sum(
+            Box::new(Expression::WitIn(1)),
+            Box::new(Expression::WitIn(2)),
+        );
+        assert_eq!(expr.degree(), 1);
+    }
+
+    #[test]
+    fn test_degree_4_scaled_sum() {
+        let expr: Expression<GoldilocksExt2> = Expression::ScaledSum(
+            Box::new(Expression::WitIn(1)),
+            Box::new(Expression::WitIn(2)),
+            Box::new(Expression::WitIn(3)),
+        );
+        assert_eq!(expr.degree(), 2);
+    }
+
+    #[test]
+    fn test_degree_5_scaled_sum() {
+        let expr: Expression<GoldilocksExt2> = Expression::ScaledSum(
+            Box::new(Expression::WitIn(1)),
+            Box::new(Expression::Product(
+                Box::new(Expression::WitIn(1)),
+                Box::new(Expression::WitIn(2)),
+            )),
+            Box::new(Expression::WitIn(3)),
+        );
+        assert_eq!(expr.degree(), 3);
+    }
 }