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; + } }