openvm_sha2_air/
utils.rs

1use ndarray::ArrayViewMut;
2pub use openvm_circuit_primitives::utils::compose;
3use openvm_circuit_primitives::{
4    encoder::Encoder,
5    utils::{not, select},
6};
7use openvm_stark_backend::{
8    p3_air::AirBuilder,
9    p3_field::{PrimeCharacteristicRing, PrimeField32},
10};
11use rand::{rngs::StdRng, Rng};
12
13use crate::{RotateRight, Sha2BlockHasherSubairConfig};
14
15/// Convert a word into a list of 8-bit limbs in little endian
16pub fn word_into_u8_limbs<C: Sha2BlockHasherSubairConfig>(num: impl Into<C::Word>) -> Vec<u32> {
17    word_into_limbs::<C>(num.into(), C::WORD_U8S)
18}
19
20/// Convert a word into a list of 16-bit limbs in little endian
21pub fn word_into_u16_limbs<C: Sha2BlockHasherSubairConfig>(num: impl Into<C::Word>) -> Vec<u32> {
22    word_into_limbs::<C>(num.into(), C::WORD_U16S)
23}
24
25/// Convert a word into a list of 1-bit limbs in little endian
26pub fn word_into_bits<C: Sha2BlockHasherSubairConfig>(num: impl Into<C::Word>) -> Vec<u32> {
27    word_into_limbs::<C>(num.into(), C::WORD_BITS)
28}
29
30/// Convert a word into a list of limbs in little endian
31pub fn word_into_limbs<C: Sha2BlockHasherSubairConfig>(num: C::Word, num_limbs: usize) -> Vec<u32> {
32    let limb_bits = std::mem::size_of::<C::Word>() * 8 / num_limbs;
33    (0..num_limbs)
34        .map(|i| {
35            let shifted = num >> (limb_bits * i);
36            let mask: C::Word = ((1u32 << limb_bits) - 1).into();
37            let masked = shifted & mask;
38            masked.try_into().unwrap()
39        })
40        .collect()
41}
42
43/// Convert a u32 into a list of 1-bit limbs in little endian
44pub fn u32_into_bits<C: Sha2BlockHasherSubairConfig>(num: u32) -> Vec<u32> {
45    let limb_bits = 32 / C::WORD_BITS;
46    (0..C::WORD_BITS)
47        .map(|i| (num >> (limb_bits * i)) & ((1 << limb_bits) - 1))
48        .collect()
49}
50
51/// Convert a list of limbs in little endian into a Word
52pub fn le_limbs_into_word<C: Sha2BlockHasherSubairConfig>(limbs: &[u32]) -> C::Word {
53    let mut limbs = limbs.to_vec();
54    limbs.reverse();
55    be_limbs_into_word::<C>(&limbs)
56}
57
58/// Convert a list of limbs in big endian into a Word
59pub fn be_limbs_into_word<C: Sha2BlockHasherSubairConfig>(limbs: &[u32]) -> C::Word {
60    let limb_bits = C::WORD_BITS / limbs.len();
61    limbs.iter().fold(C::Word::from(0), |acc, &limb| {
62        (acc << limb_bits) | limb.into()
63    })
64}
65
66/// Convert a list of limbs in little endian into a u32
67pub fn limbs_into_u32(limbs: &[u32]) -> u32 {
68    let limb_bits = 32 / limbs.len();
69    limbs
70        .iter()
71        .rev()
72        .fold(0, |acc, &limb| (acc << limb_bits) | limb)
73}
74
75/// Rotates `bits` right by `n` bits, assumes `bits` is in little-endian
76#[inline]
77pub(crate) fn rotr<F: PrimeCharacteristicRing + Clone>(
78    bits: &[impl Into<F> + Clone],
79    n: usize,
80) -> Vec<F> {
81    (0..bits.len())
82        .map(|i| bits[(i + n) % bits.len()].clone().into())
83        .collect()
84}
85
86/// Shifts `bits` right by `n` bits, assumes `bits` is in little-endian
87#[inline]
88pub(crate) fn shr<F: PrimeCharacteristicRing + Clone>(
89    bits: &[impl Into<F> + Clone],
90    n: usize,
91) -> Vec<F> {
92    (0..bits.len())
93        .map(|i| {
94            if i + n < bits.len() {
95                bits[i + n].clone().into()
96            } else {
97                F::ZERO
98            }
99        })
100        .collect()
101}
102
103/// Computes x ^ y ^ z, where x, y, z are assumed to be boolean
104#[inline]
105pub(crate) fn xor_bit<F: PrimeCharacteristicRing + Clone>(
106    x: impl Into<F>,
107    y: impl Into<F>,
108    z: impl Into<F>,
109) -> F {
110    let (x, y, z) = (x.into(), y.into(), z.into());
111    (x.clone() * y.clone() * z.clone())
112        + (x.clone() * not::<F>(y.clone()) * not::<F>(z.clone()))
113        + (not::<F>(x.clone()) * y.clone() * not::<F>(z.clone()))
114        + (not::<F>(x) * not::<F>(y) * z)
115}
116
117/// Computes x ^ y ^ z, where x, y, z are [C::WORD_BITS] bit numbers
118#[inline]
119pub(crate) fn xor<F: PrimeCharacteristicRing + Clone>(
120    x: &[impl Into<F> + Clone],
121    y: &[impl Into<F> + Clone],
122    z: &[impl Into<F> + Clone],
123) -> Vec<F> {
124    (0..x.len())
125        .map(|i| xor_bit(x[i].clone(), y[i].clone(), z[i].clone()))
126        .collect()
127}
128
129/// Choose function from the SHA spec
130#[inline]
131pub fn ch<C: Sha2BlockHasherSubairConfig>(x: C::Word, y: C::Word, z: C::Word) -> C::Word {
132    (x & y) ^ ((!x) & z)
133}
134
135/// Computes Ch(x,y,z), where x, y, z are [C::WORD_BITS] bit numbers
136#[inline]
137pub(crate) fn ch_field<F: PrimeCharacteristicRing>(
138    x: &[impl Into<F> + Clone],
139    y: &[impl Into<F> + Clone],
140    z: &[impl Into<F> + Clone],
141) -> Vec<F> {
142    (0..x.len())
143        .map(|i| select(x[i].clone(), y[i].clone(), z[i].clone()))
144        .collect()
145}
146
147/// Majority function from the SHA spec
148pub fn maj<C: Sha2BlockHasherSubairConfig>(x: C::Word, y: C::Word, z: C::Word) -> C::Word {
149    (x & y) ^ (x & z) ^ (y & z)
150}
151
152/// Computes Maj(x,y,z), where x, y, z are [C::WORD_BITS] bit numbers
153#[inline]
154pub(crate) fn maj_field<F: PrimeCharacteristicRing + Clone>(
155    x: &[impl Into<F> + Clone],
156    y: &[impl Into<F> + Clone],
157    z: &[impl Into<F> + Clone],
158) -> Vec<F> {
159    (0..x.len())
160        .map(|i| {
161            let (x, y, z) = (
162                x[i].clone().into(),
163                y[i].clone().into(),
164                z[i].clone().into(),
165            );
166            x.clone() * y.clone() + x.clone() * z.clone() + y.clone() * z.clone()
167                - F::TWO * x * y * z
168        })
169        .collect()
170}
171
172/// Big sigma_0 function from the SHA spec
173pub fn big_sig0<C: Sha2BlockHasherSubairConfig>(x: C::Word) -> C::Word {
174    if C::WORD_BITS == 32 {
175        x.rotate_right(2) ^ x.rotate_right(13) ^ x.rotate_right(22)
176    } else {
177        x.rotate_right(28) ^ x.rotate_right(34) ^ x.rotate_right(39)
178    }
179}
180
181/// Computes BigSigma0(x), where x is a [C::WORD_BITS] bit number in little-endian
182#[inline]
183pub(crate) fn big_sig0_field<F: PrimeCharacteristicRing + Clone, C: Sha2BlockHasherSubairConfig>(
184    x: &[impl Into<F> + Clone],
185) -> Vec<F> {
186    if C::WORD_BITS == 32 {
187        xor(&rotr::<F>(x, 2), &rotr::<F>(x, 13), &rotr::<F>(x, 22))
188    } else {
189        xor(&rotr::<F>(x, 28), &rotr::<F>(x, 34), &rotr::<F>(x, 39))
190    }
191}
192
193/// Big sigma_1 function from the SHA spec
194pub fn big_sig1<C: Sha2BlockHasherSubairConfig>(x: C::Word) -> C::Word {
195    if C::WORD_BITS == 32 {
196        x.rotate_right(6) ^ x.rotate_right(11) ^ x.rotate_right(25)
197    } else {
198        x.rotate_right(14) ^ x.rotate_right(18) ^ x.rotate_right(41)
199    }
200}
201
202/// Computes BigSigma1(x), where x is a [C::WORD_BITS] bit number in little-endian
203#[inline]
204pub(crate) fn big_sig1_field<F: PrimeCharacteristicRing + Clone, C: Sha2BlockHasherSubairConfig>(
205    x: &[impl Into<F> + Clone],
206) -> Vec<F> {
207    if C::WORD_BITS == 32 {
208        xor(&rotr::<F>(x, 6), &rotr::<F>(x, 11), &rotr::<F>(x, 25))
209    } else {
210        xor(&rotr::<F>(x, 14), &rotr::<F>(x, 18), &rotr::<F>(x, 41))
211    }
212}
213
214/// Small sigma_0 function from the SHA spec
215pub fn small_sig0<C: Sha2BlockHasherSubairConfig>(x: C::Word) -> C::Word {
216    if C::WORD_BITS == 32 {
217        x.rotate_right(7) ^ x.rotate_right(18) ^ (x >> 3)
218    } else {
219        x.rotate_right(1) ^ x.rotate_right(8) ^ (x >> 7)
220    }
221}
222
223/// Computes SmallSigma0(x), where x is a [C::WORD_BITS] bit number in little-endian
224#[inline]
225pub(crate) fn small_sig0_field<
226    F: PrimeCharacteristicRing + Clone,
227    C: Sha2BlockHasherSubairConfig,
228>(
229    x: &[impl Into<F> + Clone],
230) -> Vec<F> {
231    if C::WORD_BITS == 32 {
232        xor(&rotr::<F>(x, 7), &rotr::<F>(x, 18), &shr::<F>(x, 3))
233    } else {
234        xor(&rotr::<F>(x, 1), &rotr::<F>(x, 8), &shr::<F>(x, 7))
235    }
236}
237
238/// Small sigma_1 function from the SHA spec
239pub fn small_sig1<C: Sha2BlockHasherSubairConfig>(x: C::Word) -> C::Word {
240    if C::WORD_BITS == 32 {
241        x.rotate_right(17) ^ x.rotate_right(19) ^ (x >> 10)
242    } else {
243        x.rotate_right(19) ^ x.rotate_right(61) ^ (x >> 6)
244    }
245}
246
247/// Computes SmallSigma1(x), where x is a [C::WORD_BITS] bit number in little-endian
248#[inline]
249pub(crate) fn small_sig1_field<
250    F: PrimeCharacteristicRing + Clone,
251    C: Sha2BlockHasherSubairConfig,
252>(
253    x: &[impl Into<F> + Clone],
254) -> Vec<F> {
255    if C::WORD_BITS == 32 {
256        xor(&rotr::<F>(x, 17), &rotr::<F>(x, 19), &shr::<F>(x, 10))
257    } else {
258        xor(&rotr::<F>(x, 19), &rotr::<F>(x, 61), &shr::<F>(x, 6))
259    }
260}
261
262/// Generate a random message of a given length
263pub fn get_random_message(rng: &mut StdRng, len: usize) -> Vec<u8> {
264    let mut random_message: Vec<u8> = vec![0u8; len];
265    rng.fill(&mut random_message[..]);
266    random_message
267}
268
269/// Wrapper of `get_flag_pt` to get the flag pointer as an array
270pub fn get_flag_pt_array(encoder: &Encoder, flag_idx: usize) -> Vec<u32> {
271    encoder.get_flag_pt(flag_idx)
272}
273
274/// Constrain the addition of `C::WORD_BITS` bit words in 16-bit limbs
275/// It takes in the terms some in bits some in 16-bit limbs,
276/// the expected sum in bits and the carries
277pub fn constraint_word_addition<AB: AirBuilder, C: Sha2BlockHasherSubairConfig>(
278    builder: &mut AB,
279    terms_bits: &[&[impl Into<AB::Expr> + Clone]],
280    terms_limb: &[&[impl Into<AB::Expr> + Clone]],
281    expected_sum: &[impl Into<AB::Expr> + Clone],
282    carries: &[impl Into<AB::Expr> + Clone],
283) {
284    debug_assert!(terms_bits.iter().all(|x| x.len() == C::WORD_BITS));
285    debug_assert!(terms_limb.iter().all(|x| x.len() == C::WORD_U16S));
286    assert_eq!(expected_sum.len(), C::WORD_BITS);
287    assert_eq!(carries.len(), C::WORD_U16S);
288
289    for i in 0..C::WORD_U16S {
290        let mut limb_sum = if i == 0 {
291            AB::Expr::ZERO
292        } else {
293            carries[i - 1].clone().into()
294        };
295        for term in terms_bits {
296            limb_sum += compose::<AB::Expr>(&term[i * 16..(i + 1) * 16], 1);
297        }
298        for term in terms_limb {
299            limb_sum += term[i].clone().into();
300        }
301        let expected_sum_limb = compose::<AB::Expr>(&expected_sum[i * 16..(i + 1) * 16], 1)
302            + carries[i].clone().into() * AB::Expr::from_u32(1 << 16);
303        builder.assert_eq(limb_sum, expected_sum_limb);
304    }
305}
306
307pub fn set_arrayview_from_u32_slice<F: PrimeField32, D: ndarray::Dimension>(
308    arrayview: &mut ArrayViewMut<F, D>,
309    data: impl IntoIterator<Item = u32>,
310) {
311    arrayview
312        .iter_mut()
313        .zip(data.into_iter().map(|x| F::from_u32(x)))
314        .for_each(|(x, y)| *x = y);
315}
316
317pub fn set_arrayview_from_u8_slice<F: PrimeField32, D: ndarray::Dimension>(
318    arrayview: &mut ArrayViewMut<F, D>,
319    data: impl IntoIterator<Item = u8>,
320) {
321    arrayview
322        .iter_mut()
323        .zip(data.into_iter().map(|x| F::from_u8(x)))
324        .for_each(|(x, y)| *x = y);
325}
openvm_sha2_air/utils.rs

openvm_sha2_air/
utils.rs
