sp1_core_machine/operations/
add5.rs

1use p3_air::AirBuilder;
2use p3_field::{AbstractField, Field};
3use sp1_derive::AlignedBorrow;
4
5use sp1_core_executor::events::ByteRecord;
6use sp1_primitives::consts::WORD_SIZE;
7use sp1_stark::{air::SP1AirBuilder, Word};
8
9use crate::air::WordAirBuilder;
10
11/// A set of columns needed to compute the sum of five words.
12#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
13#[repr(C)]
14pub struct Add5Operation<T> {
15    /// The result of `a + b + c + d + e`.
16    pub value: Word<T>,
17
18    /// Indicates if the carry for the `i`th limb is 0.
19    pub is_carry_0: Word<T>,
20
21    /// Indicates if the carry for the `i`th limb is 1.
22    pub is_carry_1: Word<T>,
23
24    /// Indicates if the carry for the `i`th limb is 2.
25    pub is_carry_2: Word<T>,
26
27    /// Indicates if the carry for the `i`th limb is 3.
28    pub is_carry_3: Word<T>,
29
30    /// Indicates if the carry for the `i`th limb is 4. The carry when adding 5 words is at most 4.
31    pub is_carry_4: Word<T>,
32
33    /// The carry for the `i`th limb.
34    pub carry: Word<T>,
35}
36
37impl<F: Field> Add5Operation<F> {
38    #[allow(clippy::too_many_arguments)]
39    pub fn populate(
40        &mut self,
41        record: &mut impl ByteRecord,
42        a_u32: u32,
43        b_u32: u32,
44        c_u32: u32,
45        d_u32: u32,
46        e_u32: u32,
47    ) -> u32 {
48        let expected =
49            a_u32.wrapping_add(b_u32).wrapping_add(c_u32).wrapping_add(d_u32).wrapping_add(e_u32);
50
51        self.value = Word::from(expected);
52        let a = a_u32.to_le_bytes();
53        let b = b_u32.to_le_bytes();
54        let c = c_u32.to_le_bytes();
55        let d = d_u32.to_le_bytes();
56        let e = e_u32.to_le_bytes();
57
58        let base = 256;
59        let mut carry = [0u8, 0u8, 0u8, 0u8, 0u8];
60        for i in 0..WORD_SIZE {
61            let mut res =
62                (a[i] as u32) + (b[i] as u32) + (c[i] as u32) + (d[i] as u32) + (e[i] as u32);
63            if i > 0 {
64                res += carry[i - 1] as u32;
65            }
66            carry[i] = (res / base) as u8;
67            self.is_carry_0[i] = F::from_bool(carry[i] == 0);
68            self.is_carry_1[i] = F::from_bool(carry[i] == 1);
69            self.is_carry_2[i] = F::from_bool(carry[i] == 2);
70            self.is_carry_3[i] = F::from_bool(carry[i] == 3);
71            self.is_carry_4[i] = F::from_bool(carry[i] == 4);
72            self.carry[i] = F::from_canonical_u8(carry[i]);
73            debug_assert!(carry[i] <= 4);
74            debug_assert_eq!(self.value[i], F::from_canonical_u32(res % base));
75        }
76
77        // Range check.
78        {
79            record.add_u8_range_checks(&a);
80            record.add_u8_range_checks(&b);
81            record.add_u8_range_checks(&c);
82            record.add_u8_range_checks(&d);
83            record.add_u8_range_checks(&e);
84            record.add_u8_range_checks(&expected.to_le_bytes());
85        }
86
87        expected
88    }
89
90    pub fn eval<AB: SP1AirBuilder>(
91        builder: &mut AB,
92        words: &[Word<AB::Var>; 5],
93        is_real: AB::Var,
94        cols: Add5Operation<AB::Var>,
95    ) {
96        builder.assert_bool(is_real);
97        // Range check each byte.
98        {
99            words.iter().for_each(|word| builder.slice_range_check_u8(&word.0, is_real));
100            builder.slice_range_check_u8(&cols.value.0, is_real);
101        }
102        let mut builder_is_real = builder.when(is_real);
103
104        // Each value in is_carry_{0,1,2,3,4} is 0 or 1, and exactly one of them is 1 per digit.
105        {
106            for i in 0..WORD_SIZE {
107                builder_is_real.assert_bool(cols.is_carry_0[i]);
108                builder_is_real.assert_bool(cols.is_carry_1[i]);
109                builder_is_real.assert_bool(cols.is_carry_2[i]);
110                builder_is_real.assert_bool(cols.is_carry_3[i]);
111                builder_is_real.assert_bool(cols.is_carry_4[i]);
112                builder_is_real.assert_eq(
113                    cols.is_carry_0[i] +
114                        cols.is_carry_1[i] +
115                        cols.is_carry_2[i] +
116                        cols.is_carry_3[i] +
117                        cols.is_carry_4[i],
118                    AB::Expr::one(),
119                );
120            }
121        }
122
123        // Calculates carry from is_carry_{0,1,2,3,4}.
124        {
125            let one = AB::Expr::one();
126            let two = AB::F::from_canonical_u32(2);
127            let three = AB::F::from_canonical_u32(3);
128            let four = AB::F::from_canonical_u32(4);
129
130            for i in 0..WORD_SIZE {
131                builder_is_real.assert_eq(
132                    cols.carry[i],
133                    cols.is_carry_1[i] * one.clone() +
134                        cols.is_carry_2[i] * two +
135                        cols.is_carry_3[i] * three +
136                        cols.is_carry_4[i] * four,
137                );
138            }
139        }
140
141        // Compare the sum and summands by looking at carry.
142        {
143            let base = AB::F::from_canonical_u32(256);
144            // For each limb, assert that difference between the carried result and the non-carried
145            // result is the product of carry and base.
146            for i in 0..WORD_SIZE {
147                let mut overflow: AB::Expr = AB::F::zero().into();
148                for word in words {
149                    overflow = overflow.clone() + word[i].into();
150                }
151                overflow = overflow.clone() - cols.value[i].into();
152
153                if i > 0 {
154                    overflow = overflow.clone() + cols.carry[i - 1].into();
155                }
156                builder_is_real.assert_eq(cols.carry[i] * base, overflow.clone());
157            }
158        }
159    }
160}