openvm_sha2_air/
trace.rs

1use std::{marker::PhantomData, ops::Range};
2
3use openvm_circuit_primitives::{
4    bitwise_op_lookup::SharedBitwiseOperationLookupChip, encoder::Encoder, utils::compose,
5};
6use openvm_stark_backend::p3_field::PrimeField32;
7use sha2::{compress256, compress512, digest::generic_array::GenericArray};
8
9use crate::{
10    big_sig0, big_sig0_field, big_sig1, big_sig1_field, ch, ch_field, get_flag_pt_array,
11    le_limbs_into_word, maj, maj_field, set_arrayview_from_u32_slice, small_sig0, small_sig0_field,
12    small_sig1, small_sig1_field, word_into_bits, word_into_u16_limbs, word_into_u8_limbs,
13    Sha2BlockHasherSubairConfig, Sha2DigestColsRefMut, Sha2RoundColsRef, Sha2RoundColsRefMut,
14    Sha2Variant, WrappingAdd,
15};
16
17/// A helper struct for the SHA-2 trace generation.
18/// Also, separates the inner AIR from the trace generation.
19pub struct Sha2BlockHasherFillerHelper<C: Sha2BlockHasherSubairConfig> {
20    pub row_idx_encoder: Encoder,
21    _phantom: PhantomData<C>,
22}
23
24impl<C: Sha2BlockHasherSubairConfig> Default for Sha2BlockHasherFillerHelper<C> {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30/// The trace generation of SHA-2 should be done in two passes.
31/// The first pass should do `get_block_trace` for every block and generate the invalid rows through
32/// `get_default_row` The second pass should go through all the blocks and call
33/// `generate_missing_cells`
34impl<C: Sha2BlockHasherSubairConfig> Sha2BlockHasherFillerHelper<C> {
35    pub fn new() -> Self {
36        Self {
37            row_idx_encoder: Encoder::new(C::ROWS_PER_BLOCK + 1, 2, false),
38            _phantom: PhantomData,
39        }
40    }
41
42    /// This function takes the input_message (padding not handled), the previous hash,
43    /// and returns the new hash after processing the block input
44    pub fn get_block_hash(prev_hash: &[C::Word], input: Vec<u8>) -> Vec<C::Word> {
45        debug_assert!(prev_hash.len() == C::HASH_WORDS);
46        debug_assert!(input.len() == C::BLOCK_U8S);
47        let mut new_hash: [C::Word; 8] = prev_hash.try_into().unwrap();
48        match C::VARIANT {
49            Sha2Variant::Sha256 => {
50                let input_array = [*GenericArray::<u8, sha2::digest::consts::U64>::from_slice(
51                    &input,
52                )];
53                let hash_ptr: &mut [u32; 8] = unsafe { std::mem::transmute(&mut new_hash) };
54                compress256(hash_ptr, &input_array);
55            }
56            Sha2Variant::Sha512 | Sha2Variant::Sha384 => {
57                let hash_ptr: &mut [u64; 8] = unsafe { std::mem::transmute(&mut new_hash) };
58                let input_array = [*GenericArray::<u8, sha2::digest::consts::U128>::from_slice(
59                    &input,
60                )];
61                compress512(hash_ptr, &input_array);
62            }
63        }
64        new_hash.to_vec()
65    }
66
67    /// This function takes a C::BLOCK_BITS-bit chunk of the input message (padding not handled),
68    /// the previous hash, a flag indicating if it's the last block, the global block index, the
69    /// local block index, and the buffer values that will be put in rows 0..4.
70    /// Will populate the given `trace` with the trace of the block, where the width of the trace is
71    /// `trace_width` and the starting column for the `Sha2Air` is `trace_start_col`.
72    /// **Note**: this function only generates some of the required trace. Another pass is required,
73    /// refer to [`Self::generate_missing_cells`] for details.
74    #[allow(clippy::too_many_arguments)]
75    pub fn generate_block_trace<F: PrimeField32>(
76        &self,
77        trace: &mut [F],
78        trace_width: usize,
79        trace_start_col: usize,
80        input: &[C::Word],
81        bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
82        prev_hash: &[C::Word],
83        next_block_prev_hash: &[C::Word],
84        global_block_idx: u32,
85    ) {
86        #[cfg(debug_assertions)]
87        {
88            assert!(input.len() == C::BLOCK_WORDS);
89            assert!(prev_hash.len() == C::HASH_WORDS);
90            assert!(next_block_prev_hash.len() == C::HASH_WORDS);
91            assert!(trace_start_col + C::SUBAIR_WIDTH == trace_width);
92            assert!(trace.len() == trace_width * C::ROWS_PER_BLOCK);
93        }
94
95        let get_range = |start: usize, len: usize| -> Range<usize> { start..start + len };
96        let mut message_schedule = vec![C::Word::from(0); C::ROUNDS_PER_BLOCK];
97        message_schedule[..input.len()].copy_from_slice(input);
98        let mut work_vars = prev_hash.to_vec();
99        for (i, row) in trace.chunks_exact_mut(trace_width).enumerate() {
100            // do the rounds
101            if i < C::ROUND_ROWS {
102                let mut cols: Sha2RoundColsRefMut<F> = Sha2RoundColsRefMut::from::<C>(
103                    &mut row[get_range(trace_start_col, C::SUBAIR_ROUND_WIDTH)],
104                );
105                *cols.flags.is_round_row = F::ONE;
106                *cols.flags.is_first_4_rows = if i < C::MESSAGE_ROWS { F::ONE } else { F::ZERO };
107                *cols.flags.is_digest_row = F::ZERO;
108                set_arrayview_from_u32_slice(
109                    &mut cols.flags.row_idx,
110                    get_flag_pt_array(&self.row_idx_encoder, i),
111                );
112                *cols.flags.global_block_idx = F::from_u32(global_block_idx);
113
114                // W_idx = M_idx
115                if i < C::MESSAGE_ROWS {
116                    for j in 0..C::ROUNDS_PER_ROW {
117                        set_arrayview_from_u32_slice(
118                            &mut cols.message_schedule.w.row_mut(j),
119                            word_into_bits::<C>(input[i * C::ROUNDS_PER_ROW + j]),
120                        );
121                    }
122                }
123                // W_idx = SIG1(W_{idx-2}) + W_{idx-7} + SIG0(W_{idx-15}) + W_{idx-16}
124                else {
125                    for j in 0..C::ROUNDS_PER_ROW {
126                        let idx = i * C::ROUNDS_PER_ROW + j;
127                        let nums: [C::Word; 4] = [
128                            small_sig1::<C>(message_schedule[idx - 2]),
129                            message_schedule[idx - 7],
130                            small_sig0::<C>(message_schedule[idx - 15]),
131                            message_schedule[idx - 16],
132                        ];
133                        let w: C::Word = nums
134                            .iter()
135                            .fold(C::Word::from(0), |acc, &num| acc.wrapping_add(num));
136                        set_arrayview_from_u32_slice(
137                            &mut cols.message_schedule.w.row_mut(j),
138                            word_into_bits::<C>(w),
139                        );
140                        let nums_limbs = nums
141                            .iter()
142                            .map(|x| word_into_u16_limbs::<C>(*x))
143                            .collect::<Vec<_>>();
144                        let w_limbs = word_into_u16_limbs::<C>(w);
145
146                        // fill in the carrys
147                        for k in 0..C::WORD_U16S {
148                            let mut sum = nums_limbs.iter().fold(0, |acc, num| acc + num[k]);
149                            if k > 0 {
150                                sum += (cols.message_schedule.carry_or_buffer[[j, k * 2 - 2]]
151                                    + F::TWO
152                                        * cols.message_schedule.carry_or_buffer[[j, k * 2 - 1]])
153                                .as_canonical_u32();
154                            }
155                            let carry = (sum - w_limbs[k]) >> 16;
156                            cols.message_schedule.carry_or_buffer[[j, k * 2]] =
157                                F::from_u32(carry & 1);
158                            cols.message_schedule.carry_or_buffer[[j, k * 2 + 1]] =
159                                F::from_u32(carry >> 1);
160                        }
161                        // update the message schedule
162                        message_schedule[idx] = w;
163                    }
164                }
165                // fill in the work variables
166                for j in 0..C::ROUNDS_PER_ROW {
167                    // t1 = h + SIG1(e) + ch(e, f, g) + K_idx + W_idx
168                    let t1 = [
169                        work_vars[7],
170                        big_sig1::<C>(work_vars[4]),
171                        ch::<C>(work_vars[4], work_vars[5], work_vars[6]),
172                        C::get_k()[i * C::ROUNDS_PER_ROW + j],
173                        le_limbs_into_word::<C>(
174                            cols.message_schedule
175                                .w
176                                .row(j)
177                                .map(|f| f.as_canonical_u32())
178                                .as_slice()
179                                .unwrap(),
180                        ),
181                    ];
182                    let t1_sum: C::Word = t1
183                        .iter()
184                        .fold(C::Word::from(0), |acc, &num| acc.wrapping_add(num));
185
186                    // t2 = SIG0(a) + maj(a, b, c)
187                    let t2 = [
188                        big_sig0::<C>(work_vars[0]),
189                        maj::<C>(work_vars[0], work_vars[1], work_vars[2]),
190                    ];
191
192                    let t2_sum: C::Word = t2
193                        .iter()
194                        .fold(C::Word::from(0), |acc, &num| acc.wrapping_add(num));
195
196                    // e = d + t1
197                    let e = work_vars[3].wrapping_add(t1_sum);
198                    set_arrayview_from_u32_slice(
199                        &mut cols.work_vars.e.row_mut(j),
200                        word_into_bits::<C>(e),
201                    );
202                    let e_limbs = word_into_u16_limbs::<C>(e);
203                    // a = t1 + t2
204                    let a = t1_sum.wrapping_add(t2_sum);
205                    set_arrayview_from_u32_slice(
206                        &mut cols.work_vars.a.row_mut(j),
207                        word_into_bits::<C>(a),
208                    );
209                    let a_limbs = word_into_u16_limbs::<C>(a);
210                    // fill in the carrys
211                    for k in 0..C::WORD_U16S {
212                        let t1_limb = t1
213                            .iter()
214                            .fold(0, |acc, &num| acc + word_into_u16_limbs::<C>(num)[k]);
215                        let t2_limb = t2
216                            .iter()
217                            .fold(0, |acc, &num| acc + word_into_u16_limbs::<C>(num)[k]);
218
219                        let mut e_limb = t1_limb + word_into_u16_limbs::<C>(work_vars[3])[k];
220                        let mut a_limb = t1_limb + t2_limb;
221                        if k > 0 {
222                            a_limb += cols.work_vars.carry_a[[j, k - 1]].as_canonical_u32();
223                            e_limb += cols.work_vars.carry_e[[j, k - 1]].as_canonical_u32();
224                        }
225                        let carry_a = (a_limb - a_limbs[k]) >> 16;
226                        let carry_e = (e_limb - e_limbs[k]) >> 16;
227                        cols.work_vars.carry_a[[j, k]] = F::from_u32(carry_a);
228                        cols.work_vars.carry_e[[j, k]] = F::from_u32(carry_e);
229                        bitwise_lookup_chip.request_range(carry_a, carry_e);
230                    }
231
232                    // update working variables
233                    work_vars[7] = work_vars[6];
234                    work_vars[6] = work_vars[5];
235                    work_vars[5] = work_vars[4];
236                    work_vars[4] = e;
237                    work_vars[3] = work_vars[2];
238                    work_vars[2] = work_vars[1];
239                    work_vars[1] = work_vars[0];
240                    work_vars[0] = a;
241                }
242
243                // filling w_3 and intermed_4 here and the rest later
244                if i > 0 {
245                    for j in 0..C::ROUNDS_PER_ROW {
246                        let idx = i * C::ROUNDS_PER_ROW + j;
247                        let w_4 = word_into_u16_limbs::<C>(message_schedule[idx - 4]);
248                        let sig_0_w_3 =
249                            word_into_u16_limbs::<C>(small_sig0::<C>(message_schedule[idx - 3]));
250                        set_arrayview_from_u32_slice(
251                            &mut cols.schedule_helper.intermed_4.row_mut(j),
252                            (0..C::WORD_U16S)
253                                .map(|k| w_4[k] + sig_0_w_3[k])
254                                .collect::<Vec<_>>(),
255                        );
256                        if j < C::ROUNDS_PER_ROW - 1 {
257                            let w_3 = message_schedule[idx - 3];
258                            set_arrayview_from_u32_slice(
259                                &mut cols.schedule_helper.w_3.row_mut(j),
260                                word_into_u16_limbs::<C>(w_3),
261                            );
262                        }
263                    }
264                }
265            }
266            // generate the digest row
267            else {
268                let mut cols: Sha2DigestColsRefMut<F> = Sha2DigestColsRefMut::from::<C>(
269                    &mut row[get_range(trace_start_col, C::SUBAIR_DIGEST_WIDTH)],
270                );
271                for j in 0..C::ROUNDS_PER_ROW - 1 {
272                    let w_3 = message_schedule[i * C::ROUNDS_PER_ROW + j - 3];
273                    set_arrayview_from_u32_slice(
274                        &mut cols.schedule_helper.w_3.row_mut(j),
275                        word_into_u16_limbs::<C>(w_3),
276                    );
277                }
278                *cols.flags.is_round_row = F::ZERO;
279                *cols.flags.is_first_4_rows = F::ZERO;
280                *cols.flags.is_digest_row = F::ONE;
281                set_arrayview_from_u32_slice(
282                    &mut cols.flags.row_idx,
283                    get_flag_pt_array(&self.row_idx_encoder, C::ROUND_ROWS),
284                );
285                *cols.flags.global_block_idx = F::from_u32(global_block_idx);
286
287                let final_hash: Vec<C::Word> = (0..C::HASH_WORDS)
288                    .map(|i| work_vars[i].wrapping_add(prev_hash[i]))
289                    .collect();
290                let final_hash_limbs: Vec<Vec<u32>> = final_hash
291                    .iter()
292                    .map(|word| word_into_u8_limbs::<C>(*word))
293                    .collect();
294                // need to ensure final hash limbs are bytes, in order for
295                //   prev_hash[i] + work_vars[i] == final_hash[i]
296                // to be constrained correctly
297                for word in final_hash_limbs.iter() {
298                    for chunk in word.chunks(2) {
299                        bitwise_lookup_chip.request_range(chunk[0], chunk[1]);
300                    }
301                }
302                set_arrayview_from_u32_slice(
303                    &mut cols.final_hash,
304                    final_hash
305                        .iter()
306                        .flat_map(|word| word_into_u8_limbs::<C>(*word)),
307                );
308                set_arrayview_from_u32_slice(
309                    &mut cols.prev_hash,
310                    prev_hash
311                        .iter()
312                        .flat_map(|word| word_into_u16_limbs::<C>(*word)),
313                );
314                let next_block_prev_hash_bits = next_block_prev_hash
315                    .iter()
316                    .map(|x| word_into_bits::<C>(*x))
317                    .collect::<Vec<_>>();
318
319                for i in 0..C::ROUNDS_PER_ROW {
320                    set_arrayview_from_u32_slice(
321                        &mut cols.hash.a.row_mut(i),
322                        next_block_prev_hash_bits[C::ROUNDS_PER_ROW - i - 1].clone(),
323                    );
324                    set_arrayview_from_u32_slice(
325                        &mut cols.hash.e.row_mut(i),
326                        next_block_prev_hash_bits[C::ROUNDS_PER_ROW - i + 3].clone(),
327                    );
328                }
329            }
330        }
331
332        for i in 0..C::ROWS_PER_BLOCK - 1 {
333            let rows = &mut trace[i * trace_width..(i + 2) * trace_width];
334            let (local, next) = rows.split_at_mut(trace_width);
335            let mut local_cols: Sha2RoundColsRefMut<F> = Sha2RoundColsRefMut::from::<C>(
336                &mut local[get_range(trace_start_col, C::SUBAIR_ROUND_WIDTH)],
337            );
338            let mut next_cols: Sha2RoundColsRefMut<F> = Sha2RoundColsRefMut::from::<C>(
339                &mut next[get_range(trace_start_col, C::SUBAIR_ROUND_WIDTH)],
340            );
341            if i > 0 {
342                for j in 0..C::ROUNDS_PER_ROW {
343                    next_cols
344                        .schedule_helper
345                        .intermed_8
346                        .row_mut(j)
347                        .assign(&local_cols.schedule_helper.intermed_4.row(j));
348                    if (2..C::ROWS_PER_BLOCK - 3).contains(&i) {
349                        next_cols
350                            .schedule_helper
351                            .intermed_12
352                            .row_mut(j)
353                            .assign(&local_cols.schedule_helper.intermed_8.row(j));
354                    }
355                }
356            }
357            if i == C::ROWS_PER_BLOCK - 2 {
358                // `next` is a digest row.
359                // Fill in `carry_a` and `carry_e` with dummy values so the constraints on `a` and
360                // `e` hold.
361                let const_local_cols = Sha2RoundColsRef::<F>::from_mut::<C>(&local_cols);
362                Self::generate_carry_ae(const_local_cols.clone(), &mut next_cols);
363
364                // Fill in row 16's `intermed_4` with dummy values so the message schedule
365                // constraints holds on that row
366                Self::generate_intermed_4(const_local_cols, &mut next_cols);
367            }
368            if i < C::MESSAGE_ROWS - 1 {
369                // i is in 0..3.
370                // Fill in `local.intermed_12` with dummy values so the message schedule constraints
371                // hold on rows 1..4.
372                Self::generate_intermed_12(
373                    &mut local_cols,
374                    Sha2RoundColsRef::<F>::from_mut::<C>(&next_cols),
375                );
376            }
377        }
378    }
379
380    /// This function will fill in the cells that we couldn't do during the first pass.
381    /// This function should be called only after `generate_block_trace` was called for all blocks
382    /// And [`Self::generate_default_row`] is called for all invalid rows
383    /// Will populate the missing values of `trace`, where the width of the trace is `trace_width`
384    /// Note: `trace` needs to be the rows 1..C::ROWS_PER_BLOCK of a block and the first row of the
385    /// next block
386    pub fn generate_missing_cells<F: PrimeField32>(
387        &self,
388        trace: &mut [F],
389        trace_width: usize,
390        trace_start_col: usize,
391    ) {
392        let rows = &mut trace[(C::ROUND_ROWS - 2) * trace_width..(C::ROUND_ROWS + 1) * trace_width];
393        let (last_round_row, rows) = rows.split_at_mut(trace_width);
394        let (digest_row, next_block_first_row) = rows.split_at_mut(trace_width);
395        let mut cols_last_round_row: Sha2RoundColsRefMut<F> = Sha2RoundColsRefMut::from::<C>(
396            &mut last_round_row[trace_start_col..trace_start_col + C::SUBAIR_ROUND_WIDTH],
397        );
398        let mut cols_digest_row: Sha2RoundColsRefMut<F> = Sha2RoundColsRefMut::from::<C>(
399            &mut digest_row[trace_start_col..trace_start_col + C::SUBAIR_ROUND_WIDTH],
400        );
401        let mut cols_next_block_first_row: Sha2RoundColsRefMut<F> = Sha2RoundColsRefMut::from::<C>(
402            &mut next_block_first_row[trace_start_col..trace_start_col + C::SUBAIR_ROUND_WIDTH],
403        );
404        // Fill in the last round row's `intermed_12` with dummy values so the message schedule
405        // constraints holds on the last round row
406        Self::generate_intermed_12(
407            &mut cols_last_round_row,
408            Sha2RoundColsRef::from_mut::<C>(&cols_digest_row),
409        );
410        // Fill in the digest row's `intermed_12` with dummy values so the message schedule
411        // constraints holds on the next block's row 0
412        Self::generate_intermed_12(
413            &mut cols_digest_row,
414            Sha2RoundColsRef::from_mut::<C>(&cols_next_block_first_row),
415        );
416        // Fill in the next block's first row's `intermed_4` with dummy values so the message
417        // schedule constraints holds on that row
418        Self::generate_intermed_4(
419            Sha2RoundColsRef::from_mut::<C>(&cols_digest_row),
420            &mut cols_next_block_first_row,
421        );
422    }
423
424    /// Fills the `cols` as a padding row
425    /// Note: we still need to correctly fill in the hash values, carries and intermeds
426    pub fn generate_default_row<F: PrimeField32>(
427        &self,
428        cols: &mut Sha2RoundColsRefMut<F>,
429        first_block_prev_hash: &[C::Word],
430        carry_a: Option<&[F]>,
431        carry_e: Option<&[F]>,
432        global_block_idx: u32,
433    ) {
434        debug_assert!(first_block_prev_hash.len() == C::HASH_WORDS);
435        debug_assert!(carry_a.is_some() == carry_e.is_some());
436        debug_assert!(
437            carry_a.is_none() || carry_a.unwrap().len() == C::ROUNDS_PER_ROW * C::WORD_U16S
438        );
439        debug_assert!(
440            carry_e.is_none() || carry_e.unwrap().len() == C::ROUNDS_PER_ROW * C::WORD_U16S
441        );
442
443        set_arrayview_from_u32_slice(
444            &mut cols.flags.row_idx,
445            get_flag_pt_array(&self.row_idx_encoder, C::ROWS_PER_BLOCK),
446        );
447        *cols.flags.global_block_idx = F::from_u32(global_block_idx);
448
449        for i in 0..C::ROUNDS_PER_ROW {
450            // The padding rows need to have the first block's prev_hash here, to satisfy the air
451            // constraints
452            set_arrayview_from_u32_slice(
453                &mut cols.work_vars.a.row_mut(i),
454                word_into_bits::<C>(first_block_prev_hash[C::ROUNDS_PER_ROW - i - 1]).into_iter(),
455            );
456            set_arrayview_from_u32_slice(
457                &mut cols.work_vars.e.row_mut(i),
458                word_into_bits::<C>(first_block_prev_hash[C::ROUNDS_PER_ROW - i + 3]).into_iter(),
459            );
460
461            // The invalid carries are not constants anymore, so we need to fill them in here
462            if let Some(carry_a) = carry_a {
463                cols.work_vars
464                    .carry_a
465                    .iter_mut()
466                    .zip(carry_a.iter())
467                    .for_each(|(x, y)| *x = *y);
468            }
469            if let Some(carry_e) = carry_e {
470                cols.work_vars
471                    .carry_e
472                    .iter_mut()
473                    .zip(carry_e.iter())
474                    .for_each(|(x, y)| *x = *y);
475            }
476        }
477    }
478
479    /// The following functions do the calculations in native field since they will be called on
480    /// padding rows which can overflow and we need to make sure it matches the AIR constraints
481    /// Puts the correct carries in the `next_row`, the resulting carries can be out of bounds.
482    /// Assumes next.w and next.k are zero, which is the case when constraint_word_addition is
483    /// constrained on digest rows or padding rows.
484    /// It only looks at local.a, next.a, local.e, next.e.
485    pub fn generate_carry_ae<F: PrimeField32>(
486        local_cols: Sha2RoundColsRef<F>,
487        next_cols: &mut Sha2RoundColsRefMut<F>,
488    ) {
489        let a = [
490            local_cols
491                .work_vars
492                .a
493                .rows()
494                .into_iter()
495                .collect::<Vec<_>>(),
496            next_cols.work_vars.a.rows().into_iter().collect::<Vec<_>>(),
497        ]
498        .concat();
499        let e = [
500            local_cols
501                .work_vars
502                .e
503                .rows()
504                .into_iter()
505                .collect::<Vec<_>>(),
506            next_cols.work_vars.e.rows().into_iter().collect::<Vec<_>>(),
507        ]
508        .concat();
509        for i in 0..C::ROUNDS_PER_ROW {
510            let cur_a = a[i + 4];
511            let sig_a = big_sig0_field::<F, C>(a[i + 3].as_slice().unwrap());
512            let maj_abc = maj_field::<F>(
513                a[i + 3].as_slice().unwrap(),
514                a[i + 2].as_slice().unwrap(),
515                a[i + 1].as_slice().unwrap(),
516            );
517            let d = a[i];
518            let cur_e = e[i + 4];
519            let sig_e = big_sig1_field::<F, C>(e[i + 3].as_slice().unwrap());
520            let ch_efg = ch_field::<F>(
521                e[i + 3].as_slice().unwrap(),
522                e[i + 2].as_slice().unwrap(),
523                e[i + 1].as_slice().unwrap(),
524            );
525            let h = e[i];
526
527            let t1 = [h.to_vec(), sig_e, ch_efg.to_vec()];
528            let t2 = [sig_a, maj_abc];
529            for j in 0..C::WORD_U16S {
530                let t1_limb_sum = t1.iter().fold(F::ZERO, |acc, x| {
531                    acc + compose::<F>(&x[j * 16..(j + 1) * 16], 1)
532                });
533                let t2_limb_sum = t2.iter().fold(F::ZERO, |acc, x| {
534                    acc + compose::<F>(&x[j * 16..(j + 1) * 16], 1)
535                });
536                let d_limb = compose::<F>(&d.as_slice().unwrap()[j * 16..(j + 1) * 16], 1);
537                let cur_a_limb = compose::<F>(&cur_a.as_slice().unwrap()[j * 16..(j + 1) * 16], 1);
538                let cur_e_limb = compose::<F>(&cur_e.as_slice().unwrap()[j * 16..(j + 1) * 16], 1);
539                let sum = d_limb
540                    + t1_limb_sum
541                    + if j == 0 {
542                        F::ZERO
543                    } else {
544                        next_cols.work_vars.carry_e[[i, j - 1]]
545                    }
546                    - cur_e_limb;
547                let carry_e = sum * (F::from_u32(1 << 16).inverse());
548
549                let sum = t1_limb_sum
550                    + t2_limb_sum
551                    + if j == 0 {
552                        F::ZERO
553                    } else {
554                        next_cols.work_vars.carry_a[[i, j - 1]]
555                    }
556                    - cur_a_limb;
557                let carry_a = sum * (F::from_u32(1 << 16).inverse());
558                next_cols.work_vars.carry_e[[i, j]] = carry_e;
559                next_cols.work_vars.carry_a[[i, j]] = carry_a;
560            }
561        }
562    }
563
564    /// Puts the correct intermed_4 in the `next_row`
565    pub fn generate_intermed_4<F: PrimeField32>(
566        local_cols: Sha2RoundColsRef<F>,
567        next_cols: &mut Sha2RoundColsRefMut<F>,
568    ) {
569        let w = [
570            local_cols
571                .message_schedule
572                .w
573                .rows()
574                .into_iter()
575                .collect::<Vec<_>>(),
576            next_cols
577                .message_schedule
578                .w
579                .rows()
580                .into_iter()
581                .collect::<Vec<_>>(),
582        ]
583        .concat();
584        let w_limbs: Vec<Vec<F>> = w
585            .iter()
586            .map(|x| {
587                (0..C::WORD_U16S)
588                    .map(|i| compose::<F>(&x.as_slice().unwrap()[i * 16..(i + 1) * 16], 1))
589                    .collect::<Vec<F>>()
590            })
591            .collect();
592        for i in 0..C::ROUNDS_PER_ROW {
593            let sig_w = small_sig0_field::<F, C>(w[i + 1].as_slice().unwrap());
594            let sig_w_limbs: Vec<F> = (0..C::WORD_U16S)
595                .map(|j| compose::<F>(&sig_w[j * 16..(j + 1) * 16], 1))
596                .collect();
597            for (j, sig_w_limb) in sig_w_limbs.iter().enumerate() {
598                next_cols.schedule_helper.intermed_4[[i, j]] = w_limbs[i][j] + *sig_w_limb;
599            }
600        }
601    }
602
603    /// Puts the needed intermed_12 in the `local_row`
604    pub fn generate_intermed_12<F: PrimeField32>(
605        local_cols: &mut Sha2RoundColsRefMut<F>,
606        next_cols: Sha2RoundColsRef<F>,
607    ) {
608        let w = [
609            local_cols
610                .message_schedule
611                .w
612                .rows()
613                .into_iter()
614                .collect::<Vec<_>>(),
615            next_cols
616                .message_schedule
617                .w
618                .rows()
619                .into_iter()
620                .collect::<Vec<_>>(),
621        ]
622        .concat();
623        let w_limbs: Vec<Vec<F>> = w
624            .iter()
625            .map(|x| {
626                (0..C::WORD_U16S)
627                    .map(|i| compose::<F>(&x.as_slice().unwrap()[i * 16..(i + 1) * 16], 1))
628                    .collect::<Vec<F>>()
629            })
630            .collect();
631        for i in 0..C::ROUNDS_PER_ROW {
632            // sig_1(w_{t-2})
633            let sig_w_2: Vec<F> = (0..C::WORD_U16S)
634                .map(|j| {
635                    compose::<F>(
636                        &small_sig1_field::<F, C>(w[i + 2].as_slice().unwrap())
637                            [j * 16..(j + 1) * 16],
638                        1,
639                    )
640                })
641                .collect();
642            // w_{t-7}
643            let w_7 = if i < 3 {
644                local_cols.schedule_helper.w_3.row(i).to_slice().unwrap()
645            } else {
646                w_limbs[i - 3].as_slice()
647            };
648            // w_t
649            let w_cur = w_limbs[i + 4].as_slice();
650            for j in 0..C::WORD_U16S {
651                let carry = next_cols.message_schedule.carry_or_buffer[[i, j * 2]]
652                    + F::TWO * next_cols.message_schedule.carry_or_buffer[[i, j * 2 + 1]];
653                let sum = sig_w_2[j] + w_7[j] - carry * F::from_u32(1 << 16) - w_cur[j]
654                    + if j > 0 {
655                        next_cols.message_schedule.carry_or_buffer[[i, j * 2 - 2]]
656                            + F::from_u32(2)
657                                * next_cols.message_schedule.carry_or_buffer[[i, j * 2 - 1]]
658                    } else {
659                        F::ZERO
660                    };
661                local_cols.schedule_helper.intermed_12[[i, j]] = -sum;
662            }
663        }
664    }
665}
openvm_sha2_air/trace.rs

openvm_sha2_air/
trace.rs
