From 422f0d172a91f9ecaa76225ffddc1e2b4464eff0 Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Mon, 9 Sep 2024 10:45:30 -0300
Subject: [PATCH 1/2] fix: when suggesting attribute with no arguments, don't
 include parentheses

---
 .../src/requests/completion/completion_items.rs  | 10 +++++++---
 tooling/lsp/src/requests/completion/tests.rs     | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tooling/lsp/src/requests/completion/completion_items.rs b/tooling/lsp/src/requests/completion/completion_items.rs
index aa959b226b0..eb8d5f659af 100644
--- a/tooling/lsp/src/requests/completion/completion_items.rs
+++ b/tooling/lsp/src/requests/completion/completion_items.rs
@@ -226,10 +226,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
index c449466fc5c..414ba16a5a9 100644
--- a/tooling/lsp/src/requests/completion/tests.rs
+++ b/tooling/lsp/src/requests/completion/tests.rs
@@ -1952,4 +1952,20 @@ 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;
+    }
 }

From 0e2aba0bc58923a949dbc25f32eb124d7a8038b5 Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Mon, 9 Sep 2024 10:54:38 -0300
Subject: [PATCH 2/2] fix: suggest trait attributes

---
 compiler/noirc_frontend/src/ast/visitor.rs    |  5 +++++
 .../lsp/src/requests/completion/builtins.rs   |  2 +-
 .../requests/completion/completion_items.rs   |  1 +
 tooling/lsp/src/requests/completion/tests.rs  | 20 +++++++++++++++++++
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/compiler/noirc_frontend/src/ast/visitor.rs b/compiler/noirc_frontend/src/ast/visitor.rs
index 6f54ad2e3b9..94c4dc591ed 100644
--- 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
index bca1061ff47..520c158d260 100644
--- 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
index eb8d5f659af..7e7511cdaa3 100644
--- 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 {
diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs
index 414ba16a5a9..97cb145eb22 100644
--- a/tooling/lsp/src/requests/completion/tests.rs
+++ b/tooling/lsp/src/requests/completion/tests.rs
@@ -1968,4 +1968,24 @@ mod completion_tests {
         )
         .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;
+    }
 }