openvm_sha2_air/
air.rs

1use std::{iter::once, marker::PhantomData};
2
3use ndarray::s;
4use openvm_circuit_primitives::{
5    bitwise_op_lookup::BitwiseOperationLookupBus, encoder::Encoder, utils::select, ColumnsAir,
6    SubAir,
7};
8use openvm_stark_backend::{
9    interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
10    p3_air::{AirBuilder, BaseAir},
11    p3_field::{Field, PrimeCharacteristicRing},
12    p3_matrix::Matrix,
13};
14
15use super::{
16    big_sig0_field, big_sig1_field, ch_field, compose, maj_field, small_sig0_field,
17    small_sig1_field,
18};
19use crate::{
20    constraint_word_addition, word_into_u16_limbs, Sha2BlockHasherSubairConfig, Sha2DigestColsRef,
21    Sha2RoundColsRef,
22};
23
24/// Expects the message to be padded to a multiple of C::BLOCK_WORDS * C::WORD_BITS bits
25#[derive(Clone, Debug)]
26pub struct Sha2BlockHasherSubAir<C: Sha2BlockHasherSubairConfig> {
27    pub bitwise_lookup_bus: BitwiseOperationLookupBus,
28    pub row_idx_encoder: Encoder,
29    /// Internal bus for self-interactions in this AIR.
30    pub private_bus: PermutationCheckBus,
31    _phantom: PhantomData<C>,
32}
33
34// No columns provided: width is the config-dependent `C::SUBAIR_WIDTH` and rows are accessed via
35// `Sha2{Round,Digest}ColsRef` (slice-borrowing ref structs, no static `Cols` to reflect).
36impl<C: Sha2BlockHasherSubairConfig> ColumnsAir for Sha2BlockHasherSubAir<C> {}
37
38impl<C: Sha2BlockHasherSubairConfig> Sha2BlockHasherSubAir<C> {
39    pub fn new(bitwise_lookup_bus: BitwiseOperationLookupBus, private_bus_idx: BusIndex) -> Self {
40        Self {
41            bitwise_lookup_bus,
42            row_idx_encoder: Encoder::new(C::ROWS_PER_BLOCK + 1, 2, false), /* + 1 for dummy
43                                                                             *   (padding) rows */
44            private_bus: PermutationCheckBus::new(private_bus_idx),
45            _phantom: PhantomData,
46        }
47    }
48}
49
50impl<F, C: Sha2BlockHasherSubairConfig> BaseAir<F> for Sha2BlockHasherSubAir<C> {
51    fn width(&self) -> usize {
52        C::SUBAIR_WIDTH
53    }
54}
55
56impl<AB: InteractionBuilder, C: Sha2BlockHasherSubairConfig> SubAir<AB>
57    for Sha2BlockHasherSubAir<C>
58{
59    /// The start column for the sub-air to use
60    type AirContext<'a>
61        = usize
62    where
63        Self: 'a,
64        AB: 'a,
65        <AB as AirBuilder>::Var: 'a,
66        <AB as AirBuilder>::Expr: 'a;
67
68    fn eval<'a>(&'a self, builder: &'a mut AB, start_col: Self::AirContext<'a>)
69    where
70        AB::Var: 'a,
71        AB::Expr: 'a,
72    {
73        self.eval_row(builder, start_col);
74        self.eval_transitions(builder, start_col);
75    }
76}
77
78impl<C: Sha2BlockHasherSubairConfig> Sha2BlockHasherSubAir<C> {
79    /// Implements the single row constraints (i.e. imposes constraints only on local)
80    /// Implements some sanity constraints on the row index, flags, and work variables
81    fn eval_row<AB: InteractionBuilder>(&self, builder: &mut AB, start_col: usize) {
82        let main = builder.main();
83        let local = main.row_slice(0).unwrap();
84
85        // Doesn't matter which column struct we use here as we are only interested in the common
86        // columns
87        let local_cols: Sha2DigestColsRef<AB::Var> =
88            Sha2DigestColsRef::from::<C>(&local[start_col..start_col + C::SUBAIR_DIGEST_WIDTH]);
89        let flags = &local_cols.flags;
90        builder.assert_bool(*flags.is_round_row);
91        builder.assert_bool(*flags.is_first_4_rows);
92        builder.assert_bool(*flags.is_digest_row);
93        builder.assert_bool(*flags.is_round_row + *flags.is_digest_row);
94
95        self.row_idx_encoder
96            .eval(builder, local_cols.flags.row_idx.to_slice().unwrap());
97        // assert all row indices are in [0, C::ROWS_PER_BLOCK]
98        builder.assert_one(self.row_idx_encoder.contains_flag_range::<AB>(
99            local_cols.flags.row_idx.to_slice().unwrap(),
100            0..=C::ROWS_PER_BLOCK,
101        ));
102        // assert that the row indices are [0, 3] for the first 4 rows
103        builder.assert_eq(
104            self.row_idx_encoder
105                .contains_flag_range::<AB>(local_cols.flags.row_idx.to_slice().unwrap(), 0..=3),
106            *flags.is_first_4_rows,
107        );
108        // round row indices are in [0, C::ROUND_ROWS - 1]
109        builder.assert_eq(
110            self.row_idx_encoder.contains_flag_range::<AB>(
111                local_cols.flags.row_idx.to_slice().unwrap(),
112                0..=C::ROUND_ROWS - 1,
113            ),
114            *flags.is_round_row,
115        );
116        // digest row always has row index C::ROUND_ROWS
117        builder.assert_eq(
118            self.row_idx_encoder.contains_flag::<AB>(
119                local_cols.flags.row_idx.to_slice().unwrap(),
120                &[C::ROUND_ROWS],
121            ),
122            *flags.is_digest_row,
123        );
124        // If padding row we want the row_idx to be C::ROWS_PER_BLOCK
125        builder.assert_eq(
126            self.row_idx_encoder.contains_flag::<AB>(
127                local_cols.flags.row_idx.to_slice().unwrap(),
128                &[C::ROWS_PER_BLOCK],
129            ),
130            flags.is_padding_row(),
131        );
132
133        // Constrain a, e, being composed of bits: we make sure a and e are always in the same place
134        // in the trace matrix Note: this has to be true for every row, even padding rows
135        for i in 0..C::ROUNDS_PER_ROW {
136            for j in 0..C::WORD_BITS {
137                builder.assert_bool(local_cols.hash.a[[i, j]]);
138                builder.assert_bool(local_cols.hash.e[[i, j]]);
139            }
140        }
141    }
142
143    /// Evaluates the final hash for this block
144    fn eval_digest_row<AB: InteractionBuilder>(
145        &self,
146        builder: &mut AB,
147        local: Sha2RoundColsRef<AB::Var>,
148        next: Sha2DigestColsRef<AB::Var>,
149    ) {
150        // Assert that the previous hash + work vars == final hash.
151        // That is, `next.prev_hash[i] + local.work_vars[i] == next.final_hash[i]`
152        // where addition is done modulo 2^32
153        for i in 0..C::HASH_WORDS {
154            let mut carry = AB::Expr::ZERO;
155            for j in 0..C::WORD_U16S {
156                let work_var_limb = if i < C::ROUNDS_PER_ROW {
157                    compose::<AB::Expr>(
158                        local
159                            .work_vars
160                            .a
161                            .slice(s![C::ROUNDS_PER_ROW - 1 - i, j * 16..(j + 1) * 16])
162                            .as_slice()
163                            .unwrap(),
164                        1,
165                    )
166                } else {
167                    compose::<AB::Expr>(
168                        local
169                            .work_vars
170                            .e
171                            .slice(s![C::ROUNDS_PER_ROW + 3 - i, j * 16..(j + 1) * 16])
172                            .as_slice()
173                            .unwrap(),
174                        1,
175                    )
176                };
177                let final_hash_limb = compose::<AB::Expr>(
178                    next.final_hash
179                        .slice(s![i, j * 2..(j + 1) * 2])
180                        .as_slice()
181                        .unwrap(),
182                    8,
183                );
184
185                carry = AB::Expr::from(AB::F::from_u32(1 << 16).inverse())
186                    * (next.prev_hash[[i, j]] + work_var_limb + carry - final_hash_limb);
187                builder
188                    .when(*next.flags.is_digest_row)
189                    .assert_bool(carry.clone());
190            }
191            // constrain the final hash limbs two at a time since we can do two checks per
192            // interaction
193            for chunk in next.final_hash.row(i).as_slice().unwrap().chunks(2) {
194                self.bitwise_lookup_bus
195                    .send_range(chunk[0], chunk[1])
196                    .eval(builder, *next.flags.is_digest_row);
197            }
198        }
199    }
200
201    fn eval_transitions<AB: InteractionBuilder>(&self, builder: &mut AB, start_col: usize) {
202        let main = builder.main();
203        let local = main.row_slice(0).unwrap();
204        let next = main.row_slice(1).unwrap();
205
206        // Doesn't matter what column structs we use here
207        let local_cols: Sha2RoundColsRef<AB::Var> =
208            Sha2RoundColsRef::from::<C>(&local[start_col..start_col + C::SUBAIR_ROUND_WIDTH]);
209        let next_cols: Sha2RoundColsRef<AB::Var> =
210            Sha2RoundColsRef::from::<C>(&next[start_col..start_col + C::SUBAIR_ROUND_WIDTH]);
211
212        let local_is_padding_row = local_cols.flags.is_padding_row();
213        // Note that there will always be a padding row in the trace since the unpadded height is a
214        // multiple of 17 (SHA-256) or 21 (SHA-512, SHA-384). So the next row is padding iff the
215        // current block is the last block in the trace.
216        let next_is_padding_row = next_cols.flags.is_padding_row();
217
218        // If we are in a round row, the next row cannot be a padding row
219        builder
220            .when(*local_cols.flags.is_round_row)
221            .assert_zero(next_is_padding_row.clone());
222        // The first row must be a round row
223        builder
224            .when_first_row()
225            .assert_one(*local_cols.flags.is_round_row);
226        // Anchor the trace shape at the end so the wrapped evaluation window cannot
227        // cycle back from a round row without ever reaching digest/padding rows.
228        builder
229            .when_last_row()
230            .assert_one(local_is_padding_row.clone());
231        // If we are in a padding row, the next row must also be a padding row
232        builder
233            .when_transition()
234            .when(local_is_padding_row.clone())
235            .assert_one(next_is_padding_row.clone());
236        // If we are in a digest row, the next row cannot be a digest row
237        builder
238            .when(*local_cols.flags.is_digest_row)
239            .assert_zero(*next_cols.flags.is_digest_row);
240        // Constrain how much the row index changes by
241        // round->round: 1
242        // round->digest: 1
243        // digest->round: -C::ROUND_ROWS
244        // digest->padding: 1
245        // padding->padding: 0
246        // Other transitions are not allowed by the above constraints
247        let delta = *local_cols.flags.is_round_row * AB::Expr::ONE
248            + *local_cols.flags.is_digest_row
249                * *next_cols.flags.is_round_row
250                * AB::Expr::from_usize(C::ROUND_ROWS)
251                * AB::Expr::NEG_ONE
252            + *local_cols.flags.is_digest_row * next_is_padding_row.clone() * AB::Expr::ONE;
253
254        let local_row_idx = self.row_idx_encoder.flag_with_val::<AB>(
255            local_cols.flags.row_idx.to_slice().unwrap(),
256            &(0..=C::ROWS_PER_BLOCK).map(|i| (i, i)).collect::<Vec<_>>(),
257        );
258        let next_row_idx = self.row_idx_encoder.flag_with_val::<AB>(
259            next_cols.flags.row_idx.to_slice().unwrap(),
260            &(0..=C::ROWS_PER_BLOCK).map(|i| (i, i)).collect::<Vec<_>>(),
261        );
262
263        builder
264            .when_transition()
265            .assert_eq(local_row_idx.clone() + delta, next_row_idx.clone());
266        builder.when_first_row().assert_zero(local_row_idx);
267
268        // Global block index is 1 on first row
269        builder
270            .when_first_row()
271            .assert_one(*local_cols.flags.global_block_idx);
272
273        // Global block index is constant on all rows in a block
274        builder.when(*local_cols.flags.is_round_row).assert_eq(
275            *local_cols.flags.global_block_idx,
276            *next_cols.flags.global_block_idx,
277        );
278        // Global block index increases by 1 between blocks
279        builder
280            .when_transition()
281            .when(*local_cols.flags.is_digest_row)
282            .assert_eq(
283                *local_cols.flags.global_block_idx + AB::Expr::ONE,
284                *next_cols.flags.global_block_idx,
285            );
286        // Global block index is constant on padding rows
287        builder
288            .when_transition()
289            .when(local_is_padding_row.clone())
290            .assert_eq(
291                *local_cols.flags.global_block_idx,
292                *next_cols.flags.global_block_idx,
293            );
294
295        // Constrain that all the padding rows have the same work vars as the last block's digest
296        // row. We constrain elsewhere that the last block's digest row is equal to the first
297        // block's prev_hash. Together, this ensures that all the padding rows have the same
298        // work vars as the first block's prev_hash. As a result, the
299        // constraint_word_addition constraints in eval_work_vars() on the first row of the first
300        // block (i.e. when next = first_row) will correctly constrain the first 4 rounds of
301        // the first block.
302        for i in 0..C::ROUNDS_PER_ROW {
303            for j in 0..C::WORD_BITS {
304                builder.when(next_cols.flags.is_padding_row()).assert_eq(
305                    local_cols.work_vars.a[[i, j]],
306                    next_cols.work_vars.a[[i, j]],
307                );
308                builder.when(next_cols.flags.is_padding_row()).assert_eq(
309                    local_cols.work_vars.e[[i, j]],
310                    next_cols.work_vars.e[[i, j]],
311                );
312            }
313        }
314
315        self.eval_message_schedule(builder, local_cols.clone(), next_cols.clone());
316        self.eval_work_vars(builder, local_cols.clone(), next_cols);
317        let next: Sha2DigestColsRef<AB::Var> =
318            Sha2DigestColsRef::from::<C>(&next[start_col..start_col + C::SUBAIR_DIGEST_WIDTH]);
319        self.eval_digest_row(builder, local_cols, next);
320        let local_cols: Sha2DigestColsRef<AB::Var> =
321            Sha2DigestColsRef::from::<C>(&local[start_col..start_col + C::SUBAIR_DIGEST_WIDTH]);
322        self.eval_prev_hash(builder, local_cols, next_is_padding_row);
323    }
324
325    /// Constrains that the next block's `prev_hash` is equal to the current block's `hash`
326    /// Note: the constraining is done by interactions with the chip itself on every digest row
327    fn eval_prev_hash<AB: InteractionBuilder>(
328        &self,
329        builder: &mut AB,
330        local: Sha2DigestColsRef<AB::Var>,
331        is_last_block_of_trace: AB::Expr, /* note this indicates the last block of the trace,
332                                           * not the last block of the message */
333    ) {
334        // Constrain that next block's `prev_hash` is equal to the current block's `hash`
335        let composed_hash = (0..C::HASH_WORDS)
336            .map(|i| {
337                let hash_bits = if i < C::ROUNDS_PER_ROW {
338                    local
339                        .hash
340                        .a
341                        .row(C::ROUNDS_PER_ROW - 1 - i)
342                        .mapv(|x| x.into())
343                        .to_vec()
344                } else {
345                    local
346                        .hash
347                        .e
348                        .row(C::ROUNDS_PER_ROW + 3 - i)
349                        .mapv(|x| x.into())
350                        .to_vec()
351                };
352                (0..C::WORD_U16S)
353                    .map(|j| compose::<AB::Expr>(&hash_bits[j * 16..(j + 1) * 16], 1))
354                    .collect::<Vec<_>>()
355            })
356            .collect::<Vec<_>>();
357        // Need to handle the case if this is the very last block of the trace matrix
358        let next_global_block_idx = select(
359            is_last_block_of_trace,
360            AB::Expr::ONE,
361            *local.flags.global_block_idx + AB::Expr::ONE,
362        );
363        // The following interactions constrain certain values from block to block
364        self.private_bus.send(
365            builder,
366            composed_hash
367                .into_iter()
368                .flatten()
369                .chain(once(next_global_block_idx)),
370            *local.flags.is_digest_row,
371        );
372
373        self.private_bus.receive(
374            builder,
375            local
376                .prev_hash
377                .flatten()
378                .mapv(|x| x.into())
379                .into_iter()
380                .chain(once((*local.flags.global_block_idx).into())),
381            *local.flags.is_digest_row,
382        );
383    }
384
385    /// Constrain the message schedule additions for `next` row
386    /// Note: For every addition we need to constrain the following for each of [WORD_U16S] limbs
387    /// sig_1(w_{t-2})[i] + w_{t-7}[i] + sig_0(w_{t-15})[i] + w_{t-16}[i] + carry_w[t][i-1] -
388    /// carry_w[t][i] * 2^16 - w_t[i] == 0 Refer to [https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf]
389    fn eval_message_schedule<'a, AB: InteractionBuilder>(
390        &self,
391        builder: &mut AB,
392        local: Sha2RoundColsRef<'a, AB::Var>,
393        next: Sha2RoundColsRef<'a, AB::Var>,
394    ) {
395        // This `w` array contains 8 message schedule words - w_{idx}, ..., w_{idx+7} for some idx
396        let w = ndarray::concatenate(
397            ndarray::Axis(0),
398            &[local.message_schedule.w, next.message_schedule.w],
399        )
400        .unwrap();
401
402        // Constrain `w_3` for `next` row
403        for i in 0..C::ROUNDS_PER_ROW - 1 {
404            // here we constrain the w_3 of the i_th word of the next row
405            // w_3 of next is w[i+4-3] = w[i+1]
406            let w_3 = w.row(i + 1).mapv(|x| x.into()).to_vec();
407            let expected_w_3 = next.schedule_helper.w_3.row(i);
408            for j in 0..C::WORD_U16S {
409                let w_3_limb = compose::<AB::Expr>(&w_3[j * 16..(j + 1) * 16], 1);
410                builder
411                    .when(*local.flags.is_round_row)
412                    .assert_eq(w_3_limb, expected_w_3[j].into());
413            }
414        }
415
416        // Constrain intermed for `next` row
417        // We will only constrain intermed_12 for rows [3, C::ROUND_ROWS - 2], and let it
418        // unconstrained for other rows Other rows should put the needed value in
419        // intermed_12 to make the below summation constraint hold
420        let is_row_intermed_12 = self.row_idx_encoder.contains_flag_range::<AB>(
421            next.flags.row_idx.to_slice().unwrap(),
422            3..=C::ROUND_ROWS - 2,
423        );
424        // We will only constrain intermed_8 for rows [2, C::ROUND_ROWS - 3], and let it
425        // unconstrained for other rows
426        let is_row_intermed_8 = self.row_idx_encoder.contains_flag_range::<AB>(
427            next.flags.row_idx.to_slice().unwrap(),
428            2..=C::ROUND_ROWS - 3,
429        );
430        for i in 0..C::ROUNDS_PER_ROW {
431            // w_idx
432            let w_idx = w.row(i).mapv(|x| x.into()).to_vec();
433            // sig_0(w_{idx+1})
434            let sig_w = small_sig0_field::<AB::Expr, C>(w.row(i + 1).as_slice().unwrap());
435            for j in 0..C::WORD_U16S {
436                let w_idx_limb = compose::<AB::Expr>(&w_idx[j * 16..(j + 1) * 16], 1);
437                let sig_w_limb = compose::<AB::Expr>(&sig_w[j * 16..(j + 1) * 16], 1);
438
439                // We would like to constrain this only on round rows, but we can't do a conditional
440                // check because the degree is already 3. So we must fill in `intermed_4` with dummy
441                // values on the first round row and the digest row (rows 0 and 16 for SHA-256) to
442                // ensure the constraint holds on these rows.
443                builder.assert_eq(
444                    next.schedule_helper.intermed_4[[i, j]],
445                    w_idx_limb + sig_w_limb,
446                );
447
448                builder.when(is_row_intermed_8.clone()).assert_eq(
449                    next.schedule_helper.intermed_8[[i, j]],
450                    local.schedule_helper.intermed_4[[i, j]],
451                );
452
453                builder.when(is_row_intermed_12.clone()).assert_eq(
454                    next.schedule_helper.intermed_12[[i, j]],
455                    local.schedule_helper.intermed_8[[i, j]],
456                );
457            }
458        }
459
460        // Constrain the message schedule additions for `next` row
461        for i in 0..C::ROUNDS_PER_ROW {
462            // Note, here by w_{t} we mean the i_th word of the `next` row
463            // w_{t-7}
464            let w_7 = if i < 3 {
465                local.schedule_helper.w_3.row(i).mapv(|x| x.into()).to_vec()
466            } else {
467                let w_3 = w.row(i - 3).mapv(|x| x.into()).to_vec();
468                (0..C::WORD_U16S)
469                    .map(|j| compose::<AB::Expr>(&w_3[j * 16..(j + 1) * 16], 1))
470                    .collect::<Vec<_>>()
471            };
472            // sig_0(w_{t-15}) + w_{t-16}
473            let intermed_16 = local.schedule_helper.intermed_12.row(i).mapv(|x| x.into());
474
475            let carries = (0..C::WORD_U16S)
476                .map(|j| {
477                    next.message_schedule.carry_or_buffer[[i, j * 2]]
478                        + AB::Expr::TWO * next.message_schedule.carry_or_buffer[[i, j * 2 + 1]]
479                })
480                .collect::<Vec<_>>();
481
482            // Constrain `W_{idx} = sig_1(W_{idx-2}) + W_{idx-7} + sig_0(W_{idx-15}) + W_{idx-16}`
483            // We can't do a conditional check because the degree of the sum is already 3.
484            // Instead, we fill in dummy values on non-applicable rows:
485            //   - `intermed_12` with dummy values on rows 0..3, C::ROUND_ROWS-1, and C::ROUND_ROWS
486            //   - `carry_or_buffer` with dummy values on the first block's row 0 (wrap-around)
487            constraint_word_addition::<_, C>(
488                builder,
489                &[&small_sig1_field::<AB::Expr, C>(
490                    w.row(i + 2).as_slice().unwrap(),
491                )],
492                &[&w_7, intermed_16.as_slice().unwrap()],
493                w.row(i + 4).as_slice().unwrap(),
494                &carries,
495            );
496
497            for j in 0..C::WORD_U16S {
498                // When on rows 4..C::ROUND_ROWS message schedule carries should be 0 or 1
499                let is_row_4_or_more = *next.flags.is_round_row - *next.flags.is_first_4_rows;
500                builder
501                    .when(is_row_4_or_more.clone())
502                    .assert_bool(next.message_schedule.carry_or_buffer[[i, j * 2]]);
503                builder
504                    .when(is_row_4_or_more)
505                    .assert_bool(next.message_schedule.carry_or_buffer[[i, j * 2 + 1]]);
506            }
507            // Constrain w being composed of bits
508            for j in 0..C::WORD_BITS {
509                builder
510                    .when(*next.flags.is_round_row)
511                    .assert_bool(next.message_schedule.w[[i, j]]);
512            }
513        }
514    }
515
516    /// Constrain the work vars on `next` row according to the sha documentation
517    /// Refer to [https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf]
518    fn eval_work_vars<'a, AB: InteractionBuilder>(
519        &self,
520        builder: &mut AB,
521        local: Sha2RoundColsRef<'a, AB::Var>,
522        next: Sha2RoundColsRef<'a, AB::Var>,
523    ) {
524        let a =
525            ndarray::concatenate(ndarray::Axis(0), &[local.work_vars.a, next.work_vars.a]).unwrap();
526        let e =
527            ndarray::concatenate(ndarray::Axis(0), &[local.work_vars.e, next.work_vars.e]).unwrap();
528
529        for i in 0..C::ROUNDS_PER_ROW {
530            for j in 0..C::WORD_U16S {
531                // Although we need carry_a <= 6 and carry_e <= 5, constraining carry_a, carry_e in
532                // [0, 2^8) is enough to prevent overflow and ensure the soundness
533                // of the addition we want to check
534                self.bitwise_lookup_bus
535                    .send_range(
536                        local.work_vars.carry_a[[i, j]],
537                        local.work_vars.carry_e[[i, j]],
538                    )
539                    .eval(builder, *local.flags.is_round_row);
540            }
541
542            let w_limbs = (0..C::WORD_U16S)
543                .map(|j| {
544                    compose::<AB::Expr>(
545                        next.message_schedule
546                            .w
547                            .slice(s![i, j * 16..(j + 1) * 16])
548                            .as_slice()
549                            .unwrap(),
550                        1,
551                    ) * *next.flags.is_round_row
552                })
553                .collect::<Vec<_>>();
554
555            let k_limbs = (0..C::WORD_U16S)
556                .map(|j| {
557                    self.row_idx_encoder.flag_with_val::<AB>(
558                        next.flags.row_idx.to_slice().unwrap(),
559                        &(0..C::ROUND_ROWS)
560                            .map(|rw_idx| {
561                                (
562                                    rw_idx,
563                                    word_into_u16_limbs::<C>(
564                                        C::get_k()[rw_idx * C::ROUNDS_PER_ROW + i],
565                                    )[j] as usize,
566                                )
567                            })
568                            .collect::<Vec<_>>(),
569                    )
570                })
571                .collect::<Vec<_>>();
572
573            // Constrain `a = h + sig_1(e) + ch(e, f, g) + K + W + sig_0(a) + Maj(a, b, c)`
574            // We have to enforce this constraint on all rows since the degree of the constraint is
575            // already 3. So, we must fill in `carry_a` with dummy values on digest rows
576            // to ensure the constraint holds.
577            constraint_word_addition::<_, C>(
578                builder,
579                &[
580                    e.row(i).mapv(|x| x.into()).as_slice().unwrap(), // previous `h`
581                    &big_sig1_field::<AB::Expr, C>(e.row(i + 3).as_slice().unwrap()), /* sig_1 of
582                                                                     previous `e` */
583                    &ch_field::<AB::Expr>(
584                        e.row(i + 3).as_slice().unwrap(),
585                        e.row(i + 2).as_slice().unwrap(),
586                        e.row(i + 1).as_slice().unwrap(),
587                    ), /* Ch of previous `e`, `f`, `g` */
588                    &big_sig0_field::<AB::Expr, C>(a.row(i + 3).as_slice().unwrap()), /* sig_0 of `a` */
589                    &maj_field::<AB::Expr>(
590                        a.row(i + 3).as_slice().unwrap(),
591                        a.row(i + 2).as_slice().unwrap(),
592                        a.row(i + 1).as_slice().unwrap(),
593                    ), /* Maj of previous a, b, c */
594                ],
595                &[&w_limbs, &k_limbs],                             // K and W
596                a.row(i + 4).as_slice().unwrap(),                  // new `a`
597                next.work_vars.carry_a.row(i).as_slice().unwrap(), // carries of addition
598            );
599
600            // Constrain `e = d + h + sig_1(e) + ch(e, f, g) + K + W`
601            // We have to enforce this constraint on all rows since the degree of the constraint is
602            // already 3. So, we must fill in `carry_e` with dummy values on digest rows
603            // to ensure the constraint holds.
604            constraint_word_addition::<_, C>(
605                builder,
606                &[
607                    a.row(i).mapv(|x| x.into()).as_slice().unwrap(), // previous `d`
608                    e.row(i).mapv(|x| x.into()).as_slice().unwrap(), // previous `h`
609                    &big_sig1_field::<AB::Expr, C>(e.row(i + 3).as_slice().unwrap()), /* sig_1 of
610                                                                     previous `e` */
611                    &ch_field::<AB::Expr>(
612                        e.row(i + 3).as_slice().unwrap(),
613                        e.row(i + 2).as_slice().unwrap(),
614                        e.row(i + 1).as_slice().unwrap(),
615                    ), /* Ch of previous `e`, `f`, `g` */
616                ],
617                &[&w_limbs, &k_limbs],                             // K and W
618                e.row(i + 4).as_slice().unwrap(),                  // new `e`
619                next.work_vars.carry_e.row(i).as_slice().unwrap(), // carries of addition
620            );
621        }
622    }
623}
openvm_sha2_air/air.rs

openvm_sha2_air/
air.rs
