From d6f60d70dc41640ad84f7a968927b20818bcaf2a Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Mon, 9 Sep 2024 17:36:34 -0300
Subject: [PATCH] fix: suggest trait attributes in LSP (#5972)

# Description

## Problem

I noticed LSP didn't suggest trait attributes... because these weren't
visited. I also noticed when an attribute related to a function was
suggested, but the function only had one argument, the parentheses were
included in the suggestion but they aren't needed.

## Summary

Fixes the above problems.


![lsp-trait-attribute](https://github.com/user-attachments/assets/16cdd8ab-1d03-40d1-a1b1-19b26d88e322)

## Additional Context



## Documentation

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

# PR Checklist

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
---
 compiler/noirc_frontend/src/ast/visitor.rs    |  5 +++
 .../lsp/src/requests/completion/builtins.rs   |  2 +-
 .../requests/completion/completion_items.rs   | 11 ++++--
 tooling/lsp/src/requests/completion/tests.rs  | 36 +++++++++++++++++++
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/compiler/noirc_frontend/src/ast/visitor.rs b/compiler/noirc_frontend/src/ast/visitor.rs
index 359e04e41a7..fcb4e4c5dd1 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,
 }
 
@@ -617,6 +618,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 aa959b226b0..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 {
@@ -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
index c449466fc5c..97cb145eb22 100644
--- 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;
+    }
 }