fix: suggest trait attributes in LSP

compiler/noirc_frontend/src/ast/visitor.rs

@@ -16,7 +16,7 @@
         QuotedTypeId,
     },
     parser::{Item, ItemKind, ParsedSubModule},
     token::{CustomAtrribute, SecondaryAttribute, Tokens},
     ParsedModule, QuotedType,
 };
 
@@ -30,6 +30,7 @@
 pub enum AttributeTarget {
     Module,
     Struct,
+    Trait,
     Function,
 }
 
@@ -448,7 +449,7 @@
         true
     }
 
     fn visit_custom_attribute(&mut self, _: &CustomAtrribute, _target: AttributeTarget) {}
 }
 
 impl ParsedModule {
@@ -615,6 +616,10 @@
     }
 
     pub fn accept_children(&self, visitor: &mut impl Visitor) {
+        for attribute in &self.attributes {
+            attribute.accept(AttributeTarget::Trait, visitor);
+        }
+
         for item in &self.items {
             item.item.accept(visitor);
         }
@@ -1241,8 +1246,8 @@
             UnresolvedTypeData::Unspecified => visitor.visit_unspecified_type(self.span),
             UnresolvedTypeData::Quoted(typ) => visitor.visit_quoted_type(typ, self.span),
             UnresolvedTypeData::FieldElement => visitor.visit_field_element_type(self.span),
             UnresolvedTypeData::Integer(signdness, size) => {
                 visitor.visit_integer_type(*signdness, *size, self.span);
             }
             UnresolvedTypeData::Bool => visitor.visit_bool_type(self.span),
             UnresolvedTypeData::Unit => visitor.visit_unit_type(self.span),
@@ -1339,7 +1344,7 @@
     }
 }
 
 impl CustomAtrribute {
     pub fn accept(&self, target: AttributeTarget, visitor: &mut impl Visitor) {
         visitor.visit_custom_attribute(self, target);
     }

tooling/lsp/src/requests/completion/builtins.rs

@@ -87,7 +87,7 @@ impl<'a> NodeFinder<'a> {
 
     pub(super) fn suggest_builtin_attributes(&mut self, prefix: &str, target: AttributeTarget) {
         match target {
-            AttributeTarget::Module => (),
+            AttributeTarget::Module | AttributeTarget::Trait => (),
             AttributeTarget::Struct => {
                 self.suggest_one_argument_attributes(prefix, &["abi"]);
             }

tooling/lsp/src/requests/completion/completion_items.rs

@@ -49,6 +49,7 @@ impl<'a> NodeFinder<'a> {
                 match target {
                     AttributeTarget::Module => Some(Type::Quoted(QuotedType::Module)),
                     AttributeTarget::Struct => Some(Type::Quoted(QuotedType::StructDefinition)),
+                    AttributeTarget::Trait => Some(Type::Quoted(QuotedType::TraitDefinition)),
                     AttributeTarget::Function => Some(Type::Quoted(QuotedType::FunctionDefinition)),
                 }
             } else {
@@ -226,10 +227,14 @@ impl<'a> NodeFinder<'a> {
                     function_kind,
                     skip_first_argument,
                 );
-                let label = if insert_text.ends_with("()") {
-                    format!("{}()", name)
+                let (label, insert_text) = if insert_text.ends_with("()") {
+                    if skip_first_argument {
+                        (name.to_string(), insert_text.strip_suffix("()").unwrap().to_string())
+                    } else {
+                        (format!("{}()", name), insert_text)
+                    }
                 } else {
-                    format!("{}(…)", name)
+                    (format!("{}(…)", name), insert_text)
                 };
 
                 snippet_completion_item(label, kind, insert_text, Some(description))

tooling/lsp/src/requests/completion/tests.rs

@@ -122,14 +122,14 @@
     #[test]
     async fn test_use_first_segment() {
         let src = r#"
             mod foobaz {}
             mod foobar {}
             use foob>|<
         "#;
 
         assert_completion(
             src,
             vec![module_completion_item("foobaz"), module_completion_item("foobar")],
         )
         .await;
     }
@@ -299,7 +299,7 @@
             mod bar {
                 mod something {}
 
                 use super::foob>|<
             }
         "#;
 
@@ -1541,7 +1541,7 @@
     async fn test_auto_import_suggests_modules_too() {
         let src = r#"
             mod foo {
                 mod barbaz {
                     fn hello_world() {}
                 }
             }
@@ -1952,4 +1952,40 @@
         )
         .await;
     }
+
+    #[test]
+    async fn test_suggests_function_attribute_no_arguments() {
+        let src = r#"
+            #[some>|<]
+            fn foo() {}
+
+            fn some_attr(f: FunctionDefinition) {}
+        "#;
+
+        assert_completion_excluding_auto_import(
+            src,
+            vec![function_completion_item("some_attr", "some_attr", "fn(FunctionDefinition)")],
+        )
+        .await;
+    }
+
+    #[test]
+    async fn test_suggests_trait_attribute() {
+        let src = r#"
+            #[some>|<]
+            trait SomeTrait {}
+
+            fn some_attr(t: TraitDefinition, x: Field) {}
+        "#;
+
+        assert_completion_excluding_auto_import(
+            src,
+            vec![function_completion_item(
+                "some_attr(…)",
+                "some_attr(${1:x})",
+                "fn(TraitDefinition, Field)",
+            )],
+        )
+        .await;
+    }
 }