Skip to content

Commit

Permalink
Formal Expr logup variables.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Alon-Ti
Alon-Ti committed Nov 18, 2024
1 parent 3a88b3a commit b45fd93
Show file tree
Hide file tree
Showing 2 changed files with 181 additions and 10 deletions.
185 changes: 175 additions & 10 deletions crates/prover/src/constraint_framework/expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -14,7 +14,7 @@ use crate::core::lookups::utils::Fraction;
struct ColumnExpr {
interaction: usize,
idx: usize,
offset: usize,
offset: isize,
}

#[derive(Clone, Debug, PartialEq)]
Expand All @@ -27,6 +27,8 @@ enum Expr {
/// `SecureCol(Add(a0, b0), Add(a1, b1), Add(a2, b2), Add(a3, b3))`
SecureCol([Box<Expr>; 4]),
Const(BaseField),
/// Formal parameter to the AIR, for example the interaction elements of a relation.
Param(String),
Add(Box<Expr>, Box<Expr>),
Sub(Box<Expr>, Box<Expr>),
Mul(Box<Expr>, Box<Expr>),
Expand Down Expand Up @@ -148,16 +150,34 @@ impl AddAssign<BaseField> for Expr {
}
}

/// Returns the expression
/// `value[0] * <relation>_alpha0 + value[1] * <relation>_alpha1 + ... - <relation>_z.`
fn combine_formal<R: Relation<Expr, Expr>>(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<Expr>,
pub logup: LogupAtRow<Self>,
}

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(),
Expand All @@ -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)
Expand All @@ -203,21 +223,44 @@ impl EvalAtRow for ExprEvaluator {
])
}

fn add_to_relation<R: Relation<Self::F, Self::EF>>(
&mut self,
entries: &[RelationEntry<'_, Self::F, Self::EF, R>],
) {
let fracs: Vec<Fraction<Self::EF, Self::EF>> = 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(
Expand Down Expand Up @@ -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 {
Expand All @@ -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
}
}
Expand Down
6 changes: 6 additions & 0 deletions crates/prover/src/constraint_framework/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ pub trait Relation<F: Clone, EF: RelationEFTraitBound<F>>: Sized {
fn combine(&self, values: &[F]) -> EF;

fn get_name(&self) -> &str;
fn get_size(&self) -> usize;
}

/// A struct representing a relation entry.
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -277,6 +279,10 @@ macro_rules! relation {
fn get_name(&self) -> &str {
stringify!($name)
}

fn get_size(&self) -> usize {
$size
}
}
};
}
Expand Down

0 comments on commit b45fd93

Please sign in to comment.