tasm_lib/arithmetic/u160/
overflowing_add.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5/// Mimics overflowing_add but works on u160 data types.
6///
7/// ### Behavior
8///
9/// ```text
10/// BEFORE: _ [rhs: u160] [lhs: u160]
11/// AFTER:  _ [sum: u160] [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 `u160`s.
26    ///
27    /// ```text
28    /// BEFORE: _ rhs_4 rhs_3 rhs_2 rhs_1 rhs_0 lhs_4 lhs_3 lhs_2 lhs_1 lhs_0
29    /// AFTER:  _ 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 5
34            // _ rhs_4 rhs_3 rhs_2 rhs_1 lhs_4 lhs_3 lhs_2 lhs_1 lhs_0 rhs_0
35
36            add
37            split
38            // _ rhs_4 rhs_3 rhs_2 rhs_1 lhs_4 lhs_3 lhs_2 lhs_1 (lhs_0 + rhs_0)_hi (lhs_0 + rhs_0)_lo
39            // _ rhs_4 rhs_3 rhs_2 rhs_1 lhs_4 lhs_3 lhs_2 lhs_1 carry_1 sum_0
40
41            swap 6
42            add
43            add
44            split
45            // _ rhs_4 rhs_3 rhs_2 sum_0 lhs_4 lhs_3 lhs_2 carry_2 sum_1
46
47            swap 6
48            add
49            add
50            split
51            // _ rhs_4 rhs_3 sum_1 sum_0 lhs_4 lhs_3 carry_3 sum_2
52
53            swap 6
54            add
55            add
56            split
57            // _ rhs_4 sum_2 sum_1 sum_0 lhs_4 carry_4 sum_3
58
59            swap 6
60            add
61            add
62            split
63            // _ sum_3 sum_2 sum_1 sum_0 is_overflow sum_4
64
65            place 5
66            // _ sum_4 sum_3 sum_2 sum_1 sum_0 is_overflow
67
68            // _ [sum] overflow
69        )
70    }
71}
72
73impl BasicSnippet for OverflowingAdd {
74    fn parameters(&self) -> Vec<(DataType, String)> {
75        ["lhs", "rhs"]
76            .map(|s| (DataType::U160, s.to_owned()))
77            .to_vec()
78    }
79
80    fn return_values(&self) -> Vec<(DataType, String)> {
81        vec![
82            (DataType::U160, "sum".to_owned()),
83            (DataType::Bool, "overflow".to_owned()),
84        ]
85    }
86
87    fn entrypoint(&self) -> String {
88        "tasmlib_arithmetic_u160_overflowing_add".to_string()
89    }
90
91    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
92        triton_asm! { {self.entrypoint()}: {&Self::addition_code()} return }
93    }
94}
95
96#[cfg(test)]
97pub(crate) mod tests {
98    use num::BigUint;
99    use rand::rngs::StdRng;
100
101    use super::*;
102    use crate::arithmetic::u160::u128_to_u160;
103    use crate::arithmetic::u160::u128_to_u160_shl_32;
104    use crate::test_prelude::*;
105
106    impl OverflowingAdd {
107        fn assert_expected_add_behavior(&self, lhs: [u32; 5], rhs: [u32; 5]) {
108            let initial_stack = self.set_up_test_stack((rhs, lhs));
109
110            let mut expected_stack = initial_stack.clone();
111            self.rust_shadow(&mut expected_stack);
112
113            test_rust_equivalence_given_complete_state(
114                &ShadowedClosure::new(Self),
115                &initial_stack,
116                &[],
117                &NonDeterminism::default(),
118                &None,
119                Some(&expected_stack),
120            );
121        }
122
123        pub fn edge_case_points() -> Vec<[u32; 5]> {
124            [0, 0x200000002fffffffffff908f8, 1 << 127, u128::MAX]
125                .into_iter()
126                .flat_map(|p| [p.checked_sub(1), Some(p), p.checked_add(1)])
127                .flatten()
128                .map(u128_to_u160)
129                .chain([[u32::MAX; 5], [0, 0, 0, 0, 1]])
130                .collect()
131        }
132    }
133
134    impl Closure for OverflowingAdd {
135        type Args = ([u32; 5], [u32; 5]);
136
137        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
138            let left: [u32; 5] = pop_encodable(stack);
139            let left: BigUint = BigUint::new(left.to_vec());
140            let right: [u32; 5] = pop_encodable(stack);
141            let right: BigUint = BigUint::new(right.to_vec());
142            let sum = left + right;
143            let mut sum = sum.to_u32_digits();
144            assert!(
145                sum.len() <= 5 || sum.len() == 6 && *sum.last().unwrap() == 1,
146                "Value must be bounded thusly"
147            );
148
149            let is_overflow = sum.len() == 6;
150
151            sum.resize(5, 0);
152            let sum: [u32; 5] = sum.to_vec().try_into().unwrap();
153            push_encodable(stack, &sum);
154            push_encodable(stack, &is_overflow);
155        }
156
157        fn pseudorandom_args(&self, seed: [u8; 32], _: Option<BenchmarkCase>) -> Self::Args {
158            StdRng::from_seed(seed).random()
159        }
160
161        fn corner_case_args(&self) -> Vec<Self::Args> {
162            let edge_case_points = Self::edge_case_points();
163
164            edge_case_points
165                .iter()
166                .cartesian_product(&edge_case_points)
167                .map(|(&l, &r)| (l, r))
168                .collect()
169        }
170    }
171
172    #[test]
173    fn rust_shadow() {
174        ShadowedClosure::new(OverflowingAdd).test();
175    }
176
177    #[test]
178    fn unit_test() {
179        let snippet = OverflowingAdd;
180        snippet.assert_expected_add_behavior(u128_to_u160(1u128 << 67), u128_to_u160(1u128 << 67))
181    }
182
183    #[test]
184    fn overflow_test() {
185        let test_overflowing_add = |a, b| {
186            OverflowingAdd.assert_expected_add_behavior(a, b);
187            OverflowingAdd.assert_expected_add_behavior(b, a);
188        };
189
190        test_overflowing_add([1, 0, 0, 0, 0], [u32::MAX; 5]);
191        test_overflowing_add(
192            [2, 0, 0, 0, 0],
193            [u32::MAX - 1, u32::MAX, u32::MAX, u32::MAX, u32::MAX],
194        );
195        test_overflowing_add([1 << 31, 0, 0, 0, 0], [1 << 31, 0, 0, 0, 0]);
196        test_overflowing_add([u32::MAX; 5], [u32::MAX; 5]);
197
198        for a in [31, 32, 33, 63, 64, 65, 95, 96, 97].map(|p| 1 << p) {
199            test_overflowing_add([u32::MAX; 5], u128_to_u160(a));
200        }
201
202        for i in 0..128 {
203            let a = 1 << i;
204            let b = u128::MAX - a + 1;
205            debug_assert_eq!((0, true), a.overflowing_add(b), "i = {i}; a = {a}, b = {b}");
206
207            test_overflowing_add(u128_to_u160(a), u128_to_u160_shl_32(b));
208        }
209    }
210}