openvm_recursion_circuit/batch_constraint/expr_eval/
dag_commit.rs

1use core::borrow::Borrow;
2use std::{array::from_fn, sync::Arc};
3
4use itertools::{fold, Itertools};
5use openvm_circuit_primitives::{encoder::Encoder, utils::assert_array_eq, ColumnsAir, SubAir};
6use openvm_poseidon2_air::{
7    Poseidon2Config, Poseidon2SubAir, Poseidon2SubChip, Poseidon2SubCols,
8    BABY_BEAR_POSEIDON2_SBOX_DEGREE, POSEIDON2_WIDTH,
9};
10use openvm_recursion_circuit_derive::AlignedBorrow;
11use openvm_stark_backend::{air_builders::sub::SubAirBuilder, interaction::InteractionBuilder};
12use openvm_stark_sdk::config::baby_bear_poseidon2::{DIGEST_SIZE, F};
13use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
14use p3_field::{Field, InjectiveMonomial, PrimeCharacteristicRing, PrimeField};
15
16use crate::{
17    batch_constraint::expr_eval::{
18        CachedRecord, CachedSymbolicExpressionColumns, FLAG_MODULUS, NUM_FLAGS,
19    },
20    utils::assert_zeros,
21};
22
23pub(in crate::batch_constraint) const SBOX_REGISTERS: usize = 1;
24
25pub fn default_poseidon2_sub_chip<
26    F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>,
27>() -> Poseidon2SubChip<F, SBOX_REGISTERS> {
28    Poseidon2SubChip::<F, SBOX_REGISTERS>::new(Poseidon2Config::default().constants)
29}
30
31#[repr(C)]
32#[derive(AlignedBorrow)]
33pub struct DagCommitCols<T> {
34    pub inner: Poseidon2SubCols<T, SBOX_REGISTERS>,
35    pub flags: [T; NUM_FLAGS],
36    pub is_constraint: T,
37}
38
39#[repr(C)]
40#[derive(AlignedBorrow)]
41pub struct DagCommitPvs<T> {
42    pub commit: [T; DIGEST_SIZE],
43}
44
45/// Sub-AIR to compute the onion hash of one digest per row. Expects each AIR
46/// that uses it to have DagCommitPvs as its public value representation.
47pub struct DagCommitSubAir<F: Field> {
48    pub subair: Arc<Poseidon2SubAir<F, SBOX_REGISTERS>>,
49}
50
51// No columns provided: `DagCommitCols` embeds external `Poseidon2SubCols` which doesn't derive
52// `StructReflection`.
53impl<F: Field> ColumnsAir for DagCommitSubAir<F> {}
54
55impl<F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>> DagCommitSubAir<F> {
56    pub fn new() -> Self {
57        Self::default()
58    }
59}
60
61impl<F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>> Default
62    for DagCommitSubAir<F>
63{
64    fn default() -> Self {
65        Self {
66            subair: default_poseidon2_sub_chip().air,
67        }
68    }
69}
70
71impl<F: Field> BaseAir<F> for DagCommitSubAir<F> {
72    fn width(&self) -> usize {
73        DagCommitCols::<u8>::width()
74    }
75}
76
77impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> SubAir<AB>
78    for DagCommitSubAir<AB::F>
79{
80    type AirContext<'a>
81        = (&'a [AB::Var], &'a [AB::Var])
82    where
83        AB::Expr: 'a,
84        AB::Var: 'a,
85        AB: 'a;
86
87    fn eval<'a>(&'a self, builder: &'a mut AB, (local, next): (&'a [AB::Var], &'a [AB::Var]))
88    where
89        AB::Var: 'a,
90        AB::Expr: 'a,
91    {
92        let local: &DagCommitCols<AB::Var> = (*local).borrow();
93        let next: &DagCommitCols<AB::Var> = (*next).borrow();
94
95        let mut sub_builder =
96            SubAirBuilder::<AB, Poseidon2SubAir<AB::F, SBOX_REGISTERS>, AB::F>::new(
97                builder,
98                0..self.subair.width(),
99            );
100        self.subair.eval(&mut sub_builder);
101
102        debug_assert_eq!(FLAG_MODULUS, 3);
103        builder.assert_bool(local.is_constraint);
104        for flag in local.flags {
105            builder.assert_tern(flag);
106        }
107
108        let first_digest_element =
109            collapse_flags(local.flags.map(Into::into), local.is_constraint.into());
110        builder.assert_eq(local.inner.inputs[0], first_digest_element);
111
112        assert_zeros::<_, DIGEST_SIZE>(
113            &mut builder.when_first_row(),
114            from_fn(|i| local.inner.inputs[i + DIGEST_SIZE]),
115        );
116        assert_array_eq::<_, _, _, DIGEST_SIZE>(
117            &mut builder.when_transition(),
118            from_fn(|i| local.inner.ending_full_rounds.last().unwrap().post[i + DIGEST_SIZE]),
119            from_fn(|i| next.inner.inputs[i + DIGEST_SIZE]),
120        );
121
122        let &DagCommitPvs::<_> { commit: pvs_commit } = builder.public_values().borrow();
123
124        assert_array_eq::<_, _, _, DIGEST_SIZE>(
125            &mut builder.when_last_row(),
126            from_fn(|i| local.inner.ending_full_rounds.last().unwrap().post[i]),
127            pvs_commit,
128        );
129    }
130}
131
132#[derive(Debug, Clone)]
133pub struct DagCommitInfo<F: Clone> {
134    pub commit: [F; DIGEST_SIZE],
135    pub poseidon2_inputs: Vec<[F; POSEIDON2_WIDTH]>,
136}
137
138/// Collapse flags and is_constraint into a single field element, which should be the
139/// first element of a row's input digest.
140///
141/// WARNING: To use this in an AIR you MUST constrain that is_constraint is boolean
142/// and that each flag is in [0, FLAG_MODULUS). This ensures that the element doesn't
143/// overflow for any 13-bit field or higher.
144pub fn collapse_flags<F: PrimeCharacteristicRing>(flags: [F; NUM_FLAGS], is_constraint: F) -> F {
145    fold(
146        flags.iter().enumerate(),
147        is_constraint.clone(),
148        |acc, (pow_exp, flag)| {
149            acc + (flag.clone() * F::from_u32(FLAG_MODULUS.pow(pow_exp as u32) << 1))
150        },
151    )
152}
153
154/// Compresses a CachedSymbolicExpressionColumns row into a digest. Uses collapse_flags for
155/// the first element.
156pub fn cached_symbolic_expr_cols_to_digest<F: PrimeCharacteristicRing>(
157    cached_cols: &[F],
158) -> [F; DIGEST_SIZE] {
159    let cached_cols: &CachedSymbolicExpressionColumns<_> = cached_cols.borrow();
160    let mut ret = [F::ZERO; DIGEST_SIZE];
161    ret[0] = collapse_flags(cached_cols.flags.clone(), cached_cols.is_constraint.clone());
162    ret[1] = cached_cols.air_idx.clone();
163    ret[2] = cached_cols.node_or_interaction_idx.clone();
164    ret[3] = cached_cols.attrs[0].clone();
165    ret[4] = cached_cols.attrs[1].clone();
166    ret[5] = cached_cols.attrs[2].clone();
167    ret[6] = cached_cols.fanout.clone();
168    ret[7] = cached_cols.constraint_idx.clone();
169    ret
170}
171
172/// Converts an input digest (+ flags and is_constraint) into cached columns.
173pub fn digest_to_cached_symbolic_expr_cols<F: Copy>(
174    digest: [F; DIGEST_SIZE],
175    flags: [F; NUM_FLAGS],
176    is_constraint: F,
177) -> CachedSymbolicExpressionColumns<F> {
178    CachedSymbolicExpressionColumns {
179        flags,
180        air_idx: digest[1],
181        node_or_interaction_idx: digest[2],
182        attrs: from_fn(|i| digest[3..][i]),
183        fanout: digest[6],
184        is_constraint,
185        constraint_idx: digest[7],
186    }
187}
188
189pub fn dag_commit_cols_to_cached_cols<F: Copy>(
190    dag_commit_cols: &[F],
191) -> CachedSymbolicExpressionColumns<F> {
192    let cols: &DagCommitCols<_> = dag_commit_cols.borrow();
193    let digest = from_fn(|i| cols.inner.inputs[i]);
194    digest_to_cached_symbolic_expr_cols(digest, cols.flags, cols.is_constraint)
195}
196
197pub(crate) fn generate_dag_commit_info(
198    cached_records: &[CachedRecord],
199    encoder: Encoder,
200) -> DagCommitInfo<F> {
201    let sub_chip = Poseidon2SubChip::<F, SBOX_REGISTERS>::new(Poseidon2Config::default().constants);
202    let mut state = [F::ZERO; POSEIDON2_WIDTH];
203
204    let height = cached_records.len().next_power_of_two();
205    let poseidon2_inputs = (0..height)
206        .map(|row_idx| {
207            let digest: [F; DIGEST_SIZE] = if row_idx < cached_records.len() {
208                let r = &cached_records[row_idx];
209                let encoder_flag = encoder.get_flag_pt(r.kind as usize);
210                let columns = CachedSymbolicExpressionColumns {
211                    flags: from_fn(|i| F::from_u32(encoder_flag[i])),
212                    air_idx: F::from_usize(r.air_idx),
213                    node_or_interaction_idx: F::from_usize(r.node_idx),
214                    attrs: r.attrs.map(F::from_usize),
215                    fanout: F::from_usize(r.fanout),
216                    is_constraint: F::from_bool(r.is_constraint),
217                    constraint_idx: F::from_usize(r.constraint_idx),
218                };
219                cached_symbolic_expr_cols_to_digest(&columns.to_vec())
220            } else {
221                [F::ZERO; DIGEST_SIZE]
222            };
223
224            state[..DIGEST_SIZE].copy_from_slice(&digest);
225            let input = state;
226            sub_chip.permute_mut(&mut state);
227            input
228        })
229        .collect_vec();
230
231    DagCommitInfo {
232        commit: from_fn(|i| state[i]),
233        poseidon2_inputs,
234    }
235}
openvm_recursion_circuit/batch_constraint/expr_eval/dag_commit.rs

openvm_recursion_circuit/batch_constraint/expr_eval/
dag_commit.rs
