openvm_sha256_air/
utils.rs

1use std::array;
2
3pub use openvm_circuit_primitives::utils::compose;
4use openvm_circuit_primitives::{
5    encoder::Encoder,
6    utils::{not, select},
7};
8use openvm_stark_backend::{p3_air::AirBuilder, p3_field::FieldAlgebra};
9
10use super::{Sha256DigestCols, Sha256RoundCols};
11
12// ==== Do not change these constants! ====
13/// Number of bits in a SHA256 word
14pub const SHA256_WORD_BITS: usize = 32;
15/// Number of 16-bit limbs in a SHA256 word
16pub const SHA256_WORD_U16S: usize = SHA256_WORD_BITS / 16;
17/// Number of 8-bit limbs in a SHA256 word
18pub const SHA256_WORD_U8S: usize = SHA256_WORD_BITS / 8;
19/// Number of words in a SHA256 block
20pub const SHA256_BLOCK_WORDS: usize = 16;
21/// Number of cells in a SHA256 block
22pub const SHA256_BLOCK_U8S: usize = SHA256_BLOCK_WORDS * SHA256_WORD_U8S;
23/// Number of bits in a SHA256 block
24pub const SHA256_BLOCK_BITS: usize = SHA256_BLOCK_WORDS * SHA256_WORD_BITS;
25/// Number of rows per block
26pub const SHA256_ROWS_PER_BLOCK: usize = 17;
27/// Number of rounds per row
28pub const SHA256_ROUNDS_PER_ROW: usize = 4;
29/// Number of words in a SHA256 hash
30pub const SHA256_HASH_WORDS: usize = 8;
31/// Number of vars needed to encode the row index with [Encoder]
32pub const SHA256_ROW_VAR_CNT: usize = 5;
33/// Width of the Sha256RoundCols
34pub const SHA256_ROUND_WIDTH: usize = Sha256RoundCols::<u8>::width();
35/// Width of the Sha256DigestCols
36pub const SHA256_DIGEST_WIDTH: usize = Sha256DigestCols::<u8>::width();
37/// Size of the buffer of the first 4 rows of a block (each row's size)
38pub const SHA256_BUFFER_SIZE: usize = SHA256_ROUNDS_PER_ROW * SHA256_WORD_U16S * 2;
39/// Width of the Sha256Cols
40pub const SHA256_WIDTH: usize = if SHA256_ROUND_WIDTH > SHA256_DIGEST_WIDTH {
41    SHA256_ROUND_WIDTH
42} else {
43    SHA256_DIGEST_WIDTH
44};
45/// We can notice that `carry_a`'s and `carry_e`'s are always the same on invalid rows
46/// To optimize the trace generation of invalid rows, we have those values precomputed here
47pub(crate) const SHA256_INVALID_CARRY_A: [[u32; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW] = [
48    [1230919683, 1162494304],
49    [266373122, 1282901987],
50    [1519718403, 1008990871],
51    [923381762, 330807052],
52];
53pub(crate) const SHA256_INVALID_CARRY_E: [[u32; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW] = [
54    [204933122, 1994683449],
55    [443873282, 1544639095],
56    [719953922, 1888246508],
57    [194580482, 1075725211],
58];
59/// SHA256 constant K's
60pub const SHA256_K: [u32; 64] = [
61    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
62    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
63    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
64    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
65    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
66    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
67    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
68    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
69];
70
71/// SHA256 initial hash values
72pub const SHA256_H: [u32; 8] = [
73    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
74];
75
76/// Returns the number of blocks required to hash a message of length `len`
77pub fn get_sha256_num_blocks(len: u32) -> u32 {
78    // need to pad with one 1 bit, 64 bits for the message length and then pad until the length
79    // is divisible by [SHA256_BLOCK_BITS]
80    ((len << 3) as usize + 1 + 64).div_ceil(SHA256_BLOCK_BITS) as u32
81}
82
83/// Convert a u32 into a list of bits in little endian then convert each bit into a field element
84pub fn u32_into_bits_field<F: FieldAlgebra + Clone>(num: u32) -> [F; SHA256_WORD_BITS] {
85    array::from_fn(|i| F::from_bool((num >> i) & 1 == 1))
86}
87
88/// Convert a u32 into a an array of 2 16-bit limbs in little endian
89pub fn u32_into_u16s(num: u32) -> [u32; 2] {
90    [num & 0xffff, num >> 16]
91}
92
93/// Convert a list of limbs in little endian into a u32
94pub fn limbs_into_u32<const NUM_LIMBS: usize>(limbs: [u32; NUM_LIMBS]) -> u32 {
95    let limb_bits = 32 / NUM_LIMBS;
96    limbs
97        .iter()
98        .rev()
99        .fold(0, |acc, &limb| (acc << limb_bits) | limb)
100}
101
102/// Rotates `bits` right by `n` bits, assumes `bits` is in little-endian
103#[inline]
104pub(crate) fn rotr<F: FieldAlgebra + Clone>(
105    bits: &[impl Into<F> + Clone; SHA256_WORD_BITS],
106    n: usize,
107) -> [F; SHA256_WORD_BITS] {
108    array::from_fn(|i| bits[(i + n) % SHA256_WORD_BITS].clone().into())
109}
110
111/// Shifts `bits` right by `n` bits, assumes `bits` is in little-endian
112#[inline]
113pub(crate) fn shr<F: FieldAlgebra + Clone>(
114    bits: &[impl Into<F> + Clone; SHA256_WORD_BITS],
115    n: usize,
116) -> [F; SHA256_WORD_BITS] {
117    array::from_fn(|i| {
118        if i + n < SHA256_WORD_BITS {
119            bits[i + n].clone().into()
120        } else {
121            F::ZERO
122        }
123    })
124}
125
126/// Computes x ^ y ^ z, where x, y, z are assumed to be boolean
127#[inline]
128pub(crate) fn xor_bit<F: FieldAlgebra + Clone>(
129    x: impl Into<F>,
130    y: impl Into<F>,
131    z: impl Into<F>,
132) -> F {
133    let (x, y, z) = (x.into(), y.into(), z.into());
134    (x.clone() * y.clone() * z.clone())
135        + (x.clone() * not::<F>(y.clone()) * not::<F>(z.clone()))
136        + (not::<F>(x.clone()) * y.clone() * not::<F>(z.clone()))
137        + (not::<F>(x) * not::<F>(y) * z)
138}
139
140/// Computes x ^ y ^ z, where x, y, z are [SHA256_WORD_BITS] bit numbers
141#[inline]
142pub(crate) fn xor<F: FieldAlgebra + Clone>(
143    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
144    y: &[impl Into<F> + Clone; SHA256_WORD_BITS],
145    z: &[impl Into<F> + Clone; SHA256_WORD_BITS],
146) -> [F; SHA256_WORD_BITS] {
147    array::from_fn(|i| xor_bit(x[i].clone(), y[i].clone(), z[i].clone()))
148}
149
150/// Choose function from SHA256
151#[inline]
152pub fn ch(x: u32, y: u32, z: u32) -> u32 {
153    (x & y) ^ ((!x) & z)
154}
155
156/// Computes Ch(x,y,z), where x, y, z are [SHA256_WORD_BITS] bit numbers
157#[inline]
158pub(crate) fn ch_field<F: FieldAlgebra>(
159    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
160    y: &[impl Into<F> + Clone; SHA256_WORD_BITS],
161    z: &[impl Into<F> + Clone; SHA256_WORD_BITS],
162) -> [F; SHA256_WORD_BITS] {
163    array::from_fn(|i| select(x[i].clone(), y[i].clone(), z[i].clone()))
164}
165
166/// Majority function from SHA256
167pub fn maj(x: u32, y: u32, z: u32) -> u32 {
168    (x & y) ^ (x & z) ^ (y & z)
169}
170
171/// Computes Maj(x,y,z), where x, y, z are [SHA256_WORD_BITS] bit numbers
172#[inline]
173pub(crate) fn maj_field<F: FieldAlgebra + Clone>(
174    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
175    y: &[impl Into<F> + Clone; SHA256_WORD_BITS],
176    z: &[impl Into<F> + Clone; SHA256_WORD_BITS],
177) -> [F; SHA256_WORD_BITS] {
178    array::from_fn(|i| {
179        let (x, y, z) = (
180            x[i].clone().into(),
181            y[i].clone().into(),
182            z[i].clone().into(),
183        );
184        x.clone() * y.clone() + x.clone() * z.clone() + y.clone() * z.clone() - F::TWO * x * y * z
185    })
186}
187
188/// Big sigma_0 function from SHA256
189pub fn big_sig0(x: u32) -> u32 {
190    x.rotate_right(2) ^ x.rotate_right(13) ^ x.rotate_right(22)
191}
192
193/// Computes BigSigma0(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
194#[inline]
195pub(crate) fn big_sig0_field<F: FieldAlgebra + Clone>(
196    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
197) -> [F; SHA256_WORD_BITS] {
198    xor(&rotr::<F>(x, 2), &rotr::<F>(x, 13), &rotr::<F>(x, 22))
199}
200
201/// Big sigma_1 function from SHA256
202pub fn big_sig1(x: u32) -> u32 {
203    x.rotate_right(6) ^ x.rotate_right(11) ^ x.rotate_right(25)
204}
205
206/// Computes BigSigma1(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
207#[inline]
208pub(crate) fn big_sig1_field<F: FieldAlgebra + Clone>(
209    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
210) -> [F; SHA256_WORD_BITS] {
211    xor(&rotr::<F>(x, 6), &rotr::<F>(x, 11), &rotr::<F>(x, 25))
212}
213
214/// Small sigma_0 function from SHA256
215pub fn small_sig0(x: u32) -> u32 {
216    x.rotate_right(7) ^ x.rotate_right(18) ^ (x >> 3)
217}
218
219/// Computes SmallSigma0(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
220#[inline]
221pub(crate) fn small_sig0_field<F: FieldAlgebra + Clone>(
222    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
223) -> [F; SHA256_WORD_BITS] {
224    xor(&rotr::<F>(x, 7), &rotr::<F>(x, 18), &shr::<F>(x, 3))
225}
226
227/// Small sigma_1 function from SHA256
228pub fn small_sig1(x: u32) -> u32 {
229    x.rotate_right(17) ^ x.rotate_right(19) ^ (x >> 10)
230}
231
232/// Computes SmallSigma1(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
233#[inline]
234pub(crate) fn small_sig1_field<F: FieldAlgebra + Clone>(
235    x: &[impl Into<F> + Clone; SHA256_WORD_BITS],
236) -> [F; SHA256_WORD_BITS] {
237    xor(&rotr::<F>(x, 17), &rotr::<F>(x, 19), &shr::<F>(x, 10))
238}
239
240/// Wrapper of `get_flag_pt` to get the flag pointer as an array
241pub fn get_flag_pt_array<const N: usize>(encoder: &Encoder, flag_idx: usize) -> [u32; N] {
242    encoder.get_flag_pt(flag_idx).try_into().unwrap()
243}
244
245/// Constrain the addition of [SHA256_WORD_BITS] bit words in 16-bit limbs
246/// It takes in the terms some in bits some in 16-bit limbs,
247/// the expected sum in bits and the carries
248pub fn constraint_word_addition<AB: AirBuilder>(
249    builder: &mut AB,
250    terms_bits: &[&[impl Into<AB::Expr> + Clone; SHA256_WORD_BITS]],
251    terms_limb: &[&[impl Into<AB::Expr> + Clone; SHA256_WORD_U16S]],
252    expected_sum: &[impl Into<AB::Expr> + Clone; SHA256_WORD_BITS],
253    carries: &[impl Into<AB::Expr> + Clone; SHA256_WORD_U16S],
254) {
255    for i in 0..SHA256_WORD_U16S {
256        let mut limb_sum = if i == 0 {
257            AB::Expr::ZERO
258        } else {
259            carries[i - 1].clone().into()
260        };
261        for term in terms_bits {
262            limb_sum += compose::<AB::Expr>(&term[i * 16..(i + 1) * 16], 1);
263        }
264        for term in terms_limb {
265            limb_sum += term[i].clone().into();
266        }
267        let expected_sum_limb = compose::<AB::Expr>(&expected_sum[i * 16..(i + 1) * 16], 1)
268            + carries[i].clone().into() * AB::Expr::from_canonical_u32(1 << 16);
269        builder.assert_eq(limb_sum, expected_sum_limb);
270    }
271}
openvm_sha256_air/utils.rs

openvm_sha256_air/
utils.rs
