Skip to main content

tasm_lib/arithmetic/u192/
overflowing_add.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5/// Mimics overflowing_add but works on u192 data types.
6///
7/// ### Behavior
8///
9/// ```text
10/// BEFORE: _ [rhs: u192] [lhs: u192]
11/// AFTER:  _ [sum: u192] [is_overflow: bool]
12/// ```
13///
14/// ### Preconditions
15///
16/// - all input arguments are properly [`BFieldCodec`] encoded
17///
18/// ### Postconditions
19///
20/// - the output is properly [`BFieldCodec`] encoded
21#[derive(Debug, Clone)]
22pub struct OverflowingAdd;
23
24impl OverflowingAdd {
25    /// Generate code to perform an addition on `u192`s.
26    ///
27    /// ```text
28    /// BEFORE: _ r5 r4 r3 r2 r1 r0 l5 l4 l3 l2 l1 l0
29    /// AFTER:  _ sum_5 sum_4 sum_3 sum_2 sum_1 sum_0 is_overflow
30    /// ```
31    pub(crate) fn addition_code() -> Vec<LabelledInstruction> {
32        triton_asm!(
33            pick 6
34            // _ r5 r4 r3 r2 r1 l5 l4 l3 l2 l1 l0 r0
35
36            add
37            split
38            // _ r5 r4 r3 r2 r1 l5 l4 l3 l2 l1 (l0 + r0)_hi (l0 + r0)_lo
39            // _ r5 r4 r3 r2 r1 l5 l4 l3 l2 l1 carry_1 sum_0
40
41            swap 7
42            add
43            add
44            split
45            // _ r5 r4 r3 r2 sum_0 l5 l4 l3 l2 carry_2 sum_1
46
47            swap 7
48            add
49            add
50            split
51            // _ r5 r4 r3 sum_1 sum_0 l5 l4 l3 carry_3 sum_2
52
53            swap 7
54            add
55            add
56            split
57            // _ r5 r4 sum_2 sum_1 sum_0 l5 l4 carry_4 sum_3
58
59            swap 7
60            add
61            add
62            split
63            // _ r5 sum_3 sum_2 sum_1 sum_0 l5 carry_5 sum_4
64
65            swap 7
66            add
67            add
68            split
69            // _ sum_4 sum_3 sum_2 sum_1 sum_0 (is_overflow: bool) sum_5
70
71            place 6
72            // _ sum_5 sum_4 sum_3 sum_2 sum_1 sum_0 is_overflow
73
74            // _ [sum] overflow
75        )
76    }
77}
78
79impl BasicSnippet for OverflowingAdd {
80    fn parameters(&self) -> Vec<(DataType, String)> {
81        ["lhs", "rhs"]
82            .map(|s| (DataType::U192, s.to_owned()))
83            .to_vec()
84    }
85
86    fn return_values(&self) -> Vec<(DataType, String)> {
87        vec![
88            (DataType::U192, "sum".to_owned()),
89            (DataType::Bool, "overflow".to_owned()),
90        ]
91    }
92
93    fn entrypoint(&self) -> String {
94        "tasmlib_arithmetic_u192_overflowing_add".to_string()
95    }
96
97    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
98        triton_asm! { {self.entrypoint()}: {&Self::addition_code()} return }
99    }
100}
101
102#[cfg(test)]
103pub(crate) mod tests {
104    use num::BigUint;
105    use rand::rngs::StdRng;
106
107    use super::*;
108    use crate::arithmetic::u192::U192;
109    use crate::arithmetic::u192::to_u192;
110    use crate::arithmetic::u192::u128_to_u192;
111    use crate::arithmetic::u192::u128_to_u192_shl64;
112    use crate::test_prelude::*;
113
114    impl OverflowingAdd {
115        fn assert_expected_add_behavior(&self, lhs: U192, rhs: U192) {
116            let initial_stack = self.set_up_test_stack((rhs, lhs));
117
118            let mut expected_stack = initial_stack.clone();
119            self.rust_shadow(&mut expected_stack).unwrap();
120
121            test_rust_equivalence_given_complete_state(
122                &ShadowedClosure::new(Self),
123                &initial_stack,
124                &[],
125                &NonDeterminism::default(),
126                &None,
127                Some(&expected_stack),
128            );
129        }
130
131        pub fn edge_case_points() -> Vec<U192> {
132            [0, 0x200000002fffffffffff908f8, 1 << 127, u128::MAX]
133                .into_iter()
134                .flat_map(|p| [p.checked_sub(1), Some(p), p.checked_add(1)])
135                .flatten()
136                .map(u128_to_u192)
137                .chain([
138                    [u32::MAX; 6],
139                    [0, 0, 0, 0, 0, 1],
140                    [0, 0, 0, 0, 0, 1 << 31],
141                    [0, 0, 0, 0, 1 << 31, 0],
142                ])
143                .collect()
144        }
145    }
146
147    impl Closure for OverflowingAdd {
148        type Args = (U192, U192);
149
150        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
151            let left: U192 = pop_encodable(stack)?;
152            let left: BigUint = BigUint::new(left.to_vec());
153            let right: U192 = pop_encodable(stack)?;
154            let right: BigUint = BigUint::new(right.to_vec());
155            let sum = left + right;
156            let mut sum = sum.to_u32_digits();
157            let overflow_shape_is_valid =
158                sum.len() <= 6 || (sum.len() == 7 && sum.last().copied().is_some_and(|x| x == 1));
159            if !overflow_shape_is_valid {
160                return Err(RustShadowError::Other);
161            }
162
163            let is_overflow = sum.len() == 7;
164
165            sum.resize(6, 0);
166            let sum: U192 = sum.to_vec().try_into().unwrap();
167            push_encodable(stack, &sum);
168            push_encodable(stack, &is_overflow);
169            Ok(())
170        }
171
172        fn pseudorandom_args(&self, seed: [u8; 32], _: Option<BenchmarkCase>) -> Self::Args {
173            StdRng::from_seed(seed).random()
174        }
175
176        fn corner_case_args(&self) -> Vec<Self::Args> {
177            let edge_case_points = Self::edge_case_points();
178
179            edge_case_points
180                .iter()
181                .cartesian_product(&edge_case_points)
182                .map(|(&l, &r)| (l, r))
183                .collect()
184        }
185    }
186
187    #[macro_rules_attr::apply(test)]
188    fn rust_shadow() {
189        ShadowedClosure::new(OverflowingAdd).test()
190    }
191
192    #[macro_rules_attr::apply(test)]
193    fn unit_test() {
194        let snippet = OverflowingAdd;
195        snippet.assert_expected_add_behavior(u128_to_u192(1u128 << 67), u128_to_u192(1u128 << 67))
196    }
197
198    #[macro_rules_attr::apply(test)]
199    fn overflow_test() {
200        let test_overflowing_add = |a, b| {
201            OverflowingAdd.assert_expected_add_behavior(a, b);
202            OverflowingAdd.assert_expected_add_behavior(b, a);
203        };
204
205        let max = to_u192(u128::MAX, u64::MAX);
206        let max_minus_one = to_u192(u128::MAX, u64::MAX - 1);
207        test_overflowing_add(u128_to_u192(1), max);
208        test_overflowing_add(u128_to_u192(2), max_minus_one);
209        test_overflowing_add(u128_to_u192_shl64(1 << 127), u128_to_u192_shl64(1 << 127));
210        test_overflowing_add(max, max);
211
212        for a in [31, 32, 33, 63, 64, 65, 95, 96, 97].map(|p| 1 << p) {
213            test_overflowing_add(max, u128_to_u192(a));
214        }
215
216        for i in 0..128 {
217            let a = 1 << i;
218            let b = u128::MAX - a + 1;
219            debug_assert_eq!((0, true), a.overflowing_add(b), "i = {i}; a = {a}, b = {b}");
220
221            test_overflowing_add(u128_to_u192(a), u128_to_u192_shl64(b));
222        }
223    }
224}