feat: Add a limited form of arithmetic on generics (#5625)

# Description

## Problem\*

Resolves #4784

## Summary\*

Adds a limited form of arithmetic operators on numeric generic types in
Noir. This includes only 5 operations: +, -, *, /, and %.

Going forward I'm going to distinguish between two definitions for
clearer communication between devs & users:
- Numeric generics: generics that use numeric type variables like `fn
foo<let N: u32>()`
- Arithmetic generics: when arithmetic operators are used on numeric
generics in a type position: `fn foo<let N: u32>() -> [Field; N + 1]`

## Additional Context

It is explicitly a _non-goal_ to always correctly return when two
_arithmetic generics_ are equal. For example, in our current type
checking we do not check whether applying the distributive law would
make two types equal. So you'd get a type error for trying to use `2*(N
+ 1)` where `2*N + 2` was expected. Users should generally consider the
type equality here to be largely structural and try to match that as
best as possible. ~~The only check performed besides simplification of
constant terms is commutativity for `+` and `*`. Even then, the check is
only one level deep. So `N * M * 3` is not equal to `3 * N * M`.~~
Update: we canonicalize commutative operations now, so we sort this and
it now type checks.

## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [x] **[For Experimental Features]** Documentation to be submitted in a
separate PR.
- I'd like to have this in for some time before we start advertising it.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: Maxim Vezenov <[email protected]>

aztec_macros/src/utils/hir_utils.rs

@@ -195,7 +195,7 @@ pub fn inject_fn(
     let trait_id = None;
     items.functions.push(UnresolvedFunctions { file_id, functions, trait_id, self_type: None });
 
-    let mut errors = Elaborator::elaborate(context, *crate_id, items, None);
+    let mut errors = Elaborator::elaborate(context, *crate_id, items, None, false);
     errors.retain(|(error, _)| !CustomDiagnostic::from(error).is_warning());
 
     if !errors.is_empty() {
@@ -241,7 +241,7 @@ pub fn inject_global(
     let mut items = CollectedItems::default();
     items.globals.push(UnresolvedGlobal { file_id, module_id, global_id, stmt_def: global });
 
-    let _errors = Elaborator::elaborate(context, *crate_id, items, None);
+    let _errors = Elaborator::elaborate(context, *crate_id, items, None, false);
 }
 
 pub fn fully_qualified_note_path(context: &HirContext, note_id: StructId) -> Option<String> {

compiler/noirc_driver/src/abi_gen.rs

@@ -100,6 +100,7 @@ pub(super) fn abi_type_from_hir_type(context: &Context, typ: &Type) -> AbiType {
         Type::Error
         | Type::Unit
         | Type::Constant(_)
+        | Type::InfixExpr(..)
         | Type::TraitAsType(..)
         | Type::TypeVariable(_, _)
         | Type::NamedGeneric(..)

compiler/noirc_driver/src/lib.rs

@@ -113,6 +113,10 @@ pub struct CompileOptions {
     /// Outputs the paths to any modified artifacts
     #[arg(long, hide = true)]
     pub show_artifact_paths: bool,
+
+    /// Temporary flag to enable the experimental arithmetic generics feature
+    #[arg(long, hide = true)]
+    pub arithmetic_generics: bool,
 }
 
 pub fn parse_expression_width(input: &str) -> Result<ExpressionWidth, std::io::Error> {
@@ -262,21 +266,28 @@ pub fn add_dep(
 pub fn check_crate(
     context: &mut Context,
     crate_id: CrateId,
-    deny_warnings: bool,
-    disable_macros: bool,
-    debug_comptime_in_file: Option<&str>,
+    options: &CompileOptions,
 ) -> CompilationResult<()> {
-    let macros: &[&dyn MacroProcessor] =
-        if disable_macros { &[] } else { &[&aztec_macros::AztecMacro as &dyn MacroProcessor] };
+    let macros: &[&dyn MacroProcessor] = if options.disable_macros {
+        &[]
+    } else {
+        &[&aztec_macros::AztecMacro as &dyn MacroProcessor]
+    };
 
     let mut errors = vec![];
-    let diagnostics = CrateDefMap::collect_defs(crate_id, context, debug_comptime_in_file, macros);
+    let diagnostics = CrateDefMap::collect_defs(
+        crate_id,
+        context,
+        options.debug_comptime_in_file.as_deref(),
+        options.arithmetic_generics,
+        macros,
+    );
     errors.extend(diagnostics.into_iter().map(|(error, file_id)| {
         let diagnostic = CustomDiagnostic::from(&error);
         diagnostic.in_file(file_id)
     }));
 
-    if has_errors(&errors, deny_warnings) {
+    if has_errors(&errors, options.deny_warnings) {
         Err(errors)
     } else {
         Ok(((), errors))
@@ -302,13 +313,7 @@ pub fn compile_main(
     options: &CompileOptions,
     cached_program: Option<CompiledProgram>,
 ) -> CompilationResult<CompiledProgram> {
-    let (_, mut warnings) = check_crate(
-        context,
-        crate_id,
-        options.deny_warnings,
-        options.disable_macros,
-        options.debug_comptime_in_file.as_deref(),
-    )?;
+    let (_, mut warnings) = check_crate(context, crate_id, options)?;
 
     let main = context.get_main_function(&crate_id).ok_or_else(|| {
         // TODO(#2155): This error might be a better to exist in Nargo
@@ -343,13 +348,7 @@ pub fn compile_contract(
     crate_id: CrateId,
     options: &CompileOptions,
 ) -> CompilationResult<CompiledContract> {
-    let (_, warnings) = check_crate(
-        context,
-        crate_id,
-        options.deny_warnings,
-        options.disable_macros,
-        options.debug_comptime_in_file.as_deref(),
-    )?;
+    let (_, warnings) = check_crate(context, crate_id, options)?;
 
     // TODO: We probably want to error if contracts is empty
     let contracts = context.get_all_contracts(&crate_id);

compiler/noirc_driver/tests/stdlib_warnings.rs

@@ -25,7 +25,7 @@ fn stdlib_does_not_produce_constant_warnings() -> Result<(), ErrorsAndWarnings>
     let root_crate_id = prepare_crate(&mut context, file_name);
 
     let ((), warnings) =
-        noirc_driver::check_crate(&mut context, root_crate_id, false, false, None)?;
+        noirc_driver::check_crate(&mut context, root_crate_id, &Default::default())?;
 
     assert_eq!(warnings, Vec::new(), "stdlib is producing {} warnings", warnings.len());
 

compiler/noirc_frontend/src/ast/mod.rs

@@ -303,7 +303,7 @@ impl UnresolvedTypeData {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
+#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash, PartialOrd, Ord)]
 pub enum Signedness {
     Unsigned,
     Signed,

compiler/noirc_frontend/src/elaborator/comptime.rs

@@ -45,6 +45,7 @@ impl<'context> Elaborator<'context> {
             self.def_maps,
             self.crate_id,
             self.debug_comptime_in_file,
+            self.enable_arithmetic_generics,
         );
 
         elaborator.function_context.push(FunctionContext::default());

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -161,6 +161,9 @@ pub struct Elaborator<'context> {
     /// This map is used to lazily evaluate these globals if they're encountered before
     /// they are elaborated (e.g. in a function's type or another global's RHS).
     unresolved_globals: BTreeMap<GlobalId, UnresolvedGlobal>,
+
+    /// Temporary flag to enable the experimental arithmetic generics feature
+    enable_arithmetic_generics: bool,
 }
 
 #[derive(Default)]
@@ -183,6 +186,7 @@ impl<'context> Elaborator<'context> {
         def_maps: &'context mut DefMaps,
         crate_id: CrateId,
         debug_comptime_in_file: Option<FileId>,
+        enable_arithmetic_generics: bool,
     ) -> Self {
         Self {
             scopes: ScopeForest::default(),
@@ -203,19 +207,22 @@ impl<'context> Elaborator<'context> {
             current_trait_impl: None,
             debug_comptime_in_file,
             unresolved_globals: BTreeMap::new(),
+            enable_arithmetic_generics,
         }
     }
 
     pub fn from_context(
         context: &'context mut Context,
         crate_id: CrateId,
         debug_comptime_in_file: Option<FileId>,
+        enable_arithmetic_generics: bool,
     ) -> Self {
         Self::new(
             &mut context.def_interner,
             &mut context.def_maps,
             crate_id,
             debug_comptime_in_file,
+            enable_arithmetic_generics,
         )
     }
 
@@ -224,17 +231,31 @@ impl<'context> Elaborator<'context> {
         crate_id: CrateId,
         items: CollectedItems,
         debug_comptime_in_file: Option<FileId>,
+        enable_arithmetic_generics: bool,
     ) -> Vec<(CompilationError, FileId)> {
-        Self::elaborate_and_return_self(context, crate_id, items, debug_comptime_in_file).errors
+        Self::elaborate_and_return_self(
+            context,
+            crate_id,
+            items,
+            debug_comptime_in_file,
+            enable_arithmetic_generics,
+        )
+        .errors
     }
 
     pub fn elaborate_and_return_self(
         context: &'context mut Context,
         crate_id: CrateId,
         items: CollectedItems,
         debug_comptime_in_file: Option<FileId>,
+        enable_arithmetic_generics: bool,
     ) -> Self {
-        let mut this = Self::from_context(context, crate_id, debug_comptime_in_file);
+        let mut this = Self::from_context(
+            context,
+            crate_id,
+            debug_comptime_in_file,
+            enable_arithmetic_generics,
+        );
         this.elaborate_items(items);
         this.check_and_pop_function_context();
         this

compiler/noirc_frontend/src/elaborator/types.rs

@@ -403,13 +403,16 @@ impl<'context> Elaborator<'context> {
 
                 match (lhs, rhs) {
                     (Type::Constant(lhs), Type::Constant(rhs)) => {
-                        Type::Constant(op.function()(lhs, rhs))
+                        Type::Constant(op.function(lhs, rhs))
                     }
-                    (lhs, _) => {
-                        let span =
-                            if !matches!(lhs, Type::Constant(_)) { lhs_span } else { rhs_span };
-                        self.push_err(ResolverError::InvalidArrayLengthExpr { span });
-                        Type::Constant(0)
+                    (lhs, rhs) => {
+                        if !self.enable_arithmetic_generics {
+                            let span =
+                                if !matches!(lhs, Type::Constant(_)) { lhs_span } else { rhs_span };
+                            self.push_err(ResolverError::InvalidArrayLengthExpr { span });
+                        }
+
+                        Type::InfixExpr(Box::new(lhs), op, Box::new(rhs)).canonicalize()
                     }
                 }
             }
@@ -1614,6 +1617,10 @@ impl<'context> Elaborator<'context> {
                 }
                 Self::find_numeric_generics_in_type(fields, found);
             }
+            Type::InfixExpr(lhs, _op, rhs) => {
+                Self::find_numeric_generics_in_type(lhs, found);
+                Self::find_numeric_generics_in_type(rhs, found);
+            }
         }
     }
 

compiler/noirc_frontend/src/hir/comptime/hir_to_display_ast.rs

@@ -358,6 +358,13 @@ impl Type {
             Type::Constant(_) => panic!("Type::Constant where a type was expected: {self:?}"),
             Type::Quoted(quoted_type) => UnresolvedTypeData::Quoted(*quoted_type),
             Type::Error => UnresolvedTypeData::Error,
+            Type::InfixExpr(lhs, op, rhs) => {
+                let lhs = Box::new(lhs.to_type_expression());
+                let rhs = Box::new(rhs.to_type_expression());
+                let span = Span::default();
+                let expr = UnresolvedTypeExpression::BinaryOperation(lhs, *op, rhs, span);
+                UnresolvedTypeData::Expression(expr)
+            }
         };
 
         UnresolvedType { typ, span: None }

compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs

@@ -662,6 +662,7 @@ fn zeroed(return_type: Type) -> IResult<Value> {
         Type::TypeVariable(_, _)
         | Type::Forall(_, _)
         | Type::Constant(_)
+        | Type::InfixExpr(..)
         | Type::Quoted(_)
         | Type::Error
         | Type::TraitAsType(_, _, _)

compiler/noirc_frontend/src/hir/comptime/tests.rs

@@ -45,7 +45,7 @@ fn interpret_helper(src: &str) -> Result<Value, InterpreterError> {
 
     let main = context.get_main_function(&krate).expect("Expected 'main' function");
     let mut elaborator =
-        Elaborator::elaborate_and_return_self(&mut context, krate, collector.items, None);
+        Elaborator::elaborate_and_return_self(&mut context, krate, collector.items, None, false);
     assert_eq!(elaborator.errors.len(), 0);
 
     let mut interpreter = elaborator.setup_interpreter();

compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs

@@ -247,6 +247,7 @@ impl DefCollector {
         ast: SortedModule,
         root_file_id: FileId,
         debug_comptime_in_file: Option<&str>,
+        enable_arithmetic_generics: bool,
         macro_processors: &[&dyn MacroProcessor],
     ) -> Vec<(CompilationError, FileId)> {
         let mut errors: Vec<(CompilationError, FileId)> = vec![];
@@ -264,6 +265,7 @@ impl DefCollector {
                 dep.crate_id,
                 context,
                 debug_comptime_in_file,
+                enable_arithmetic_generics,
                 macro_processors,
             ));
 
@@ -386,8 +388,14 @@ impl DefCollector {
             })
         });
 
-        let mut more_errors =
-            Elaborator::elaborate(context, crate_id, def_collector.items, debug_comptime_in_file);
+        let mut more_errors = Elaborator::elaborate(
+            context,
+            crate_id,
+            def_collector.items,
+            debug_comptime_in_file,
+            enable_arithmetic_generics,
+        );
+
         errors.append(&mut more_errors);
 
         for macro_processor in macro_processors {

compiler/noirc_frontend/src/hir/def_map/mod.rs

@@ -76,6 +76,7 @@ impl CrateDefMap {
         crate_id: CrateId,
         context: &mut Context,
         debug_comptime_in_file: Option<&str>,
+        enable_arithmetic_generics: bool,
         macro_processors: &[&dyn MacroProcessor],
     ) -> Vec<(CompilationError, FileId)> {
         // Check if this Crate has already been compiled
@@ -125,6 +126,7 @@ impl CrateDefMap {
             ast,
             root_file_id,
             debug_comptime_in_file,
+            enable_arithmetic_generics,
             macro_processors,
         ));