sp1_core_machine/air/
word.rs

1use std::array;
2
3use itertools::Itertools;
4use p3_field::AbstractField;
5use sp1_core_executor::ByteOpcode;
6use sp1_primitives::consts::WORD_SIZE;
7use sp1_stark::{air::ByteAirBuilder, Word};
8
9pub trait WordAirBuilder: ByteAirBuilder {
10    /// Asserts that the two words are equal.
11    fn assert_word_eq(
12        &mut self,
13        left: Word<impl Into<Self::Expr>>,
14        right: Word<impl Into<Self::Expr>>,
15    ) {
16        for (left, right) in left.0.into_iter().zip(right.0) {
17            self.assert_eq(left, right);
18        }
19    }
20
21    /// Asserts that the word is zero.
22    fn assert_word_zero(&mut self, word: Word<impl Into<Self::Expr>>) {
23        for limb in word.0 {
24            self.assert_zero(limb);
25        }
26    }
27
28    /// Index an array of words using an index bitmap.
29    fn index_word_array(
30        &mut self,
31        array: &[Word<impl Into<Self::Expr> + Clone>],
32        index_bitmap: &[impl Into<Self::Expr> + Clone],
33    ) -> Word<Self::Expr> {
34        let mut result = Word::default();
35        for i in 0..WORD_SIZE {
36            result[i] = self.index_array(
37                array.iter().map(|word| word[i].clone()).collect_vec().as_slice(),
38                index_bitmap,
39            );
40        }
41        result
42    }
43
44    /// Same as `if_else` above, but arguments are `Word` instead of individual expressions.
45    fn select_word(
46        &mut self,
47        condition: impl Into<Self::Expr> + Clone,
48        a: Word<impl Into<Self::Expr> + Clone>,
49        b: Word<impl Into<Self::Expr> + Clone>,
50    ) -> Word<Self::Expr> {
51        Word(array::from_fn(|i| self.if_else(condition.clone(), a[i].clone(), b[i].clone())))
52    }
53
54    /// Check that each limb of the given slice is a u8.
55    fn slice_range_check_u8(
56        &mut self,
57        input: &[impl Into<Self::Expr> + Clone],
58        mult: impl Into<Self::Expr> + Clone,
59    ) {
60        let mut index = 0;
61        while index + 1 < input.len() {
62            self.send_byte(
63                Self::Expr::from_canonical_u8(ByteOpcode::U8Range as u8),
64                Self::Expr::zero(),
65                input[index].clone(),
66                input[index + 1].clone(),
67                mult.clone(),
68            );
69            index += 2;
70        }
71        if index < input.len() {
72            self.send_byte(
73                Self::Expr::from_canonical_u8(ByteOpcode::U8Range as u8),
74                Self::Expr::zero(),
75                input[index].clone(),
76                Self::Expr::zero(),
77                mult.clone(),
78            );
79        }
80    }
81
82    /// Check that each limb of the given slice is a u16.
83    fn slice_range_check_u16(
84        &mut self,
85        input: &[impl Into<Self::Expr> + Copy],
86        mult: impl Into<Self::Expr> + Clone,
87    ) {
88        input.iter().for_each(|limb| {
89            self.send_byte(
90                Self::Expr::from_canonical_u8(ByteOpcode::U16Range as u8),
91                *limb,
92                Self::Expr::zero(),
93                Self::Expr::zero(),
94                mult.clone(),
95            );
96        });
97    }
98}