openvm_sha256_air/
utils.rs

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

openvm_sha256_air/
utils.rs
