pub mod add;
pub mod mul;
pub mod require;
pub mod store_mul;
pub mod trace;
use alloc::{vec, vec::Vec};
use core::array;
use miden_core::{
Felt,
field::{Algebra, PrimeCharacteristicRing, QuadFelt},
utils::RowMajorMatrix,
};
use miden_crypto::stark::air::ExtensionBuilder;
use miden_lifted_air::{AirBuilder, BaseAir, LiftedAir, LiftedAirBuilder};
pub use require::{UintRequire, UintStores};
use crate::{
logup::{
Challenges, CyclicConstraintLookupBuilder, Deg, LookupAir, LookupBatch, LookupBuilder,
LookupColumn, LookupGroup, LookupMessage, NUM_PUBLIC_VALUES, NUM_RANDOMNESS,
NUM_SIGMA_VALUES,
},
primitives::byte_pair_lut::Range16Msg,
relations::{BusId, MAX_MESSAGE_WIDTH, NUM_BUS_IDS},
utils::{current_main, next_main},
};
#[derive(Debug, Clone)]
pub struct UintValMsg<E> {
pub ptr: E,
pub bound_ptr: E,
pub limbs: [E; 8],
}
impl<E, EF> LookupMessage<E, EF> for UintValMsg<E>
where
E: Algebra<E>,
EF: Algebra<E>,
{
fn encode(&self, challenges: &Challenges<EF>) -> EF {
let [c0, c1, c2, c3, c4, c5, c6, c7] = self.limbs.clone();
challenges.encode(
BusId::UintVal as usize,
[self.ptr.clone(), self.bound_ptr.clone(), c0, c1, c2, c3, c4, c5, c6, c7],
)
}
}
#[derive(Debug, Clone)]
pub struct UintLimbsMsg<E> {
pub ptr: E,
pub bound_ptr: E,
pub limbs: [E; 16],
}
impl<E, EF> LookupMessage<E, EF> for UintLimbsMsg<E>
where
E: Algebra<E>,
EF: Algebra<E>,
{
fn encode(&self, challenges: &Challenges<EF>) -> EF {
let [l0, l1, l2, l3, l4, l5, l6, l7, l8, l9, l10, l11, l12, l13, l14, l15] =
self.limbs.clone();
challenges.encode(
BusId::UintLimbs as usize,
[
self.ptr.clone(),
self.bound_ptr.clone(),
l0,
l1,
l2,
l3,
l4,
l5,
l6,
l7,
l8,
l9,
l10,
l11,
l12,
l13,
l14,
l15,
],
)
}
}
pub const NUM_CELLS: usize = 16;
pub const COL_PTR: usize = NUM_CELLS;
pub const COL_BOUND_PTR: usize = NUM_CELLS + 1;
pub const NUM_MAIN_COLS: usize = NUM_CELLS + 2;
pub const HUB_CELL_UINTVAL_MULT: usize = 8;
pub const HUB_CELL_UINTLIMBS_MULT: usize = 9;
pub const CARRY_LO_BEGIN: usize = 4;
pub const CARRY_HI_BEGIN: usize = 12;
pub const TERM_CELL_GAP: usize = 15;
pub const PERIOD: usize = 4;
const PCOL_V_LO: usize = 0;
const PCOL_V_HI: usize = 1;
const PCOL_COMP: usize = 2;
const PCOL_BOUND: usize = 3;
pub(crate) const NUM_LOGUP_COLS: usize = 1 + 1 + NUM_CELLS.div_ceil(2) + 1; const REGISTER_COL: usize = NUM_LOGUP_COLS;
const AUX_WIDTH: usize = NUM_LOGUP_COLS + 1;
pub(crate) const COLUMN_SHAPE: [usize; NUM_LOGUP_COLS] = {
let mut shape = [2usize; NUM_LOGUP_COLS];
shape[0] = 1;
shape[NUM_LOGUP_COLS - 1] = 1;
shape
};
#[derive(Debug, Default, Clone, Copy)]
pub struct UintStoreAir;
impl BaseAir<Felt> for UintStoreAir {
fn width(&self) -> usize {
NUM_MAIN_COLS
}
fn num_public_values(&self) -> usize {
NUM_PUBLIC_VALUES
}
fn periodic_columns(&self) -> Vec<Vec<Felt>> {
let o = Felt::ONE;
let z = Felt::ZERO;
vec![
vec![o, z, z, z], vec![z, o, z, z], vec![z, z, o, z], vec![z, z, z, o], ]
}
}
impl LiftedAir<Felt, QuadFelt> for UintStoreAir {
fn num_randomness(&self) -> usize {
NUM_RANDOMNESS
}
fn aux_width(&self) -> usize {
AUX_WIDTH
}
fn num_aux_values(&self) -> usize {
NUM_SIGMA_VALUES
}
fn build_aux_trace(
&self,
main: &RowMajorMatrix<Felt>,
_air_inputs: &[Felt],
_aux_inputs: &[Felt],
challenges: &[QuadFelt],
) -> (RowMajorMatrix<QuadFelt>, Vec<QuadFelt>) {
trace::build_aux(main, challenges)
}
fn eval<AB: LiftedAirBuilder<F = Felt>>(&self, builder: &mut AB) {
let local: [AB::Var; NUM_MAIN_COLS] = current_main(builder.main(), 0);
let next: [AB::Var; NUM_MAIN_COLS] = next_main(builder.main(), 0);
let (v_lo_sel, v_hi_sel, comp_sel, bound_sel): (AB::Expr, AB::Expr, AB::Expr, AB::Expr) = {
let p = builder.periodic_values();
(
p[PCOL_V_LO].into(),
p[PCOL_V_HI].into(),
p[PCOL_COMP].into(),
p[PCOL_BOUND].into(),
)
};
let beta: AB::ExprEF = builder.permutation_randomness()[1].into();
let mut bp: Vec<AB::ExprEF> = Vec::with_capacity(8);
bp.push(AB::ExprEF::ONE);
for i in 1..8 {
bp.push(bp[i - 1].clone() * beta.clone());
}
let id: AB::ExprEF =
current_main::<_, AB::VarEF, 1>(builder.permutation(), REGISTER_COL)[0].into();
let id_next: AB::ExprEF =
next_main::<_, AB::VarEF, 1>(builder.permutation(), REGISTER_COL)[0].into();
let two16: AB::Expr = AB::Expr::from(Felt::from(1u32 << 16));
let mut recomb_lo07 = AB::ExprEF::ZERO;
let mut recomb_hi07 = AB::ExprEF::ZERO;
let mut recomb_hi815 = AB::ExprEF::ZERO;
let mut direct_lo = AB::ExprEF::ZERO;
let mut direct_hi = AB::ExprEF::ZERO;
for k in 0..4 {
let r07: AB::Expr =
AB::Expr::from(local[2 * k]) + two16.clone() * AB::Expr::from(local[2 * k + 1]);
let r815: AB::Expr = AB::Expr::from(local[8 + 2 * k])
+ two16.clone() * AB::Expr::from(local[8 + 2 * k + 1]);
recomb_lo07 += bp[k].clone() * r07.clone();
recomb_hi07 += bp[4 + k].clone() * r07;
recomb_hi815 += bp[4 + k].clone() * r815;
direct_lo += bp[k].clone() * AB::Expr::from(local[k]);
direct_hi += bp[4 + k].clone() * AB::Expr::from(local[8 + k]);
}
let t32: AB::Expr = AB::Expr::from(Felt::new(1u64 << 32).expect("2^32 < Goldilocks p"));
let mut carry_lo_term = AB::ExprEF::ZERO;
for j in 0..4 {
let weight: AB::ExprEF = bp[j + 1].clone() - bp[j].clone() * t32.clone();
carry_lo_term += weight * AB::Expr::from(local[CARRY_LO_BEGIN + j]);
}
let mut carry_hi_term = AB::ExprEF::ZERO;
for j in 4..7 {
let weight: AB::ExprEF = bp[j + 1].clone() - bp[j].clone() * t32.clone();
carry_hi_term += weight * AB::Expr::from(local[CARRY_HI_BEGIN + (j - 4)]);
}
let contrib: AB::ExprEF = recomb_lo07.clone() * (v_lo_sel + comp_sel.clone())
+ recomb_hi07 * v_hi_sel
+ recomb_hi815 * comp_sel
+ (carry_lo_term.clone() - direct_lo.clone() + carry_hi_term.clone()
- direct_hi.clone())
* bound_sel.clone();
builder.when_first_row().assert_zero_ext(id.clone());
builder.when_transition().assert_zero_ext(id_next - id.clone() - contrib);
let bound_own: AB::ExprEF = carry_lo_term - direct_lo + carry_hi_term - direct_hi;
builder.assert_zero_ext((id + bound_own) * bound_sel.clone());
let ptr: AB::Expr = local[COL_PTR].into();
builder.when_first_row().assert_zero(ptr - AB::Expr::ONE);
for &cell in
[CARRY_LO_BEGIN, CARRY_LO_BEGIN + 1, CARRY_LO_BEGIN + 2, CARRY_LO_BEGIN + 3].iter()
{
let lj: AB::Expr = local[cell].into();
builder.assert_zero(bound_sel.clone() * lj.clone() * (AB::Expr::ONE - lj));
}
for &cell in [CARRY_HI_BEGIN, CARRY_HI_BEGIN + 1, CARRY_HI_BEGIN + 2].iter() {
let lj: AB::Expr = local[cell].into();
builder.assert_zero(bound_sel.clone() * lj.clone() * (AB::Expr::ONE - lj));
}
let not_term: AB::Expr = AB::Expr::ONE - bound_sel.clone();
for col in [COL_PTR, COL_BOUND_PTR] {
let here: AB::Expr = local[col].into();
let there: AB::Expr = next[col].into();
builder.assert_zero(not_term.clone() * (there - here));
}
let gap: AB::Expr = local[TERM_CELL_GAP].into();
let ptr_here: AB::Expr = local[COL_PTR].into();
let ptr_next: AB::Expr = next[COL_PTR].into();
builder
.when_transition()
.assert_zero(bound_sel * (gap + ptr_here + AB::Expr::ONE - ptr_next));
let mut lb =
CyclicConstraintLookupBuilder::new(builder, self, self.preprocessed_width() > 0);
<Self as LookupAir<_>>::eval(self, &mut lb);
}
}
impl<LB> LookupAir<LB> for UintStoreAir
where
LB: LookupBuilder<F = Felt>,
{
fn num_columns(&self) -> usize {
NUM_LOGUP_COLS
}
fn column_shape(&self) -> &[usize] {
&COLUMN_SHAPE
}
fn max_message_width(&self) -> usize {
MAX_MESSAGE_WIDTH
}
fn num_bus_ids(&self) -> usize {
NUM_BUS_IDS
}
fn eval(&self, builder: &mut LB) {
let local: [LB::Var; NUM_MAIN_COLS] = current_main(builder.main(), 0);
let next: [LB::Var; NUM_MAIN_COLS] = next_main(builder.main(), 0);
let (v_lo_sel, v_hi_sel, comp_sel, bound_sel): (LB::Expr, LB::Expr, LB::Expr, LB::Expr) = {
let p = builder.periodic_values();
(
p[PCOL_V_LO].into(),
p[PCOL_V_HI].into(),
p[PCOL_COMP].into(),
p[PCOL_BOUND].into(),
)
};
let ptr: LB::Expr = local[COL_PTR].into();
let bound_ptr: LB::Expr = local[COL_BOUND_PTR].into();
let neg_mult: LB::Expr = LB::Expr::ZERO - next[HUB_CELL_UINTVAL_MULT].into();
let neg_limbs_mult: LB::Expr = LB::Expr::ZERO - next[HUB_CELL_UINTLIMBS_MULT].into();
let two16: LB::Expr = LB::Expr::from(Felt::from(1u32 << 16));
let recomb: [LB::Expr; 8] = array::from_fn(|k| {
if k < 4 {
local[2 * k].into() + two16.clone() * local[2 * k + 1].into()
} else {
let k = k - 4;
next[2 * k].into() + two16.clone() * next[2 * k + 1].into()
}
});
let direct: [LB::Expr; 8] =
array::from_fn(|k| if k < 4 { local[k].into() } else { local[4 + k].into() });
let raw: [LB::Expr; 16] =
array::from_fn(|j| if j < 8 { local[j].into() } else { next[j - 8].into() });
let provide_deg = Deg { v: 2, u: 1 };
let consume_deg = Deg { v: 1, u: 1 };
let pair_deg = Deg { v: 3, u: 2 };
let rc_deg = Deg { v: 1, u: 1 };
builder.next_column(
|col| {
col.group(
"uintval",
|g| {
g.batch(
"f",
LB::Expr::ONE,
|b| {
b.insert(
"provide",
neg_mult * v_lo_sel.clone(),
UintValMsg {
ptr: ptr.clone(),
bound_ptr: bound_ptr.clone(),
limbs: recomb,
},
provide_deg,
);
},
provide_deg,
);
},
provide_deg,
);
},
provide_deg,
);
builder.next_column(
|col| {
col.group(
"uintval",
|g| {
g.batch(
"f",
LB::Expr::ONE,
|b| {
b.insert(
"consume",
bound_sel.clone(),
UintValMsg {
ptr: bound_ptr.clone(),
bound_ptr: bound_ptr.clone(),
limbs: direct,
},
consume_deg,
);
b.insert(
"range16-gap",
bound_sel.clone(),
Range16Msg { w: local[TERM_CELL_GAP].into() },
consume_deg,
);
},
pair_deg,
);
},
pair_deg,
);
},
pair_deg,
);
let cell_gate = |cell: usize| -> LB::Expr {
if cell < 8 {
v_lo_sel.clone() + v_hi_sel.clone() + comp_sel.clone()
} else {
comp_sel.clone()
}
};
let cell_specs: Vec<(LB::Expr, usize)> =
(0..NUM_CELLS).map(|cell| (cell_gate(cell), cell)).collect();
for group in cell_specs
.chunks(2)
.map(<[(<LB as LookupBuilder>::Expr, usize)]>::to_vec)
.collect::<Vec<_>>()
{
builder.next_column(
|col| {
col.group(
"range16",
|g| {
g.batch(
"f",
LB::Expr::ONE,
|b| {
for (mult, cell) in group {
b.insert(
"range16-limb",
mult,
Range16Msg { w: local[cell].into() },
rc_deg,
);
}
},
pair_deg,
);
},
pair_deg,
);
},
pair_deg,
);
}
builder.next_column(
|col| {
col.group(
"uintlimbs",
|g| {
g.batch(
"f",
LB::Expr::ONE,
|b| {
b.insert(
"provide-raw",
neg_limbs_mult * v_lo_sel,
UintLimbsMsg { ptr, bound_ptr, limbs: raw },
provide_deg,
);
},
provide_deg,
);
},
provide_deg,
);
},
provide_deg,
);
}
}