feat: format trait impl functions

aztec_macros/src/transforms/events.rs

@@ -1,4 +1,7 @@
-use noirc_frontend::ast::{Documented, ItemVisibility, NoirFunction, NoirTraitImpl, TraitImplItem};
+use noirc_errors::Span;
+use noirc_frontend::ast::{
+    Documented, ItemVisibility, NoirFunction, NoirTraitImpl, TraitImplItem, TraitImplItemKind,
+};
 use noirc_frontend::macros_api::{NodeInterner, StructId};
 use noirc_frontend::token::SecondaryAttribute;
 use noirc_frontend::{
@@ -67,30 +70,37 @@ pub fn generate_event_impls(
                 event_byte_len,
                 empty_spans,
             )?;
-            event_interface_trait_impl.items.push(Documented::not_documented(
-                TraitImplItem::Function(generate_fn_get_event_type_id(
+            event_interface_trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(generate_fn_get_event_type_id(
                     event_type.as_str(),
                     event_len,
                     empty_spans,
                 )?),
-            ));
-            event_interface_trait_impl.items.push(Documented::not_documented(
-                TraitImplItem::Function(generate_fn_private_to_be_bytes(
+                span: Span::default(),
+            }));
+            event_interface_trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(generate_fn_private_to_be_bytes(
                     event_type.as_str(),
                     event_byte_len,
                     empty_spans,
                 )?),
-            ));
-            event_interface_trait_impl.items.push(Documented::not_documented(
-                TraitImplItem::Function(generate_fn_to_be_bytes(
+                span: Span::default(),
+            }));
+            event_interface_trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(generate_fn_to_be_bytes(
                     event_type.as_str(),
                     event_byte_len,
                     empty_spans,
                 )?),
-            ));
-            event_interface_trait_impl.items.push(Documented::not_documented(
-                TraitImplItem::Function(generate_fn_emit(event_type.as_str(), empty_spans)?),
-            ));
+                span: Span::default(),
+            }));
+            event_interface_trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(generate_fn_emit(
+                    event_type.as_str(),
+                    empty_spans,
+                )?),
+                span: Span::default(),
+            }));
             submodule.contents.trait_impls.push(event_interface_trait_impl);
 
             let serialize_trait_impl = generate_trait_impl_serialize(

aztec_macros/src/transforms/note_interface.rs

@@ -1,7 +1,7 @@
 use noirc_errors::Span;
 use noirc_frontend::ast::{
     Documented, ItemVisibility, LetStatement, NoirFunction, NoirStruct, PathKind, StructField,
-    TraitImplItem, TypeImpl, UnresolvedTypeData, UnresolvedTypeExpression,
+    TraitImplItem, TraitImplItemKind, TypeImpl, UnresolvedTypeData, UnresolvedTypeExpression,
 };
 use noirc_frontend::{
     graph::CrateId,
@@ -153,9 +153,10 @@ pub fn generate_note_interface_impl(
                 note_interface_impl_span,
                 empty_spans,
             )?;
-            trait_impl.items.push(Documented::not_documented(TraitImplItem::Function(
-                note_serialize_content_fn,
-            )));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(note_serialize_content_fn),
+                span: note_interface_impl_span,
+            }));
 
             let note_deserialize_content_fn = generate_note_deserialize_content(
                 &note_type,
@@ -165,9 +166,10 @@ pub fn generate_note_interface_impl(
                 note_interface_impl_span,
                 empty_spans,
             )?;
-            trait_impl.items.push(Documented::not_documented(TraitImplItem::Function(
-                note_deserialize_content_fn,
-            )));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(note_deserialize_content_fn),
+                span: note_interface_impl_span,
+            }));
 
             let note_properties_struct = generate_note_properties_struct(
                 &note_type,
@@ -196,9 +198,10 @@ pub fn generate_note_interface_impl(
                 note_interface_impl_span,
                 empty_spans,
             )?;
-            trait_impl
-                .items
-                .push(Documented::not_documented(TraitImplItem::Function(get_header_fn)));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(get_header_fn),
+                span: note_interface_impl_span,
+            }));
         }
         if !check_trait_method_implemented(trait_impl, "set_header") {
             let set_header_fn = generate_note_set_header(
@@ -207,18 +210,20 @@ pub fn generate_note_interface_impl(
                 note_interface_impl_span,
                 empty_spans,
             )?;
-            trait_impl
-                .items
-                .push(Documented::not_documented(TraitImplItem::Function(set_header_fn)));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(set_header_fn),
+                span: note_interface_impl_span,
+            }));
         }
 
         if !check_trait_method_implemented(trait_impl, "get_note_type_id") {
             let note_type_id = compute_note_type_id(&note_type);
             let get_note_type_id_fn =
                 generate_get_note_type_id(note_type_id, note_interface_impl_span, empty_spans)?;
-            trait_impl
-                .items
-                .push(Documented::not_documented(TraitImplItem::Function(get_note_type_id_fn)));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(get_note_type_id_fn),
+                span: note_interface_impl_span,
+            }));
         }
 
         if !check_trait_method_implemented(trait_impl, "compute_note_hiding_point") {
@@ -227,9 +232,10 @@ pub fn generate_note_interface_impl(
                 note_interface_impl_span,
                 empty_spans,
             )?;
-            trait_impl.items.push(Documented::not_documented(TraitImplItem::Function(
-                compute_note_hiding_point_fn,
-            )));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(compute_note_hiding_point_fn),
+                span: note_interface_impl_span,
+            }));
         }
 
         if !check_trait_method_implemented(trait_impl, "to_be_bytes") {
@@ -240,9 +246,10 @@ pub fn generate_note_interface_impl(
                 note_interface_impl_span,
                 empty_spans,
             )?;
-            trait_impl
-                .items
-                .push(Documented::not_documented(TraitImplItem::Function(to_be_bytes_fn)));
+            trait_impl.items.push(Documented::not_documented(TraitImplItem {
+                kind: TraitImplItemKind::Function(to_be_bytes_fn),
+                span: note_interface_impl_span,
+            }));
         }
     }
 

aztec_macros/src/utils/ast_utils.rs

@@ -3,7 +3,7 @@ use noirc_frontend::ast::{
     BinaryOpKind, CallExpression, CastExpression, Expression, ExpressionKind, FunctionReturnType,
     Ident, IndexExpression, InfixExpression, Lambda, MemberAccessExpression, MethodCallExpression,
     NoirTraitImpl, Path, PathSegment, Pattern, PrefixExpression, Statement, StatementKind,
-    TraitImplItem, UnaryOp, UnresolvedType, UnresolvedTypeData,
+    TraitImplItemKind, UnaryOp, UnresolvedType, UnresolvedTypeData,
 };
 use noirc_frontend::token::SecondaryAttribute;
 
@@ -179,8 +179,8 @@ pub fn index_array(array: Ident, index: &str) -> Expression {
 }
 
 pub fn check_trait_method_implemented(trait_impl: &NoirTraitImpl, method_name: &str) -> bool {
-    trait_impl.items.iter().any(|item| match &item.item {
-        TraitImplItem::Function(func) => func.def.name.0.contents == method_name,
+    trait_impl.items.iter().any(|item| match &item.item.kind {
+        TraitImplItemKind::Function(func) => func.def.name.0.contents == method_name,
         _ => false,
     })
 }

aztec_macros/src/utils/parse_utils.rs

@@ -6,9 +6,9 @@ use noirc_frontend::{
         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,
+        StatementKind, TraitImplItem, TraitImplItemKind, TraitItem, TypeImpl, UnresolvedGeneric,
+        UnresolvedGenerics, UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData,
+        UnresolvedTypeExpression, UseTree, UseTreeKind,
     },
     parser::{Item, ItemKind, ParsedSubModule, ParserError},
     ParsedModule,
@@ -137,14 +137,20 @@ fn empty_trait_item(trait_item: &mut TraitItem) {
 }
 
 fn empty_trait_impl_item(trait_impl_item: &mut TraitImplItem) {
+    trait_impl_item.span = Default::default();
+
+    empty_trait_impl_item_kind(&mut trait_impl_item.kind);
+}
+
+fn empty_trait_impl_item_kind(trait_impl_item: &mut TraitImplItemKind) {
     match trait_impl_item {
-        TraitImplItem::Function(noir_function) => empty_noir_function(noir_function),
-        TraitImplItem::Constant(name, typ, default_value) => {
+        TraitImplItemKind::Function(noir_function) => empty_noir_function(noir_function),
+        TraitImplItemKind::Constant(name, typ, default_value) => {
             empty_ident(name);
             empty_unresolved_type(typ);
             empty_expression(default_value);
         }
-        TraitImplItem::Type { name, alias } => {
+        TraitImplItemKind::Type { name, alias } => {
             empty_ident(name);
             empty_unresolved_type(alias);
         }

compiler/noirc_frontend/src/ast/traits.rs

@@ -95,7 +95,13 @@ pub struct TraitBound {
 }
 
 #[derive(Clone, Debug)]
-pub enum TraitImplItem {
+pub struct TraitImplItem {
+    pub kind: TraitImplItemKind,
+    pub span: Span,
+}
+
+#[derive(Clone, Debug)]
+pub enum TraitImplItemKind {
     Function(NoirFunction),
     Constant(Ident, UnresolvedType, Expression),
     Type { name: Ident, alias: UnresolvedType },
@@ -202,11 +208,17 @@ impl Display for NoirTraitImpl {
 }
 
 impl Display for TraitImplItem {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.kind.fmt(f)
+    }
+}
+
+impl Display for TraitImplItemKind {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            TraitImplItem::Function(function) => function.fmt(f),
-            TraitImplItem::Type { name, alias } => write!(f, "type {name} = {alias};"),
-            TraitImplItem::Constant(name, typ, value) => {
+            TraitImplItemKind::Function(function) => function.fmt(f),
+            TraitImplItemKind::Type { name, alias } => write!(f, "type {name} = {alias};"),
+            TraitImplItemKind::Constant(name, typ, value) => {
                 write!(f, "let {name}: {typ} = {value};")
             }
         }

compiler/noirc_frontend/src/ast/visitor.rs

@@ -22,8 +22,8 @@ use crate::{
 
 use super::{
     FunctionReturnType, GenericTypeArgs, IntegerBitSize, ItemVisibility, Pattern, Signedness,
-    UnresolvedGenerics, UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData,
-    UnresolvedTypeExpression,
+    TraitImplItemKind, UnresolvedGenerics, UnresolvedTraitConstraint, UnresolvedType,
+    UnresolvedTypeData, UnresolvedTypeExpression,
 };
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -69,6 +69,10 @@ pub trait Visitor {
         true
     }
 
+    fn visit_trait_impl_item_kind(&mut self, _: &TraitImplItemKind, _span: Span) -> bool {
+        true
+    }
+
     fn visit_trait_impl_item_function(&mut self, _: &NoirFunction, _span: Span) -> bool {
         true
     }
@@ -78,11 +82,17 @@ pub trait Visitor {
         _name: &Ident,
         _typ: &UnresolvedType,
         _expression: &Expression,
+        _span: Span,
     ) -> bool {
         true
     }
 
-    fn visit_trait_impl_item_type(&mut self, _name: &Ident, _alias: &UnresolvedType) -> bool {
+    fn visit_trait_impl_item_type(
+        &mut self,
+        _name: &Ident,
+        _alias: &UnresolvedType,
+        _span: Span,
+    ) -> bool {
         true
     }
 
@@ -569,24 +579,32 @@ impl TraitImplItem {
     }
 
     pub fn accept_children(&self, visitor: &mut impl Visitor) {
-        match self {
-            TraitImplItem::Function(noir_function) => {
-                let span = Span::from(
-                    noir_function.name_ident().span().start()..noir_function.span().end(),
-                );
+        self.kind.accept(self.span, visitor);
+    }
+}
+
+impl TraitImplItemKind {
+    pub fn accept(&self, span: Span, visitor: &mut impl Visitor) {
+        if visitor.visit_trait_impl_item_kind(self, span) {
+            self.accept_children(span, visitor);
+        }
+    }
 
+    pub fn accept_children(&self, span: Span, visitor: &mut impl Visitor) {
+        match self {
+            TraitImplItemKind::Function(noir_function) => {
                 if visitor.visit_trait_impl_item_function(noir_function, span) {
                     noir_function.accept(span, visitor);
                 }
             }
-            TraitImplItem::Constant(name, unresolved_type, expression) => {
-                if visitor.visit_trait_impl_item_constant(name, unresolved_type, expression) {
+            TraitImplItemKind::Constant(name, unresolved_type, expression) => {
+                if visitor.visit_trait_impl_item_constant(name, unresolved_type, expression, span) {
                     unresolved_type.accept(visitor);
                     expression.accept(visitor);
                 }
             }
-            TraitImplItem::Type { name, alias } => {
-                if visitor.visit_trait_impl_item_type(name, alias) {
+            TraitImplItemKind::Type { name, alias } => {
+                if visitor.visit_trait_impl_item_type(name, alias, span) {
                     alias.accept(visitor);
                 }
             }

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

@@ -12,7 +12,7 @@ use rustc_hash::FxHashMap as HashMap;
 
 use crate::ast::{
     Documented, FunctionDefinition, Ident, ItemVisibility, LetStatement, ModuleDeclaration,
-    NoirFunction, NoirStruct, NoirTrait, NoirTraitImpl, NoirTypeAlias, Pattern, TraitImplItem,
+    NoirFunction, NoirStruct, NoirTrait, NoirTraitImpl, NoirTypeAlias, Pattern, TraitImplItemKind,
     TraitItem, TypeImpl,
 };
 use crate::hir::resolution::errors::ResolverError;
@@ -1124,18 +1124,18 @@ pub(crate) fn collect_trait_impl_items(
     let module = ModuleId { krate, local_id };
 
     for item in std::mem::take(&mut trait_impl.items) {
-        match item.item {
-            TraitImplItem::Function(impl_method) => {
+        match item.item.kind {
+            TraitImplItemKind::Function(impl_method) => {
                 let func_id = interner.push_empty_fn();
                 let location = Location::new(impl_method.span(), file_id);
                 interner.push_function(func_id, &impl_method.def, module, location);
                 interner.set_doc_comments(ReferenceId::Function(func_id), item.doc_comments);
                 unresolved_functions.push_fn(local_id, func_id, impl_method);
             }
-            TraitImplItem::Constant(name, typ, expr) => {
+            TraitImplItemKind::Constant(name, typ, expr) => {
                 associated_constants.push((name, typ, expr));
             }
-            TraitImplItem::Type { name, alias } => {
+            TraitImplItemKind::Type { name, alias } => {
                 associated_types.push((name, alias));
             }
         }