feat: Explicit Associated Types & Constants (#5739)

# Description

## Problem\*

Resolves #3470

(does not fix #5544)

## Summary\*

Implements a basic form of associated types in Noir. 
- The main limitation of this work is that all associated types must be
explicitly specified in trait constraints (e.g. `Foo: Iterator<Elem =
i32>`) Instead of `Foo: Iterator`
- Eliding these requires inserting additional implicit generics anywhere
these constraints are used which I considered for this PR but opted to
limit the size instead and leave it for later work.

## Additional Context

This is a large PR which has quite a few non-trivial changes so I'll try
to summarize as best I can to help with reviews:
- Made the `Span` on `UnresolvedType` non-optional for better error
messages.
- Added `GenericTypeArgs` struct to hold both "ordered" and "named"
generics. Ordered generics are the ones we are used to where they are
specified in order. Named generics are associated types where they are
prefixed with `Name =` and may be given in any order in a generics list.
- Added a `Generic` trait to abstract over something that can be generic
(struct, type alias, trait, etc) so that we can share the same method
for resolving all generic things. This saved some code since we were
duplicating checks on these before, and have much more now that we have
to check named args as well.
- Added `TraitGenerics` as the resolved version of `GenericTypeArgs`.
This is used in a couple places including the `impl Trait` type in the
type system.
- Added associated types to the `lookup_trait_implementation` functions
in the NodeInterner. These are just treated as normal generics more or
less. When we allow these to be omitted in the future we should probably
fill these with fresh type variables so that they are still present but
always succeed unification.
- Stored associated types in trait impls in a separate map in the
NodeInterner since they are needed to be referenced within the type
signature of methods in the impl, which is before the impl as a whole
finishes resolution.
- Resolved `Self::AssociatedType` paths. Not that
#5544 is not fixed,
`Trait::AssociatedType` is still a panic.
- Added `current_trait` and `current_trait_impl` fields to the
elaborator to resolve these. Otherwise we don't know what `Self` refers
to.
- Resolved `<Type as Trait>::Type` paths. These are called
`AsTraitPath`s. Caveat: these are not allowed/resolved in arithmetic
generics still. This prevents e.g. `let Size = <T as Serialize>::Size +
<U as Serialize>::Size;`. This is more of an issue for when we allow
implicit associated types. For now since they must be explicit, `<T as
Serialize>::Size` must already be a named generic in scope so the
example would look more like `let Size = TSize + USize;`, see the
`serialize` test as an example.
- Additionally I had to make a couple more updates to arithmetic
generics to get the `serialize` test working:
- The existing rules were mostly good until monomorphization when we'd
have to re-do the same trait constraints but now we'd know the actual
end sizes (constants) for the constraints. This turned out to make it
more difficult as we had to solve `4 = a + b` when serializing a pair
where `a` and `b` are the serialized sizes of each item. In this case
`1` and `3` were the correct answers, and since these refer to specific
items in the pair, even an answer of `3` and `1` would be incorrect. I
ended up solving `<constant> = a + b` to `<constant> - b = a` so that
when we recur on the trait constraints for each item in the pair we'd
get `T: Serialize<Size = 4 - b>` and `U: Serialize<Size = a>`. ~~I have
a feeling the ordering could be wrong here. It type checks but requires
type annotations on the result of `serialize`.~~. We could add more
cases besides just subtraction but since this issue isn't directly
related to associated generics though I'm considering this future work.

## Documentation\*

Check one:
- [ ] No documentation needed.
- [x] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

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

---------

Co-authored-by: Ary Borenszweig <[email protected]>
Co-authored-by: Tom <[email protected]>

aztec_macros/src/transforms/note_interface.rs

@@ -88,6 +88,7 @@ pub fn generate_note_interface_impl(
         let mut note_fields = vec![];
         let note_interface_generics = trait_impl
             .trait_generics
+            .ordered_args
             .iter()
             .map(|gen| match gen.typ.clone() {
                 UnresolvedTypeData::Named(path, _, _) => Ok(path.last_name().to_string()),
@@ -120,7 +121,7 @@ pub fn generate_note_interface_impl(
                 ident("header"),
                 make_type(UnresolvedTypeData::Named(
                     chained_dep!("aztec", "note", "note_header", "NoteHeader"),
-                    vec![],
+                    Default::default(),
                     false,
                 )),
             );

aztec_macros/src/transforms/storage.rs

@@ -1,8 +1,9 @@
 use acvm::acir::AcirField;
 use noirc_errors::Span;
 use noirc_frontend::ast::{
-    BlockExpression, Expression, ExpressionKind, FunctionDefinition, Ident, Literal, NoirFunction,
-    NoirStruct, Pattern, StatementKind, TypeImpl, UnresolvedType, UnresolvedTypeData,
+    BlockExpression, Expression, ExpressionKind, FunctionDefinition, GenericTypeArgs, Ident,
+    Literal, NoirFunction, NoirStruct, Pattern, StatementKind, TypeImpl, UnresolvedType,
+    UnresolvedTypeData,
 };
 use noirc_frontend::{
     graph::CrateId,
@@ -54,13 +55,13 @@ pub fn check_for_storage_definition(
 fn inject_context_in_storage_field(field: &mut UnresolvedType) -> Result<(), AztecMacroError> {
     match &mut field.typ {
         UnresolvedTypeData::Named(path, generics, _) => {
-            generics.push(make_type(UnresolvedTypeData::Named(
+            generics.ordered_args.push(make_type(UnresolvedTypeData::Named(
                 ident_path("Context"),
-                vec![],
+                GenericTypeArgs::default(),
                 false,
             )));
             match path.last_name() {
-                "Map" => inject_context_in_storage_field(&mut generics[1]),
+                "Map" => inject_context_in_storage_field(&mut generics.ordered_args[1]),
                 _ => Ok(()),
             }
         }
@@ -144,7 +145,10 @@ pub fn generate_storage_field_constructor(
                             generate_storage_field_constructor(
                                 // Map is expected to have three generic parameters: key, value and context (i.e.
                                 // Map<K, V, Context>. Here `get(1)` fetches the value type.
-                                &(type_ident.clone(), generics.get(1).unwrap().clone()),
+                                &(
+                                    type_ident.clone(),
+                                    generics.ordered_args.get(1).unwrap().clone(),
+                                ),
                                 variable("slot"),
                             )?,
                         ),
@@ -219,8 +223,11 @@ pub fn generate_storage_implementation(
 
     // This is the type over which the impl is generic.
     let generic_context_ident = ident("Context");
-    let generic_context_type =
-        make_type(UnresolvedTypeData::Named(ident_path("Context"), vec![], true));
+    let generic_context_type = make_type(UnresolvedTypeData::Named(
+        ident_path("Context"),
+        GenericTypeArgs::default(),
+        true,
+    ));
 
     let init = NoirFunction::normal(FunctionDefinition::normal(
         &ident("init"),
@@ -231,13 +238,12 @@ pub fn generate_storage_implementation(
         &return_type(chained_path!("Self")),
     ));
 
+    let ordered_args = vec![generic_context_type.clone()];
+    let generics = GenericTypeArgs { ordered_args, named_args: Vec::new() };
+
     let storage_impl = TypeImpl {
         object_type: UnresolvedType {
-            typ: UnresolvedTypeData::Named(
-                chained_path!(storage_struct_name),
-                vec![generic_context_type.clone()],
-                true,
-            ),
+            typ: UnresolvedTypeData::Named(chained_path!(storage_struct_name), generics, true),
             span: Span::default(),
         },
         type_span: Span::default(),

aztec_macros/src/utils/parse_utils.rs

@@ -2,12 +2,13 @@ use noirc_frontend::{
     ast::{
         ArrayLiteral, AssignStatement, BlockExpression, CallExpression, CastExpression,
         ConstrainStatement, ConstructorExpression, Expression, ExpressionKind, ForLoopStatement,
-        ForRange, FunctionReturnType, Ident, IfExpression, IndexExpression, InfixExpression,
-        LValue, Lambda, LetStatement, Literal, MemberAccessExpression, MethodCallExpression,
-        ModuleDeclaration, NoirFunction, NoirStruct, NoirTrait, NoirTraitImpl, NoirTypeAlias, Path,
-        PathSegment, Pattern, PrefixExpression, Statement, StatementKind, TraitImplItem, TraitItem,
-        TypeImpl, UnresolvedGeneric, UnresolvedGenerics, UnresolvedTraitConstraint, UnresolvedType,
-        UnresolvedTypeData, UnresolvedTypeExpression, UseTree, UseTreeKind,
+        ForRange, FunctionReturnType, GenericTypeArgs, Ident, IfExpression, IndexExpression,
+        InfixExpression, LValue, Lambda, LetStatement, Literal, MemberAccessExpression,
+        MethodCallExpression, ModuleDeclaration, NoirFunction, NoirStruct, NoirTrait,
+        NoirTraitImpl, NoirTypeAlias, Path, PathSegment, Pattern, PrefixExpression, Statement,
+        StatementKind, TraitImplItem, TraitItem, TypeImpl, UnresolvedGeneric, UnresolvedGenerics,
+        UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
+        UseTree, UseTreeKind,
     },
     parser::{Item, ItemKind, ParsedSubModule, ParserError},
     ParsedModule,
@@ -297,6 +298,14 @@ fn empty_unresolved_types(unresolved_types: &mut [UnresolvedType]) {
     }
 }
 
+fn empty_type_args(generics: &mut GenericTypeArgs) {
+    empty_unresolved_types(&mut generics.ordered_args);
+    for (name, typ) in &mut generics.named_args {
+        empty_ident(name);
+        empty_unresolved_type(typ);
+    }
+}
+
 fn empty_unresolved_type(unresolved_type: &mut UnresolvedType) {
     unresolved_type.span = Default::default();
 
@@ -318,11 +327,11 @@ fn empty_unresolved_type(unresolved_type: &mut UnresolvedType) {
         }
         UnresolvedTypeData::Named(path, unresolved_types, _) => {
             empty_path(path);
-            empty_unresolved_types(unresolved_types);
+            empty_type_args(unresolved_types);
         }
         UnresolvedTypeData::TraitAsType(path, unresolved_types) => {
             empty_path(path);
-            empty_unresolved_types(unresolved_types);
+            empty_type_args(unresolved_types);
         }
         UnresolvedTypeData::MutableReference(unresolved_type) => {
             empty_unresolved_type(unresolved_type)
@@ -543,5 +552,10 @@ fn empty_unresolved_type_expression(unresolved_type_expression: &mut UnresolvedT
             empty_unresolved_type_expression(rhs);
         }
         UnresolvedTypeExpression::Constant(_, _) => (),
+        UnresolvedTypeExpression::AsTraitPath(path) => {
+            empty_unresolved_type(&mut path.typ);
+            empty_path(&mut path.trait_path);
+            empty_ident(&mut path.impl_item);
+        }
     }
 }

compiler/noirc_frontend/src/ast/mod.rs

@@ -112,10 +112,10 @@ pub enum UnresolvedTypeData {
     Parenthesized(Box<UnresolvedType>),
 
     /// A Named UnresolvedType can be a struct type or a type variable
-    Named(Path, Vec<UnresolvedType>, /*is_synthesized*/ bool),
+    Named(Path, GenericTypeArgs, /*is_synthesized*/ bool),
 
     /// A Trait as return type or parameter of function, including its generics
-    TraitAsType(Path, Vec<UnresolvedType>),
+    TraitAsType(Path, GenericTypeArgs),
 
     /// &mut T
     MutableReference(Box<UnresolvedType>),
@@ -151,6 +151,46 @@ pub struct UnresolvedType {
     pub span: Span,
 }
 
+/// An argument to a generic type or trait.
+#[derive(Debug, PartialEq, Eq, Clone, Hash)]
+pub enum GenericTypeArg {
+    /// An ordered argument, e.g. `<A, B, C>`
+    Ordered(UnresolvedType),
+
+    /// A named argument, e.g. `<A = B, C = D, E = F>`.
+    /// Used for associated types.
+    Named(Ident, UnresolvedType),
+}
+
+#[derive(Debug, Default, PartialEq, Eq, Clone, Hash)]
+pub struct GenericTypeArgs {
+    /// Each ordered argument, e.g. `<A, B, C>`
+    pub ordered_args: Vec<UnresolvedType>,
+
+    /// All named arguments, e.g. `<A = B, C = D, E = F>`.
+    /// Used for associated types.
+    pub named_args: Vec<(Ident, UnresolvedType)>,
+}
+
+impl GenericTypeArgs {
+    pub fn is_empty(&self) -> bool {
+        self.ordered_args.is_empty() && self.named_args.is_empty()
+    }
+}
+
+impl From<Vec<GenericTypeArg>> for GenericTypeArgs {
+    fn from(args: Vec<GenericTypeArg>) -> Self {
+        let mut this = GenericTypeArgs::default();
+        for arg in args {
+            match arg {
+                GenericTypeArg::Ordered(typ) => this.ordered_args.push(typ),
+                GenericTypeArg::Named(name, typ) => this.named_args.push((name, typ)),
+            }
+        }
+        this
+    }
+}
+
 /// Type wrapper for a member access
 pub struct UnaryRhsMemberAccess {
     pub method_or_field: Ident,
@@ -176,6 +216,7 @@ pub enum UnresolvedTypeExpression {
         Box<UnresolvedTypeExpression>,
         Span,
     ),
+    AsTraitPath(Box<AsTraitPath>),
 }
 
 impl Recoverable for UnresolvedType {
@@ -184,6 +225,32 @@ impl Recoverable for UnresolvedType {
     }
 }
 
+impl std::fmt::Display for GenericTypeArg {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            GenericTypeArg::Ordered(typ) => typ.fmt(f),
+            GenericTypeArg::Named(name, typ) => write!(f, "{name} = {typ}"),
+        }
+    }
+}
+
+impl std::fmt::Display for GenericTypeArgs {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.is_empty() {
+            Ok(())
+        } else {
+            let mut args = vecmap(&self.ordered_args, ToString::to_string).join(", ");
+
+            if !self.ordered_args.is_empty() && !self.named_args.is_empty() {
+                args += ", ";
+            }
+
+            args += &vecmap(&self.named_args, |(name, typ)| format!("{name} = {typ}")).join(", ");
+            write!(f, "<{args}>")
+        }
+    }
+}
+
 impl std::fmt::Display for UnresolvedTypeData {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         use UnresolvedTypeData::*;
@@ -195,22 +262,8 @@ impl std::fmt::Display for UnresolvedTypeData {
                 Signedness::Signed => write!(f, "i{num_bits}"),
                 Signedness::Unsigned => write!(f, "u{num_bits}"),
             },
-            Named(s, args, _) => {
-                let args = vecmap(args, |arg| ToString::to_string(&arg.typ));
-                if args.is_empty() {
-                    write!(f, "{s}")
-                } else {
-                    write!(f, "{}<{}>", s, args.join(", "))
-                }
-            }
-            TraitAsType(s, args) => {
-                let args = vecmap(args, |arg| ToString::to_string(&arg.typ));
-                if args.is_empty() {
-                    write!(f, "impl {s}")
-                } else {
-                    write!(f, "impl {}<{}>", s, args.join(", "))
-                }
-            }
+            Named(s, args, _) => write!(f, "{s}{args}"),
+            TraitAsType(s, args) => write!(f, "impl {s}{args}"),
             Tuple(elements) => {
                 let elements = vecmap(elements, ToString::to_string);
                 write!(f, "({})", elements.join(", "))
@@ -263,6 +316,7 @@ impl std::fmt::Display for UnresolvedTypeExpression {
             UnresolvedTypeExpression::BinaryOperation(lhs, op, rhs, _) => {
                 write!(f, "({lhs} {op} {rhs})")
             }
+            UnresolvedTypeExpression::AsTraitPath(path) => write!(f, "{path}"),
         }
     }
 }
@@ -334,6 +388,9 @@ impl UnresolvedTypeExpression {
             UnresolvedTypeExpression::Variable(path) => path.span(),
             UnresolvedTypeExpression::Constant(_, span) => *span,
             UnresolvedTypeExpression::BinaryOperation(_, _, _, span) => *span,
+            UnresolvedTypeExpression::AsTraitPath(path) => {
+                path.trait_path.span.merge(path.impl_item.span())
+            }
         }
     }
 
@@ -376,6 +433,9 @@ impl UnresolvedTypeExpression {
                 };
                 Ok(UnresolvedTypeExpression::BinaryOperation(lhs, op, rhs, expr.span))
             }
+            ExpressionKind::AsTraitPath(path) => {
+                Ok(UnresolvedTypeExpression::AsTraitPath(Box::new(path)))
+            }
             _ => Err(expr),
         }
     }

compiler/noirc_frontend/src/ast/statement.rs

@@ -7,8 +7,8 @@ use iter_extended::vecmap;
 use noirc_errors::{Span, Spanned};
 
 use super::{
-    BlockExpression, Expression, ExpressionKind, IndexExpression, MemberAccessExpression,
-    MethodCallExpression, UnresolvedType,
+    BlockExpression, Expression, ExpressionKind, GenericTypeArgs, IndexExpression,
+    MemberAccessExpression, MethodCallExpression, UnresolvedType,
 };
 use crate::elaborator::types::SELF_TYPE_NAME;
 use crate::lexer::token::SpannedToken;
@@ -371,6 +371,7 @@ impl UseTree {
 pub struct AsTraitPath {
     pub typ: UnresolvedType,
     pub trait_path: Path,
+    pub trait_generics: GenericTypeArgs,
     pub impl_item: Ident,
 }
 

compiler/noirc_frontend/src/ast/traits.rs

@@ -10,6 +10,8 @@ use crate::ast::{
 use crate::macros_api::SecondaryAttribute;
 use crate::node_interner::TraitId;
 
+use super::GenericTypeArgs;
+
 /// AST node for trait definitions:
 /// `trait name<generics> { ... items ... }`
 #[derive(Clone, Debug)]
@@ -62,7 +64,8 @@ pub struct NoirTraitImpl {
     pub impl_generics: UnresolvedGenerics,
 
     pub trait_name: Path,
-    pub trait_generics: Vec<UnresolvedType>,
+
+    pub trait_generics: GenericTypeArgs,
 
     pub object_type: UnresolvedType,
 
@@ -88,7 +91,7 @@ pub struct UnresolvedTraitConstraint {
 pub struct TraitBound {
     pub trait_path: Path,
     pub trait_id: Option<TraitId>, // initially None, gets assigned during DC
-    pub trait_generics: Vec<UnresolvedType>,
+    pub trait_generics: GenericTypeArgs,
 }
 
 #[derive(Clone, Debug)]
@@ -179,21 +182,13 @@ impl Display for UnresolvedTraitConstraint {
 
 impl Display for TraitBound {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let generics = vecmap(&self.trait_generics, |generic| generic.to_string());
-        if !generics.is_empty() {
-            write!(f, "{}<{}>", self.trait_path, generics.join(", "))
-        } else {
-            write!(f, "{}", self.trait_path)
-        }
+        write!(f, "{}{}", self.trait_path, self.trait_generics)
     }
 }
 
 impl Display for NoirTraitImpl {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let generics = vecmap(&self.trait_generics, |generic| generic.to_string());
-        let generics = generics.join(", ");
-
-        writeln!(f, "impl {}<{}> for {} {{", self.trait_name, generics, self.object_type)?;
+        writeln!(f, "impl {}{} for {} {{", self.trait_name, self.trait_generics, self.object_type)?;
 
         for item in self.items.iter() {
             let item = item.to_string();