openvm_sha2_air/
columns.rs

1//! WARNING: the order of fields in the structs is important, do not change it
2
3use core::ops::Add;
4
5use openvm_circuit_primitives::utils::not;
6use openvm_circuit_primitives_derive::ColsRef;
7use openvm_stark_backend::p3_field::PrimeCharacteristicRing;
8
9use crate::Sha2BlockHasherSubairConfig;
10
11/// In each SHA block:
12/// - First C::ROUND_ROWS rows use Sha2RoundCols
13/// - Final row uses Sha2DigestCols
14///
15/// For soundness, the AIR enforces that the trace ends in padding rows after the last digest row.
16/// Honest traces already satisfy this because the unpadded height is a multiple of 17 (SHA-256)
17/// or 21 (SHA-512), and thus not a power of 2.
18///
19/// Sha2RoundCols and Sha2DigestCols share the same first 3 fields:
20/// - flags
21/// - work_vars/hash (same type, different name)
22/// - schedule_helper
23///
24/// This design allows for:
25/// 1. Common constraints to work on either struct type by accessing these shared fields
26/// 2. Specific constraints to use the appropriate struct, with flags helping to do conditional
27///    constraints
28///
29/// Note that the `Sha2WorkVarsCols` field is used for different purposes in the two structs.
30#[repr(C)]
31#[derive(Clone, Copy, Debug, ColsRef)]
32#[config(Sha2BlockHasherSubairConfig)]
33pub struct Sha2RoundCols<
34    T,
35    const WORD_BITS: usize,
36    const WORD_U8S: usize,
37    const WORD_U16S: usize,
38    const ROUNDS_PER_ROW: usize,
39    const ROUNDS_PER_ROW_MINUS_ONE: usize,
40    const ROW_VAR_CNT: usize,
41> {
42    pub flags: Sha2FlagsCols<T, ROW_VAR_CNT>,
43    pub work_vars: Sha2WorkVarsCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U16S>,
44    pub schedule_helper:
45        Sha2MessageHelperCols<T, WORD_U16S, ROUNDS_PER_ROW, ROUNDS_PER_ROW_MINUS_ONE>,
46    pub message_schedule: Sha2MessageScheduleCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U8S>,
47}
48
49#[repr(C)]
50#[derive(Clone, Copy, Debug, ColsRef)]
51#[config(Sha2BlockHasherSubairConfig)]
52pub struct Sha2DigestCols<
53    T,
54    const WORD_BITS: usize,
55    const WORD_U8S: usize,
56    const WORD_U16S: usize,
57    const HASH_WORDS: usize,
58    const ROUNDS_PER_ROW: usize,
59    const ROUNDS_PER_ROW_MINUS_ONE: usize,
60    const ROW_VAR_CNT: usize,
61> {
62    pub flags: Sha2FlagsCols<T, ROW_VAR_CNT>,
63    /// Will serve as previous hash values for the next block
64    pub hash: Sha2WorkVarsCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U16S>,
65    pub schedule_helper:
66        Sha2MessageHelperCols<T, WORD_U16S, ROUNDS_PER_ROW, ROUNDS_PER_ROW_MINUS_ONE>,
67    /// The actual final hash values of the given block
68    /// Note: the above `hash` will be equal to `final_hash` unless we are on the last block
69    pub final_hash: [[T; WORD_U8S]; HASH_WORDS],
70    /// The final hash of the previous block
71    /// Note: will be constrained using interactions with the chip itself
72    pub prev_hash: [[T; WORD_U16S]; HASH_WORDS],
73}
74
75#[repr(C)]
76#[derive(Clone, Copy, Debug, ColsRef)]
77#[config(Sha2BlockHasherSubairConfig)]
78pub struct Sha2MessageScheduleCols<
79    T,
80    const WORD_BITS: usize,
81    const ROUNDS_PER_ROW: usize,
82    const WORD_U8S: usize,
83> {
84    /// The message schedule words as bits
85    /// The first 16 words will be the message data
86    pub w: [[T; WORD_BITS]; ROUNDS_PER_ROW],
87    /// Will be message schedule carries for rows 4..C::ROUND_ROWS and a buffer for rows 0..4 to be
88    /// used freely by wrapper chips Note: carries are 2 bit numbers represented using 2 cells
89    /// as individual bits
90    /// Note: carry_or_buffer is left unconstrained on rounds 0..3
91    pub carry_or_buffer: [[T; WORD_U8S]; ROUNDS_PER_ROW],
92}
93
94#[repr(C)]
95#[derive(Clone, Copy, Debug, ColsRef)]
96#[config(Sha2BlockHasherSubairConfig)]
97pub struct Sha2WorkVarsCols<
98    T,
99    const WORD_BITS: usize,
100    const ROUNDS_PER_ROW: usize,
101    const WORD_U16S: usize,
102> {
103    /// `a` and `e` after each iteration as 32-bits
104    pub a: [[T; WORD_BITS]; ROUNDS_PER_ROW],
105    pub e: [[T; WORD_BITS]; ROUNDS_PER_ROW],
106    /// The carry's used for addition during each iteration when computing `a` and `e`
107    pub carry_a: [[T; WORD_U16S]; ROUNDS_PER_ROW],
108    pub carry_e: [[T; WORD_U16S]; ROUNDS_PER_ROW],
109}
110
111/// These are the columns that are used to help with the message schedule additions
112/// Note: these need to be correctly assigned for every row even on padding rows
113#[repr(C)]
114#[derive(Clone, Copy, Debug, ColsRef)]
115#[config(Sha2BlockHasherSubairConfig)]
116pub struct Sha2MessageHelperCols<
117    T,
118    const WORD_U16S: usize,
119    const ROUNDS_PER_ROW: usize,
120    const ROUNDS_PER_ROW_MINUS_ONE: usize,
121> {
122    /// The following are used to move data forward to constrain the message schedule additions
123    /// The value of `w` from 3 rounds ago
124    pub w_3: [[T; WORD_U16S]; ROUNDS_PER_ROW_MINUS_ONE],
125    /// Here intermediate(i) =  w_i + sig_0(w_{i+1})
126    /// Intermed_t represents the intermediate t rounds ago
127    /// This is needed to constrain the message schedule, since we can only constrain on two rows
128    /// at a time
129    pub intermed_4: [[T; WORD_U16S]; ROUNDS_PER_ROW],
130    pub intermed_8: [[T; WORD_U16S]; ROUNDS_PER_ROW],
131    pub intermed_12: [[T; WORD_U16S]; ROUNDS_PER_ROW],
132}
133
134#[repr(C)]
135#[derive(Clone, Copy, Debug, ColsRef)]
136#[config(Sha2BlockHasherSubairConfig)]
137pub struct Sha2FlagsCols<T, const ROW_VAR_CNT: usize> {
138    pub is_round_row: T,
139    /// A flag that indicates if the current row is among the first 4 rows of a block (the message
140    /// rows)
141    pub is_first_4_rows: T,
142    pub is_digest_row: T,
143    /// We will encode the row index [0..C::ROWS_PER_BLOCK] using ROW_VAR_CNT cells
144    pub row_idx: [T; ROW_VAR_CNT],
145    /// The global index of the current block, starts at 1 for the first block
146    /// and increments by 1 for each block. Padding rows use one more than the
147    /// last real block's index and remain constant across all padding rows.
148    pub global_block_idx: T,
149}
150
151impl<O, T> Sha2FlagsColsRef<'_, T>
152where
153    T: Add<Output = O> + Copy,
154{
155    // This refers to the padding rows that are added to the air to make the trace length a power of
156    // 2. Not to be confused with the padding added to messages as part of the SHA hash
157    // function.
158    pub fn is_not_padding_row(&self) -> O {
159        *self.is_round_row + *self.is_digest_row
160    }
161
162    // This refers to the padding rows that are added to the air to make the trace length a power of
163    // 2. Not to be confused with the padding added to messages as part of the SHA hash
164    // function.
165    pub fn is_padding_row(&self) -> O
166    where
167        O: PrimeCharacteristicRing,
168    {
169        not(self.is_not_padding_row())
170    }
171}
openvm_sha2_air/columns.rs

openvm_sha2_air/
columns.rs
