chore: Add missing cases to arithmetic generics

compiler/noirc_frontend/src/elaborator/types.rs

@@ -445,14 +445,19 @@ impl<'context> Elaborator<'context> {
                 })
             }
             UnresolvedTypeExpression::Constant(int, _) => Type::Constant(int),
-            UnresolvedTypeExpression::BinaryOperation(lhs, op, rhs, _) => {
+            UnresolvedTypeExpression::BinaryOperation(lhs, op, rhs, span) => {
                 let (lhs_span, rhs_span) = (lhs.span(), rhs.span());
                 let lhs = self.convert_expression_type(*lhs);
                 let rhs = self.convert_expression_type(*rhs);
 
                 match (lhs, rhs) {
                     (Type::Constant(lhs), Type::Constant(rhs)) => {
-                        Type::Constant(op.function(lhs, rhs))
+                        if let Some(result) = op.function(lhs, rhs) {
+                            Type::Constant(result)
+                        } else {
+                            self.push_err(ResolverError::OverflowInType { lhs, op, rhs, span });
+                            Type::Error
+                        }
                     }
                     (lhs, rhs) => {
                         if !self.enable_arithmetic_generics {

compiler/noirc_frontend/src/hir/resolution/errors.rs

@@ -120,6 +120,8 @@ pub enum ResolverError {
     NamedTypeArgs { span: Span, item_kind: &'static str },
     #[error("Associated constants may only be a field or integer type")]
     AssociatedConstantsMustBeNumeric { span: Span },
+    #[error("Overflow in `{lhs} {op} {rhs}`")]
+    OverflowInType { lhs: u32, op: crate::BinaryTypeOperator, rhs: u32, span: Span },
 }
 
 impl ResolverError {
@@ -480,6 +482,13 @@ impl<'a> From<&'a ResolverError> for Diagnostic {
                     *span,
                 )
             }
+            ResolverError::OverflowInType { lhs, op, rhs, span } => {
+                Diagnostic::simple_error(
+                    format!("Overflow in `{lhs} {op} {rhs}`"),
+                    "Overflow here".to_string(),
+                    *span,
+                )
+            }
         }
     }
 }

compiler/noirc_frontend/src/hir_def/types.rs

@@ -1672,16 +1672,20 @@
                 // evaluate_to_u32 also calls canonicalize so if we just called
                 // `self.evaluate_to_u32()` we'd get infinite recursion.
                 if let (Some(lhs), Some(rhs)) = (lhs.evaluate_to_u32(), rhs.evaluate_to_u32()) {
-                    return Type::Constant(op.function(lhs, rhs));
+                    if let Some(result) = op.function(lhs, rhs) {
+                        return Type::Constant(result);
+                    }
                 }
 
                 let lhs = lhs.canonicalize();
                 let rhs = rhs.canonicalize();
-                if let Some(result) = Self::try_simplify_addition(&lhs, op, &rhs) {
+                if let Some(result) = Self::try_simplify_non_constants(&lhs, op, &rhs) {
                     return result;
                 }
 
-                if let Some(result) = Self::try_simplify_subtraction(&lhs, op, &rhs) {
+                // Try to simplify partially constant expressions in the form `(N op1 C1) op2 C2`
+                // where C1 and C2 are constants that can be combined (e.g. N + 5 - 3 = N + 2)
+                if let Some(result) = Self::try_simplify_partial_constants(&lhs, op, &rhs) {
                     return result;
                 }
 
@@ -1711,7 +1715,11 @@
                     queue.push(*rhs);
                 }
                 Type::Constant(new_constant) => {
-                    constant = op.function(constant, new_constant);
+                    if let Some(result) = op.function(constant, new_constant) {
+                        constant = result;
+                    } else {
+                        sorted.insert(Type::Constant(new_constant));
+                    }
                 }
                 other => {
                     sorted.insert(other);
@@ -1737,47 +1745,78 @@
         }
     }
 
-    /// Try to simplify an addition expression of `lhs + rhs`.
+    /// Try to simplify non-constant expressions in the form `(N op1 M) op2 M`
+    /// where the two `M` terms are expected to cancel out.
     ///
-    /// - Simplifies `(a - b) + b` to `a`.
-    fn try_simplify_addition(lhs: &Type, op: BinaryTypeOperator, rhs: &Type) -> Option<Type> {
-        use BinaryTypeOperator::*;
-        match lhs {
-            Type::InfixExpr(l_lhs, l_op, l_rhs) => {
-                if op == Addition && *l_op == Subtraction {
-                    // TODO: Propagate type bindings. Can do in another PR, this one is large enough.
-                    let unifies = l_rhs.try_unify(rhs, &mut TypeBindings::new());
-                    if unifies.is_ok() {
-                        return Some(l_lhs.as_ref().clone());
-                    }
-                }
-                None
-            }
-            _ => None,
+    /// - Simplifies `(N +/- M) -/+ M` to `a`
+    /// - Simplifies `(N */÷ M) ÷/* M` to `a`
+    fn try_simplify_non_constants(lhs: &Type, op: BinaryTypeOperator, rhs: &Type) -> Option<Type> {
+        let Type::InfixExpr(l_type, l_op, l_rhs) = lhs.follow_bindings() else {
+            return None;
+        };
+
+        // Note that this is exact, syntactic equality, not unification.
+        if l_op.inverse() != Some(op) || l_rhs.canonicalize() != rhs.canonicalize() {
+            return None;
         }
+
+        Some(*l_type)
+    }
+
+    /// Given:
+    ///   lhs = `N op C1`
+    ///   rhs = C2
+    /// Returns: `(N, op, C1, C2)` if C1 and C2 are constants.
+    ///   Note that the operator here is within the `lhs` term, the operator
+    ///   separating lhs and rhs is not needed.
+    fn parse_constant_expr(
+        lhs: &Type,
+        rhs: &Type,
+    ) -> Option<(Box<Type>, BinaryTypeOperator, u32, u32)> {
+        let rhs = rhs.evaluate_to_u32()?;
+
+        let Type::InfixExpr(l_type, l_op, l_rhs) = lhs.follow_bindings() else {
+            return None;
+        };
+
+        let l_rhs = l_rhs.evaluate_to_u32()?;
+        Some((l_type, l_op, l_rhs, rhs))
     }
 
-    /// Try to simplify a subtraction expression of `lhs - rhs`.
+    /// Try to simplify partially constant expressions in the form `(N op1 C1) op2 C2`
+    /// where C1 and C2 are constants that can be combined (e.g. N + 5 - 3 = N + 2)
     ///
-    /// - Simplifies `(a + C1) - C2` to `a + (C1 - C2)` if C1 and C2 are constants.
-    fn try_simplify_subtraction(lhs: &Type, op: BinaryTypeOperator, rhs: &Type) -> Option<Type> {
+    /// - Simplifies `(N +/- C1) +/- C2` to `N +/- (C1 +/- C2)` if C1 and C2 are constants.
+    /// - Simplifies `(N */÷ C1) */÷ C2` to `N */÷ (C1 */÷ C2)` if C1 and C2 are constants.
+    fn try_simplify_partial_constants(
+        lhs: &Type,
+        mut op: BinaryTypeOperator,
+        rhs: &Type,
+    ) -> Option<Type> {
         use BinaryTypeOperator::*;
-        match lhs {
-            Type::InfixExpr(l_lhs, l_op, l_rhs) => {
-                // Simplify `(N + 2) - 1`
-                if op == Subtraction && *l_op == Addition {
-                    if let (Some(lhs_const), Some(rhs_const)) =
-                        (l_rhs.evaluate_to_u32(), rhs.evaluate_to_u32())
-                    {
-                        if lhs_const > rhs_const {
-                            let constant = Box::new(Type::Constant(lhs_const - rhs_const));
-                            return Some(
-                                Type::InfixExpr(l_lhs.clone(), *l_op, constant).canonicalize(),
-                            );
-                        }
-                    }
+        let (l_type, l_op, l_const, r_const) = Type::parse_constant_expr(lhs, rhs)?;
+
+        match (l_op, op) {
+            (Addition | Subtraction, Addition | Subtraction) => {
+                // If l_op is a subtraction we want to inverse the rhs operator.
+                if l_op == Subtraction {
+                    op = op.inverse()?;
+                }
+                let result = op.function(l_const, r_const)?;
+                Some(Type::InfixExpr(l_type, l_op, Box::new(Type::Constant(result))))
+            }
+            (Multiplication | Division, Multiplication | Division) => {
+                // If l_op is a subtraction we want to inverse the rhs operator.
+                if l_op == Division {
+                    op = op.inverse()?;
+                }
+                // If op is a division we need to ensure it divides evenly
+                if op == Division && l_const % r_const != 0 {
+                    None
+                } else {
+                    let result = op.function(l_const, r_const)?;
+                    Some(Type::InfixExpr(l_type, l_op, Box::new(Type::Constant(result))))
                 }
-                None
             }
             _ => None,
         }
@@ -1926,7 +1965,7 @@
             Type::InfixExpr(lhs, op, rhs) => {
                 let lhs = lhs.evaluate_to_u32()?;
                 let rhs = rhs.evaluate_to_u32()?;
-                Some(op.function(lhs, rhs))
+                op.function(lhs, rhs)
             }
             _ => None,
         }
@@ -2030,17 +2069,13 @@
             Type::Forall(typevars, typ) => {
                 assert_eq!(types.len() + implicit_generic_count, typevars.len(), "Turbofish operator used with incorrect generic count which was not caught by name resolution");
 
+                let bindings =
+                    (0..implicit_generic_count).map(|_| interner.next_type_variable()).chain(types);
+
                 let replacements = typevars
                     .iter()
-                    .enumerate()
-                    .map(|(i, var)| {
-                        let binding = if i < implicit_generic_count {
-                            interner.next_type_variable()
-                        } else {
-                            types[i - implicit_generic_count].clone()
-                        };
-                        (var.id(), (var.clone(), binding))
-                    })
+                    .zip(bindings)
+                    .map(|(var, binding)| (var.id(), (var.clone(), binding)))
                     .collect();
 
                 let instantiated = typ.substitute(&replacements);
@@ -2094,7 +2129,7 @@
         }
 
         let recur_on_binding = |id, replacement: &Type| {
             // Prevent recuring forever if there's a `T := T` binding
             if replacement.type_variable_id() == Some(id) {
                 replacement.clone()
             } else {
@@ -2165,7 +2200,7 @@
                 Type::Tuple(fields)
             }
             Type::Forall(typevars, typ) => {
                 // Trying to substitute_helper a variable de, substitute_bound_typevarsfined within a nested Forall
                 // is usually impossible and indicative of an error in the type checker somewhere.
                 for var in typevars {
                     assert!(!type_bindings.contains_key(&var.id()));
@@ -2332,7 +2367,7 @@
 
     /// Replace any `Type::NamedGeneric` in this type with a `Type::TypeVariable`
     /// using to the same inner `TypeVariable`. This is used during monomorphization
     /// to bind to named generics since they are unbindable during type checking.
     pub fn replace_named_generics_with_type_variables(&mut self) {
         match self {
             Type::FieldElement
@@ -2457,13 +2492,13 @@
 
 impl BinaryTypeOperator {
     /// Perform the actual rust numeric operation associated with this operator
-    pub fn function(self, a: u32, b: u32) -> u32 {
+    pub fn function(self, a: u32, b: u32) -> Option<u32> {
         match self {
-            BinaryTypeOperator::Addition => a.wrapping_add(b),
-            BinaryTypeOperator::Subtraction => a.wrapping_sub(b),
-            BinaryTypeOperator::Multiplication => a.wrapping_mul(b),
-            BinaryTypeOperator::Division => a.wrapping_div(b),
-            BinaryTypeOperator::Modulo => a.wrapping_rem(b),
+            BinaryTypeOperator::Addition => a.checked_add(b),
+            BinaryTypeOperator::Subtraction => a.checked_sub(b),
+            BinaryTypeOperator::Multiplication => a.checked_mul(b),
+            BinaryTypeOperator::Division => a.checked_div(b),
+            BinaryTypeOperator::Modulo => a.checked_rem(b),
         }
     }
 

test_programs/compile_success_empty/arithmetic_generics/src/main.nr

@@ -7,6 +7,9 @@ fn main() {
     let _ = split_first([1, 2, 3]);
 
     let _ = push_multiple([1, 2, 3]);
+
+    test_constant_folding::<10>();
+    test_non_constant_folding::<10, 20>();
 }
 
 fn split_first<T, let N: u32>(array: [T; N]) -> (T, [T; N - 1]) {
@@ -101,3 +104,31 @@ fn demo_proof<let N: u32>() -> Equiv<W<(N * (N + 1))>, (Equiv<W<N>, (), W<N>, ()
     let p3: Equiv<W<N * N + N>, (), W<N * N + N>, ()> = add_equiv_r::<N * N, N, N, _, _>(p3_sub);
     equiv_trans(equiv_trans(p1, p2), p3)
 }
+
+fn test_constant_folding<let N: u32>() {
+    // N + C1 - C2 = N + (C1 - C2)
+    let _: W<N + 5 - 2> = W::<N + 3> {};
+
+    // N - C1 + C2 = N - (C1 - C2)
+    let _: W<N - 3 + 2> = W::<N - 1> {};
+
+    // N * C1 / C2 = N * (C1 / C2)
+    let _: W<N * 10 / 2> = W::<N * 5> {};
+
+    // N / C1 * C2 = N / (C1 / C2)
+    let _: W<N / 10 * 2> = W::<N / 5> {};
+}
+
+fn test_non_constant_folding<let N: u32, let M: u32>() {
+    // N + M - M = N
+    let _: W<N + M - M> = W::<N> {};
+
+    // N - M + M = N
+    let _: W<N - M + M> = W::<N> {};
+
+    // N * M / M = N
+    let _: W<N * M / M> = W::<N> {};
+
+    // N / M * M = N
+    let _: W<N / M * M> = W::<N> {};
+}