sp1_core_machine/air/
word.rs1use 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 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 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 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 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 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 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}