Merge 0e2aba0 into 45f4ae0

noir-lang · Sep 9, 2024 · 9d48b35 · 9d48b35
2 parents 45f4ae0 + 0e2aba0
commit 9d48b35
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 4 deletions.
diff --git a/compiler/noirc_frontend/src/ast/visitor.rs b/compiler/noirc_frontend/src/ast/visitor.rs
@@ -30,6 +30,7 @@ use super::{
 pub enum AttributeTarget {
     Module,
     Struct,
+    Trait,
     Function,
 }
 
@@ -615,6 +616,10 @@ impl NoirTrait {
     }
 
     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);
         }

diff --git a/tooling/lsp/src/requests/completion/builtins.rs b/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"]);
             }

diff --git a/tooling/lsp/src/requests/completion/completion_items.rs b/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))

diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs
@@ -1952,4 +1952,40 @@ mod completion_tests {
         )
         .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;
+    }
 }