From 9037ea465cd1e1e2c0f21cac289aa2c4ba054c86 Mon Sep 17 00:00:00 2001 From: Alon Titelman Date: Thu, 14 Nov 2024 15:57:30 +0200 Subject: [PATCH] Formal Expr logup variables. --- .../prover/src/constraint_framework/expr.rs | 185 +++++++++++++++++- crates/prover/src/constraint_framework/mod.rs | 6 + 2 files changed, 181 insertions(+), 10 deletions(-) diff --git a/crates/prover/src/constraint_framework/expr.rs b/crates/prover/src/constraint_framework/expr.rs index d5724b125..fb86edac8 100644 --- a/crates/prover/src/constraint_framework/expr.rs +++ b/crates/prover/src/constraint_framework/expr.rs @@ -3,7 +3,7 @@ use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub}; use num_traits::{One, Zero}; use super::logup::{LogupAtRow, LogupSums}; -use super::{EvalAtRow, INTERACTION_TRACE_IDX}; +use super::{EvalAtRow, Relation, RelationEntry, INTERACTION_TRACE_IDX}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::FieldExpOps; @@ -14,7 +14,7 @@ use crate::core::lookups::utils::Fraction; struct ColumnExpr { interaction: usize, idx: usize, - offset: usize, + offset: isize, } #[derive(Clone, Debug, PartialEq)] @@ -27,6 +27,8 @@ enum Expr { /// `SecureCol(Add(a0, b0), Add(a1, b1), Add(a2, b2), Add(a3, b3))` SecureCol([Box; 4]), Const(BaseField), + /// Formal parameter to the AIR, for example the interaction elements of a relation. + Param(String), Add(Box, Box), Sub(Box, Box), Mul(Box, Box), @@ -148,8 +150,25 @@ impl AddAssign for Expr { } } +/// Returns the expression +/// `value[0] * _alpha0 + value[1] * _alpha1 + ... - _z.` +fn combine_formal>(relation: &R, values: &[Expr]) -> Expr { + const Z_SUFFIX: &str = "_z"; + const ALPHA_SUFFIX: &str = "_alpha"; + + let z = Expr::Param(relation.get_name().to_owned() + Z_SUFFIX); + let alpha_powers = (0..relation.get_size()) + .map(|i| Expr::Param(relation.get_name().to_owned() + ALPHA_SUFFIX + &i.to_string())); + values + .iter() + .zip(alpha_powers) + .fold(Expr::zero(), |acc, (value, power)| { + acc + power * value.clone() + }) + - z +} + /// An Evaluator that saves all constraint expressions. -#[derive(Default)] struct ExprEvaluator { pub cur_var_index: usize, pub constraints: Vec, @@ -157,7 +176,8 @@ struct ExprEvaluator { } impl ExprEvaluator { - pub fn _new(log_size: u32, logup_sums: LogupSums) -> Self { + #[allow(dead_code)] + pub fn new(log_size: u32, logup_sums: LogupSums) -> Self { Self { cur_var_index: Default::default(), constraints: Default::default(), @@ -180,7 +200,7 @@ impl EvalAtRow for ExprEvaluator { let col = ColumnExpr { interaction, idx: self.cur_var_index, - offset: offsets[i] as usize, + offset: offsets[i], }; self.cur_var_index += 1; Expr::Col(col) @@ -203,21 +223,44 @@ impl EvalAtRow for ExprEvaluator { ]) } + fn add_to_relation>( + &mut self, + entries: &[RelationEntry<'_, Self::F, Self::EF, R>], + ) { + let fracs: Vec> = entries + .iter() + .map( + |RelationEntry { + relation, + multiplicity, + values, + }| { + Fraction::new(multiplicity.clone(), combine_formal(*relation, values)) + }, + ) + .collect(); + self.write_logup_frac(fracs.into_iter().sum()); + } + super::logup_proxy!(); } #[cfg(test)] mod tests { - use num_traits::One; + use num_traits::{One, Zero}; use crate::constraint_framework::expr::{ColumnExpr, Expr, ExprEvaluator}; - use crate::constraint_framework::{EvalAtRow, FrameworkEval, ORIGINAL_TRACE_IDX}; + use crate::constraint_framework::{ + relation, EvalAtRow, FrameworkEval, RelationEntry, ORIGINAL_TRACE_IDX, + }; + use crate::core::fields::m31::M31; + use crate::core::fields::qm31::SecureField; use crate::core::fields::FieldExpOps; #[test] fn test_expr_eval() { let test_struct = TestStruct {}; - let eval = test_struct.evaluate(ExprEvaluator::default()); - assert_eq!(eval.constraints.len(), 1); + let eval = test_struct.evaluate(ExprEvaluator::new(16, (SecureField::zero(), None))); + assert_eq!(eval.constraints.len(), 2); assert_eq!( eval.constraints[0], Expr::Mul( @@ -257,8 +300,122 @@ mod tests { )) ) ); + + assert_eq!( + eval.constraints[1], + Expr::Mul( + Box::new(Expr::Const(M31(1))), + Box::new(Expr::Sub( + Box::new(Expr::Mul( + Box::new(Expr::Sub( + Box::new(Expr::Sub( + Box::new(Expr::SecureCol([ + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 4, + offset: 0 + })), + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 6, + offset: 0 + })), + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 8, + offset: 0 + })), + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 10, + offset: 0 + })) + ])), + Box::new(Expr::Sub( + Box::new(Expr::SecureCol([ + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 5, + offset: -1 + })), + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 7, + offset: -1 + })), + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 9, + offset: -1 + })), + Box::new(Expr::Col(ColumnExpr { + interaction: 2, + idx: 11, + offset: -1 + })) + ])), + Box::new(Expr::Mul( + Box::new(Expr::Col(ColumnExpr { + interaction: 0, + idx: 3, + offset: 0 + })), + Box::new(Expr::SecureCol([ + Box::new(Expr::Const(M31(0))), + Box::new(Expr::Const(M31(0))), + Box::new(Expr::Const(M31(0))), + Box::new(Expr::Const(M31(0))) + ])) + )) + )) + )), + Box::new(Expr::Const(M31(0))) + )), + Box::new(Expr::Sub( + Box::new(Expr::Add( + Box::new(Expr::Add( + Box::new(Expr::Add( + Box::new(Expr::Const(M31(0))), + Box::new(Expr::Mul( + Box::new(Expr::Param( + "TestRelation_alpha0".to_string() + )), + Box::new(Expr::Col(ColumnExpr { + interaction: 1, + idx: 0, + offset: 0 + })) + )) + )), + Box::new(Expr::Mul( + Box::new(Expr::Param("TestRelation_alpha1".to_string())), + Box::new(Expr::Col(ColumnExpr { + interaction: 1, + idx: 1, + offset: 0 + })) + )) + )), + Box::new(Expr::Mul( + Box::new(Expr::Param("TestRelation_alpha2".to_string())), + Box::new(Expr::Col(ColumnExpr { + interaction: 1, + idx: 2, + offset: 0 + })) + )) + )), + Box::new(Expr::Param("TestRelation_z".to_string())) + )) + )), + Box::new(Expr::Const(M31(1))) + )) + ) + ); } + relation!(TestRelation, 3); + struct TestStruct {} impl FrameworkEval for TestStruct { fn log_size(&self) -> u32 { @@ -271,7 +428,15 @@ mod tests { let x0 = eval.next_trace_mask(); let x1 = eval.next_trace_mask(); let x2 = eval.next_trace_mask(); - eval.add_constraint(x0.clone() * x1.clone() * x2 * (x0 + x1).inverse()); + eval.add_constraint( + x0.clone() * x1.clone() * x2.clone() * (x0.clone() + x1.clone()).inverse(), + ); + eval.add_to_relation(&[RelationEntry::new( + &TestRelation::dummy(), + E::EF::one(), + &[x0, x1, x2], + )]); + eval.finalize_logup(); eval } } diff --git a/crates/prover/src/constraint_framework/mod.rs b/crates/prover/src/constraint_framework/mod.rs index e560b94ee..40b145f20 100644 --- a/crates/prover/src/constraint_framework/mod.rs +++ b/crates/prover/src/constraint_framework/mod.rs @@ -222,6 +222,7 @@ pub trait Relation>: Sized { fn combine(&self, values: &[F]) -> EF; fn get_name(&self) -> &str; + fn get_size(&self) -> usize; } /// A struct representing a relation entry. @@ -250,6 +251,7 @@ macro_rules! relation { #[derive(Clone, Debug, PartialEq)] pub struct $name(crate::constraint_framework::logup::LookupElements<$size>); + #[allow(dead_code)] impl $name { pub fn dummy() -> Self { Self(crate::constraint_framework::logup::LookupElements::dummy()) @@ -277,6 +279,10 @@ macro_rules! relation { fn get_name(&self) -> &str { stringify!($name) } + + fn get_size(&self) -> usize { + $size + } } }; }