fix: (LSP) make goto and hover work well for attributes

compiler/noirc_frontend/src/node_interner.rs

@@ -222,12 +222,12 @@
     interned_statement_kinds: noirc_arena::Arena<StatementKind>,
 
     // Interned `UnresolvedTypeData`s during comptime code.
     interned_unresolved_type_datas: noirc_arena::Arena<UnresolvedTypeData>,
 
     // Interned `Pattern`s during comptime code.
     interned_patterns: noirc_arena::Arena<Pattern>,
 
     /// Determins whether to run in LSP mode. In LSP mode references are tracked.
     pub(crate) lsp_mode: bool,
 
     /// Store the location of the references in the graph.
@@ -666,7 +666,7 @@
             quoted_types: Default::default(),
             interned_expression_kinds: Default::default(),
             interned_statement_kinds: Default::default(),
             interned_unresolved_type_datas: Default::default(),
             interned_patterns: Default::default(),
             lsp_mode: false,
             location_indices: LocationIndices::default(),
@@ -676,7 +676,7 @@
             auto_import_names: HashMap::default(),
             comptime_scopes: vec![HashMap::default()],
             trait_impl_associated_types: HashMap::default(),
-            usage_tracker: UsageTracker::new(),
+            usage_tracker: UsageTracker::default(),
             doc_comments: HashMap::default(),
         }
     }
@@ -2145,11 +2145,11 @@
         &mut self,
         typ: UnresolvedTypeData,
     ) -> InternedUnresolvedTypeData {
         InternedUnresolvedTypeData(self.interned_unresolved_type_datas.insert(typ))
     }
 
     pub fn get_unresolved_type_data(&self, id: InternedUnresolvedTypeData) -> &UnresolvedTypeData {
         &self.interned_unresolved_type_datas[id.0]
     }
 
     /// Returns the type of an operator (which is always a function), along with its return type.

compiler/noirc_frontend/src/usage_tracker.rs

@@ -26,16 +26,12 @@ impl UnusedItem {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct UsageTracker {
     unused_items: HashMap<ModuleId, HashMap<Ident, UnusedItem>>,
 }
 
 impl UsageTracker {
-    pub(crate) fn new() -> Self {
-        Self { unused_items: HashMap::new() }
-    }
-
     pub(crate) fn add_unused_item(
         &mut self,
         module_id: ModuleId,

test_programs/compile_success_empty/field_or_integer_static_trait_method/src/main.nr

@@ -23,3 +23,10 @@ fn main() {
     let value: Field = Field::read(data);
     assert_eq(value, 10);
 }
+
+#[attr]
+pub fn foo() {}
+
+comptime fn attr(_: FunctionDefinition) -> Quoted {
+    quote { pub fn hello() {} }
+}

tooling/lsp/Cargo.toml

@@ -14,6 +14,7 @@ workspace = true
 
 [dependencies]
 acvm.workspace = true
+chumsky.workspace = true
 codespan-lsp.workspace = true
 lsp-types.workspace = true
 nargo.workspace = true

tooling/lsp/src/attribute_reference_finder.rs

@@ -0,0 +1,115 @@
+/// If the cursor is on an custom attribute, this struct will try to resolve its
+/// underlying function and return a ReferenceId to it.
+/// This is needed in hover and go-to-definition because when an annotation generates
+/// code, that code ends up residing in the attribute definition (it ends up having the
+/// attribute's span) so using the usual graph to locate what points to that location
+/// will give not only the attribute function but also any type generated by it.
+use std::collections::BTreeMap;
+
+use chumsky::Parser;
+use fm::FileId;
+use noirc_errors::Span;
+use noirc_frontend::{
+    ast::{AttributeTarget, Visitor},
+    graph::CrateId,
+    hir::{
+        def_map::{CrateDefMap, LocalModuleId, ModuleId},
+        resolution::path_resolver::{PathResolver, StandardPathResolver},
+    },
+    lexer::Lexer,
+    node_interner::ReferenceId,
+    parser::{path_no_turbofish, ParsedSubModule},
+    token::CustomAttribute,
+    usage_tracker::UsageTracker,
+    ParsedModule,
+};
+
+use crate::modules::module_def_id_to_reference_id;
+
+pub(crate) struct AttributeReferenceFinder<'a> {
+    byte_index: usize,
+    /// The module ID in scope. This might change as we traverse the AST
+    /// if we are analyzing something inside an inline module declaration.
+    module_id: ModuleId,
+    def_maps: &'a BTreeMap<CrateId, CrateDefMap>,
+    reference_id: Option<ReferenceId>,
+}
+
+impl<'a> AttributeReferenceFinder<'a> {
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn new(
+        file: FileId,
+        byte_index: usize,
+        krate: CrateId,
+        def_maps: &'a BTreeMap<CrateId, CrateDefMap>,
+    ) -> Self {
+        // Find the module the current file belongs to
+        let def_map = &def_maps[&krate];
+        let local_id = if let Some((module_index, _)) =
+            def_map.modules().iter().find(|(_, module_data)| module_data.location.file == file)
+        {
+            LocalModuleId(module_index)
+        } else {
+            def_map.root()
+        };
+        let module_id = ModuleId { krate, local_id };
+        Self { byte_index, module_id, def_maps, reference_id: None }
+    }
+
+    pub(crate) fn find(&mut self, parsed_module: &ParsedModule) -> Option<ReferenceId> {
+        parsed_module.accept(self);
+
+        self.reference_id
+    }
+
+    fn includes_span(&self, span: Span) -> bool {
+        span.start() as usize <= self.byte_index && self.byte_index <= span.end() as usize
+    }
+}
+
+impl<'a> Visitor for AttributeReferenceFinder<'a> {
+    fn visit_parsed_submodule(&mut self, parsed_sub_module: &ParsedSubModule, _span: Span) -> bool {
+        // Switch `self.module_id` to the submodule
+        let previous_module_id = self.module_id;
+
+        let def_map = &self.def_maps[&self.module_id.krate];
+        if let Some(module_data) = def_map.modules().get(self.module_id.local_id.0) {
+            if let Some(child_module) = module_data.children.get(&parsed_sub_module.name) {
+                self.module_id = ModuleId { krate: self.module_id.krate, local_id: *child_module };
+            }
+        }
+
+        parsed_sub_module.accept_children(self);
+
+        // Restore the old module before continuing
+        self.module_id = previous_module_id;
+
+        false
+    }
+
+    fn visit_custom_attribute(&mut self, attribute: &CustomAttribute, _target: AttributeTarget) {
+        if !self.includes_span(attribute.contents_span) {
+            return;
+        }
+
+        let name = match attribute.contents.split_once('(') {
+            Some((left, _right)) => left.to_string(),
+            None => attribute.contents.to_string(),
+        };
+        let (tokens, _) = Lexer::lex(&name);
+
+        let parser = path_no_turbofish();
+        let Ok(path) = parser.parse(tokens) else {
+            return;
+        };
+
+        let resolver = StandardPathResolver::new(self.module_id);
+        let mut usage_tracker = UsageTracker::default();
+        let Ok(result) = resolver.resolve(self.def_maps, path, &mut usage_tracker, &mut None)
+        else {
+            return;
+        };
+
+        self.reference_id = Some(module_def_id_to_reference_id(result.module_def_id));
+    }
+}

tooling/lsp/src/lib.rs

@@ -62,6 +62,7 @@ use serde_json::Value as JsonValue;
 use thiserror::Error;
 use tower::Service;
 
+mod attribute_reference_finder;
 mod modules;
 mod notifications;
 mod requests;

tooling/lsp/src/requests/goto_definition.rs

@@ -1,8 +1,11 @@
 use std::future::{self, Future};
 
+use crate::attribute_reference_finder::AttributeReferenceFinder;
+use crate::utils;
 use crate::{types::GotoDefinitionResult, LspState};
 use async_lsp::ResponseError;
 
+use fm::PathString;
 use lsp_types::request::GotoTypeDefinitionParams;
 use lsp_types::{GotoDefinitionParams, GotoDefinitionResponse};
 
@@ -29,21 +32,42 @@ fn on_goto_definition_inner(
     params: GotoDefinitionParams,
     return_type_location_instead: bool,
 ) -> Result<GotoDefinitionResult, ResponseError> {
+    let uri = params.text_document_position_params.text_document.uri.clone();
+    let position = params.text_document_position_params.position;
     process_request(state, params.text_document_position_params, |args| {
-        args.interner
-            .get_definition_location_from(args.location, return_type_location_instead)
-            .or_else(|| {
-                args.interner
-                    .reference_at_location(args.location)
-                    .map(|reference| args.interner.reference_location(reference))
-            })
-            .and_then(|found_location| {
-                let file_id = found_location.file;
-                let definition_position =
-                    to_lsp_location(args.files, file_id, found_location.span)?;
-                let response = GotoDefinitionResponse::from(definition_position).to_owned();
-                Some(response)
+        let path = PathString::from_path(uri.to_file_path().unwrap());
+        let reference_id = args.files.get_file_id(&path).and_then(|file_id| {
+            utils::position_to_byte_index(args.files, file_id, &position).and_then(|byte_index| {
+                let file = args.files.get_file(file_id).unwrap();
+                let source = file.source();
+                let (parsed_module, _errors) = noirc_frontend::parse_program(source);
+
+                let mut finder = AttributeReferenceFinder::new(
+                    file_id,
+                    byte_index,
+                    args.crate_id,
+                    args.def_maps,
+                );
+                finder.find(&parsed_module)
             })
+        });
+        let location = if let Some(reference_id) = reference_id {
+            Some(args.interner.reference_location(reference_id))
+        } else {
+            args.interner
+                .get_definition_location_from(args.location, return_type_location_instead)
+                .or_else(|| {
+                    args.interner
+                        .reference_at_location(args.location)
+                        .map(|reference| args.interner.reference_location(reference))
+                })
+        };
+        location.and_then(|found_location| {
+            let file_id = found_location.file;
+            let definition_position = to_lsp_location(args.files, file_id, found_location.span)?;
+            let response = GotoDefinitionResponse::from(definition_position).to_owned();
+            Some(response)
+        })
     })
 }
 
@@ -249,4 +273,18 @@ mod goto_definition_tests {
         )
         .await;
     }
+
+    #[test]
+    async fn goto_attribute_function() {
+        expect_goto(
+            "go_to_definition",
+            Position { line: 31, character: 3 }, // "attr"
+            "src/main.nr",
+            Range {
+                start: Position { line: 34, character: 12 },
+                end: Position { line: 34, character: 16 },
+            },
+        )
+        .await;
+    }
 }

tooling/lsp/src/requests/hover.rs

@@ -1,7 +1,7 @@
 use std::future::{self, Future};
 
 use async_lsp::ResponseError;
-use fm::FileMap;
+use fm::{FileMap, PathString};
 use lsp_types::{Hover, HoverContents, HoverParams, MarkupContent, MarkupKind};
 use noirc_frontend::{
     ast::Visibility,
@@ -15,26 +15,52 @@
     Generics, Shared, StructType, Type, TypeAlias, TypeBinding, TypeVariable,
 };
 
-use crate::{modules::module_full_path, LspState};
+use crate::{
+    attribute_reference_finder::AttributeReferenceFinder, modules::module_full_path, utils,
+    LspState,
+};
 
 use super::{process_request, to_lsp_location, ProcessRequestCallbackArgs};
 
 pub(crate) fn on_hover_request(
     state: &mut LspState,
     params: HoverParams,
 ) -> impl Future<Output = Result<Option<Hover>, ResponseError>> {
+    let uri = params.text_document_position_params.text_document.uri.clone();
+    let position = params.text_document_position_params.position;
     let result = process_request(state, params.text_document_position_params, |args| {
-        args.interner.reference_at_location(args.location).and_then(|reference| {
-            let location = args.interner.reference_location(reference);
-            let lsp_location = to_lsp_location(args.files, location.file, location.span);
-            format_reference(reference, &args).map(|formatted| Hover {
-                range: lsp_location.map(|location| location.range),
-                contents: HoverContents::Markup(MarkupContent {
-                    kind: MarkupKind::Markdown,
-                    value: formatted,
-                }),
+        let path = PathString::from_path(uri.to_file_path().unwrap());
+        args.files
+            .get_file_id(&path)
+            .and_then(|file_id| {
+                utils::position_to_byte_index(args.files, file_id, &position).and_then(
+                    |byte_index| {
+                        let file = args.files.get_file(file_id).unwrap();
+                        let source = file.source();
+                        let (parsed_module, _errors) = noirc_frontend::parse_program(source);
+
+                        let mut finder = AttributeReferenceFinder::new(
+                            file_id,
+                            byte_index,
+                            args.crate_id,
+                            args.def_maps,
+                        );
+                        finder.find(&parsed_module)
+                    },
+                )
+            })
+            .or_else(|| args.interner.reference_at_location(args.location))
+            .and_then(|reference| {
+                let location = args.interner.reference_location(reference);
+                let lsp_location = to_lsp_location(args.files, location.file, location.span);
+                format_reference(reference, &args).map(|formatted| Hover {
+                    range: lsp_location.map(|location| location.range),
+                    contents: HoverContents::Markup(MarkupContent {
+                        kind: MarkupKind::Markdown,
+                        value: formatted,
+                    }),
+                })
             })
-        })
     });
 
     future::ready(result)
@@ -653,7 +679,7 @@
             "two/src/lib.nr",
             Position { line: 6, character: 9 },
             r#"    one
     mod subone"#,
         )
         .await;
     }
@@ -664,7 +690,7 @@
             "workspace",
             "two/src/lib.nr",
             Position { line: 9, character: 20 },
             r#"    one::subone
     struct SubOneStruct {
         some_field: i32,
         some_other_field: Field,
@@ -679,7 +705,7 @@
             "workspace",
             "two/src/lib.nr",
             Position { line: 46, character: 17 },
             r#"    one::subone
     struct GenericStruct<A, B> {
     }"#,
         )
@@ -692,7 +718,7 @@
             "workspace",
             "two/src/lib.nr",
             Position { line: 9, character: 35 },
             r#"    one::subone::SubOneStruct
     some_field: i32"#,
         )
         .await;
@@ -704,7 +730,7 @@
             "workspace",
             "two/src/lib.nr",
             Position { line: 12, character: 17 },
             r#"    one::subone
     trait SomeTrait"#,
         )
         .await;
@@ -918,4 +944,15 @@
         )
         .await;
     }
+
+    #[test]
+    async fn hover_on_attribute_function() {
+        assert_hover(
+            "workspace",
+            "two/src/lib.nr",
+            Position { line: 54, character: 2 },
+            "    two\n    fn attr(_: FunctionDefinition) -> Quoted",
+        )
+        .await;
+    }
 }

tooling/lsp/test_programs/go_to_definition/src/main.nr

@@ -28,3 +28,10 @@ trait Trait {
 }
 
 use dependency::something;
+
+#[attr]
+pub fn foo() {}
+
+comptime fn attr(_: FunctionDefinition) -> Quoted {
+    quote { pub fn hello() {} }
+}

tooling/lsp/test_programs/workspace/two/src/lib.nr

@@ -51,3 +51,10 @@ use std::collections::bounded_vec::BoundedVec;
 fn instantiate_generic() {
     let x: BoundedVec<subone::SubOneStruct, 3> = BoundedVec::new();
 }
+
+#[attr]
+pub fn foo() {}
+
+comptime fn attr(_: FunctionDefinition) -> Quoted {
+    quote { pub fn hello() {} }
+}