add tests, make uint optimizer complete§

zokrates_core/src/static_analysis/uint_optimizer.rs

@@ -284,35 +284,56 @@ impl<'ast, T: Field> Folder<'ast, T> for UintOptimizer<'ast, T> {
                 let e = self.fold_uint_expression(e);
                 let by = self.fold_uint_expression(by);
 
-                let by_u = match by.as_inner() {
+                let by_max: usize = by
-                    UExpressionInner::Value(ref by) => by,
+                    .metadata
-                    _ => unreachable!(),
+                    .clone()
-                };
+                    .unwrap()
+                    .max
+                    .to_dec_string()
+                    .parse()
+                    .unwrap();
+                let e_max: usize = e
+                    .metadata
+                    .clone()
+                    .unwrap()
+                    .max
+                    .to_dec_string()
+                    .parse()
+                    .unwrap();
 
-                let bitwidth = e.metadata.clone().unwrap().bitwidth();
+                let max = T::from((e_max << by_max) & (2_usize.pow(range as u32) - 1));
 
-                let max = T::from(2).pow(std::cmp::min(bitwidth as usize + *by_u as usize, range))
+                UExpression::left_shift(force_reduce(e), force_reduce(by)).with_max(max)
-                    - T::from(1);
-
-                UExpression::left_shift(force_reduce(e), by).with_max(max)
             }
             RightShift(box e, box by) => {
-                // reduce the left term
+                // reduce both terms
                 let e = self.fold_uint_expression(e);
                 let by = self.fold_uint_expression(by);
 
+                // if we don't know the amount by which we shift, the most conservative case (which leads to the biggest value) is 0
                 let by_u = match by.as_inner() {
-                    UExpressionInner::Value(ref by) => by,
+                    UExpressionInner::Value(by) => *by,
-                    _ => unreachable!(),
+                    _ => 0,
                 };
 
-                let bitwidth = e.metadata.clone().unwrap().bitwidth();
+                let e_max: usize = e
+                    .metadata
+                    .clone()
+                    .unwrap()
+                    .max
+                    .to_dec_string()
+                    .parse()
+                    .unwrap();
 
-                let max = T::from(2)
+                println!("e_max: {:x?}, by_u: {}", e_max, by_u);
-                    .pow(bitwidth as usize - std::cmp::min(*by_u as usize, bitwidth as usize))
-                    - T::from(1);
 
-                UExpression::right_shift(force_reduce(e), by).with_max(max)
+                let max = (e_max & (2_usize.pow(range as u32) - 1)) >> by_u;
+
+                println!("MAX {:x?}", max);
+
+                let max = T::from(max);
+
+                UExpression::right_shift(force_reduce(e), force_reduce(by)).with_max(max)
             }
             IfElse(box condition, box consequence, box alternative) => {
                 let condition = self.fold_boolean_expression(condition);
@@ -465,19 +486,28 @@ mod tests {
 
     use pretty_assertions::assert_eq;
 
+    /// generate a test for a binary operator
+    ///
+    /// # Arguments
+    ///
+    /// * `left_max` the max value of the left argument
+    /// * `expected_reduce_left` whether we expect the optimizer to reduce the left argument
+    /// * `right_max` the max value of the right argument
+    /// * `expected_reduce_right` whether we expect the optimizer to reduce the left argument
+    /// * `res_max` the expected max value of the output
     macro_rules! uint_test {
-        ( $left_max:expr, $left_reduce:expr, $right_max:expr, $right_reduce:expr, $method:ident, $res_max:expr  ) => {{
+        ( $left_max:expr, $expected_reduce_left:expr, $right_max:expr, $expected_reduce_right:expr, $op:ident, $res_max:expr ) => {{
             let left = e_with_max($left_max);
 
             let right = e_with_max($right_max);
 
-            let left_expected = if $left_reduce {
+            let left_expected = if $expected_reduce_left {
                 force_reduce(left.clone())
             } else {
                 force_no_reduce(left.clone())
             };
 
-            let right_expected = if $right_reduce {
+            let right_expected = if $expected_reduce_right {
                 force_reduce(right.clone())
             } else {
                 force_no_reduce(right.clone())
@@ -485,8 +515,8 @@ mod tests {
 
             assert_eq!(
                 UintOptimizer::new()
-                    .fold_uint_expression(UExpression::$method(left.clone(), right.clone())),
+                    .fold_uint_expression(UExpression::$op(left.clone(), right.clone())),
-                UExpression::$method(left_expected, right_expected).with_max($res_max)
+                UExpression::$op(left_expected, right_expected).with_max($res_max)
             );
         }};
     }
@@ -617,16 +647,64 @@ mod tests {
 
     #[test]
     fn right_shift() {
-        uint_test!(255, true, 2, true, right_shift, 63_u32);
+        // left argument in range, we reduce (no effect) and the max is the original max, as we could be shifting by 0
-        uint_test!(2, true, 2, true, right_shift, 0_u32);
+        uint_test!(0xff_u32, true, 2, true, right_shift, 0xff_u32);
-        unimplemented!("non reduction");
+        uint_test!(2, true, 2, true, right_shift, 2_u32);
+
+        // left argument out of range, we reduce and the max is the type max, shifted
+        uint_test!(
+            0xffffffffffff_u128,
+            true,
+            2,
+            true,
+            right_shift,
+            0xffffffff_u32
+        );
+
+        fn right_shift_test(e_max: u128, by: u32, output_max: u32) {
+            let left = e_with_max(e_max);
+
+            let right = UExpressionInner::Value(by as u128)
+                .annotate(crate::zir::types::UBitwidth::B32)
+                .with_max(by);
+
+            let left_expected = force_reduce(left.clone());
+
+            let right_expected = force_reduce(right.clone());
+
+            assert_eq!(
+                UintOptimizer::new()
+                    .fold_uint_expression(UExpression::right_shift(left.clone(), right.clone())),
+                UExpression::right_shift(left_expected, right_expected).with_max(output_max)
+            );
+        }
+
+        right_shift_test(0xff_u128, 2, 0xff >> 2);
+        right_shift_test(2, 2, 2 >> 2);
+        right_shift_test(0xffffffffffff_u128, 2, 0xffffffff >> 2);
     }
 
     #[test]
     fn left_shift() {
-        uint_test!(255, true, 2, true, right_shift, 1023_u32);
+        uint_test!(0xff_u32, true, 2, true, left_shift, 0xff_u32 << 2);
-        uint_test!(0xffffffff_u32, true, 2, true, right_shift, 0xffffffff_u32);
+        uint_test!(
-        unimplemented!("non reduction");
+            0xffffffff_u32,
+            true,
+            2,
+            true,
+            left_shift,
+            0xffffffff_u32 << 2
+        );
+
+        // left argument out of range, we reduce and the max is the type max, shifted
+        uint_test!(
+            0xffffffffffff_u128,
+            true,
+            2,
+            true,
+            left_shift,
+            0xffffffff_u32 << 2
+        )
     }
 
     #[test]