diff --git a/Cargo.lock b/Cargo.lock index 6a469bd67f4..cdc1d75a421 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2753,6 +2753,7 @@ version = "0.34.0" dependencies = [ "acvm", "async-lsp", + "chumsky", "codespan-lsp", "convert_case 0.6.0", "fm", diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs index 75178df319d..23ca3387f09 100644 --- a/compiler/noirc_frontend/src/node_interner.rs +++ b/compiler/noirc_frontend/src/node_interner.rs @@ -676,7 +676,7 @@ impl Default for NodeInterner { 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(), } } diff --git a/compiler/noirc_frontend/src/usage_tracker.rs b/compiler/noirc_frontend/src/usage_tracker.rs index 23c97370b46..4224ba62df9 100644 --- a/compiler/noirc_frontend/src/usage_tracker.rs +++ b/compiler/noirc_frontend/src/usage_tracker.rs @@ -28,16 +28,12 @@ impl UnusedItem { } } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct UsageTracker { unused_items: HashMap>, } impl UsageTracker { - pub(crate) fn new() -> Self { - Self { unused_items: HashMap::new() } - } - pub(crate) fn add_unused_item( &mut self, module_id: ModuleId, diff --git a/test_programs/compile_success_empty/field_or_integer_static_trait_method/src/main.nr b/test_programs/compile_success_empty/field_or_integer_static_trait_method/src/main.nr index f5a54c82ce5..64eb564cda1 100644 --- a/test_programs/compile_success_empty/field_or_integer_static_trait_method/src/main.nr +++ b/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() {} } +} diff --git a/tooling/lsp/Cargo.toml b/tooling/lsp/Cargo.toml index c15895d801f..209f2afe4a4 100644 --- a/tooling/lsp/Cargo.toml +++ b/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 diff --git a/tooling/lsp/src/attribute_reference_finder.rs b/tooling/lsp/src/attribute_reference_finder.rs new file mode 100644 index 00000000000..f08c8073a79 --- /dev/null +++ b/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, + reference_id: Option, +} + +impl<'a> AttributeReferenceFinder<'a> { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + file: FileId, + byte_index: usize, + krate: CrateId, + def_maps: &'a BTreeMap, + ) -> 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 { + 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)); + } +} diff --git a/tooling/lsp/src/lib.rs b/tooling/lsp/src/lib.rs index 39d4c3faa61..771f67e1fa2 100644 --- a/tooling/lsp/src/lib.rs +++ b/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; diff --git a/tooling/lsp/src/requests/goto_definition.rs b/tooling/lsp/src/requests/goto_definition.rs index 6538e64dc90..a2443ea165d 100644 --- a/tooling/lsp/src/requests/goto_definition.rs +++ b/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 { + 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; + } } diff --git a/tooling/lsp/src/requests/hover.rs b/tooling/lsp/src/requests/hover.rs index 46d2a5cfc8f..2628c9b2ab6 100644 --- a/tooling/lsp/src/requests/hover.rs +++ b/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,7 +15,10 @@ use noirc_frontend::{ 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}; @@ -23,18 +26,41 @@ pub(crate) fn on_hover_request( state: &mut LspState, params: HoverParams, ) -> impl Future, 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) @@ -918,4 +944,15 @@ mod hover_tests { ) .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; + } } diff --git a/tooling/lsp/test_programs/go_to_definition/src/main.nr b/tooling/lsp/test_programs/go_to_definition/src/main.nr index 4550324c182..cb04fe04eaa 100644 --- a/tooling/lsp/test_programs/go_to_definition/src/main.nr +++ b/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() {} } +} \ No newline at end of file diff --git a/tooling/lsp/test_programs/workspace/two/src/lib.nr b/tooling/lsp/test_programs/workspace/two/src/lib.nr index 3f0f0f117b7..c6b516d88ab 100644 --- a/tooling/lsp/test_programs/workspace/two/src/lib.nr +++ b/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 = BoundedVec::new(); } + +#[attr] +pub fn foo() {} + +comptime fn attr(_: FunctionDefinition) -> Quoted { + quote { pub fn hello() {} } +} \ No newline at end of file