Skip to content

Commit

Permalink
fix: (LSP) make goto and hover work well for attributes (#6152)
Browse files Browse the repository at this point in the history
# Description

## Problem

Hovering or using "Go to definition" on an attribute that was linked to
a function that generated code wasn't working well (it would jump to one
of the generated definitions instead of to the attribute function).

## Summary

Before (note how the docs are for the generated code, not the attr
function):


![lsp-hover-before](https://github.com/user-attachments/assets/2cb325c9-b73f-4f1c-a885-b842fd276c68)

After:


![lsp-hover-after](https://github.com/user-attachments/assets/d75a849e-8410-4e71-9d71-324c19714964)


## 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.
  • Loading branch information
asterite authored Sep 25, 2024
1 parent 66d2a07 commit c679bc6
Show file tree
Hide file tree
Showing 11 changed files with 241 additions and 31 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion compiler/noirc_frontend/src/node_interner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
Expand Down
6 changes: 1 addition & 5 deletions compiler/noirc_frontend/src/usage_tracker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,16 +28,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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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() {} }
}
1 change: 1 addition & 0 deletions tooling/lsp/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ workspace = true

[dependencies]
acvm.workspace = true
chumsky.workspace = true
codespan-lsp.workspace = true
lsp-types.workspace = true
nargo.workspace = true
Expand Down
115 changes: 115 additions & 0 deletions tooling/lsp/src/attribute_reference_finder.rs
Original file line number Diff line number Diff line change
@@ -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));
}
}
1 change: 1 addition & 0 deletions tooling/lsp/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
64 changes: 51 additions & 13 deletions tooling/lsp/src/requests/goto_definition.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand All @@ -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)
})
})
}

Expand Down Expand Up @@ -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;
}
}
61 changes: 49 additions & 12 deletions tooling/lsp/src/requests/hover.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -15,26 +15,52 @@ 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};

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)
Expand Down Expand Up @@ -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;
}
}
7 changes: 7 additions & 0 deletions tooling/lsp/test_programs/go_to_definition/src/main.nr
Original file line number Diff line number Diff line change
Expand Up @@ -28,3 +28,10 @@ trait Trait {
}

use dependency::something;

#[attr]
pub fn foo() {}

comptime fn attr(_: FunctionDefinition) -> Quoted {
quote { pub fn hello() {} }
}
7 changes: 7 additions & 0 deletions tooling/lsp/test_programs/workspace/two/src/lib.nr
Original file line number Diff line number Diff line change
Expand Up @@ -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() {} }
}

0 comments on commit c679bc6

Please sign in to comment.