fix: let LSP work well in Noir workspaces

compiler/noirc_frontend/src/locations.rs

@@ -178,14 +178,8 @@ impl NodeInterner {
         include_referenced: bool,
         include_self_type_name: bool,
     ) -> Option<Vec<Location>> {
-        let node_index = self.location_indices.get_node_from_location(location)?;
-
-        let reference_node = self.reference_graph[node_index];
-        let referenced_node_index = if let ReferenceId::Reference(_, _) = reference_node {
-            self.referenced_index(node_index)?
-        } else {
-            node_index
-        };
+        let referenced_node = self.find_referenced(location)?;
+        let referenced_node_index = self.reference_graph_indices[&referenced_node];
 
         let found_locations = self.find_all_references_for_index(
             referenced_node_index,
@@ -196,6 +190,19 @@ impl NodeInterner {
         Some(found_locations)
     }
 
+    // Returns the `ReferenceId` that is referenced by the given location, if any.
+    pub fn find_referenced(&self, location: Location) -> Option<ReferenceId> {
+        let node_index = self.location_indices.get_node_from_location(location)?;
+
+        let reference_node = self.reference_graph[node_index];
+        if let ReferenceId::Reference(_, _) = reference_node {
+            let node_index = self.referenced_index(node_index)?;
+            Some(self.reference_graph[node_index])
+        } else {
+            Some(reference_node)
+        }
+    }
+
     // Given a referenced node index, find all references to it and return their locations, optionally together
     // with the reference node's location if `include_referenced` is true.
     // If `include_self_type_name` is true, references where "Self" is written are returned,

tooling/lsp/src/lib.rs

@@ -229,43 +229,74 @@ fn byte_span_to_range<'a, F: files::Files<'a> + ?Sized>(
     }
 }
 
-pub(crate) fn resolve_workspace_for_source_path(file_path: &Path) -> Result<Workspace, LspError> {
-    if let Some(toml_path) = find_file_manifest(file_path) {
-        resolve_workspace_from_toml(
-            &toml_path,
-            PackageSelection::All,
-            Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
-        )
-        .map_err(|err| LspError::WorkspaceResolutionError(err.to_string()))
-    } else {
-        let Some(parent_folder) = file_path
-            .parent()
-            .and_then(|f| f.file_name())
-            .and_then(|file_name_os_str| file_name_os_str.to_str())
-        else {
-            return Err(LspError::WorkspaceResolutionError(format!(
-                "Could not resolve parent folder for file: {:?}",
-                file_path
-            )));
-        };
-        let assumed_package = Package {
-            version: None,
-            compiler_required_version: Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
-            root_dir: PathBuf::from(parent_folder),
-            package_type: PackageType::Binary,
-            entry_path: PathBuf::from(file_path),
-            name: CrateName::from_str(parent_folder)
-                .map_err(|err| LspError::WorkspaceResolutionError(err.to_string()))?,
-            dependencies: BTreeMap::new(),
-        };
-        let workspace = Workspace {
-            root_dir: PathBuf::from(parent_folder),
-            members: vec![assumed_package],
-            selected_package_index: Some(0),
-            is_assumed: true,
-        };
-        Ok(workspace)
+pub(crate) fn resolve_workspace_for_source_path(
+    file_path: &Path,
+    root_path: &Option<PathBuf>,
+) -> Result<Workspace, LspError> {
+    // If there's a LSP root path, starting from file_path go up the directory tree
+    // searching for Nargo.toml files. The last one we find is the one we'll use
+    // (we'll assume Noir workspaces aren't nested)
+    if let Some(root_path) = root_path {
+        let mut current_path = file_path;
+        let mut current_toml_path = None;
+        while current_path.starts_with(root_path) {
+            if let Some(toml_path) = find_file_manifest(current_path) {
+                current_toml_path = Some(toml_path);
+
+                if let Some(next_path) = current_path.parent() {
+                    current_path = next_path;
+                } else {
+                    break;
+                }
+            } else {
+                break;
+            }
+        }
+
+        if let Some(toml_path) = current_toml_path {
+            return resolve_workspace_from_toml(
+                &toml_path,
+                PackageSelection::All,
+                Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
+            )
+            .map_err(|err| LspError::WorkspaceResolutionError(err.to_string()));
+        }
     }
+
+    let Some(parent_folder) = file_path
+        .parent()
+        .and_then(|f| f.file_name())
+        .and_then(|file_name_os_str| file_name_os_str.to_str())
+    else {
+        return Err(LspError::WorkspaceResolutionError(format!(
+            "Could not resolve parent folder for file: {:?}",
+            file_path
+        )));
+    };
+    let assumed_package = Package {
+        version: None,
+        compiler_required_version: Some(NOIR_ARTIFACT_VERSION_STRING.to_string()),
+        root_dir: PathBuf::from(parent_folder),
+        package_type: PackageType::Binary,
+        entry_path: PathBuf::from(file_path),
+        name: CrateName::from_str(parent_folder)
+            .map_err(|err| LspError::WorkspaceResolutionError(err.to_string()))?,
+        dependencies: BTreeMap::new(),
+    };
+    let workspace = Workspace {
+        root_dir: PathBuf::from(parent_folder),
+        members: vec![assumed_package],
+        selected_package_index: Some(0),
+        is_assumed: true,
+    };
+    Ok(workspace)
+}
+
+pub(crate) fn workspace_package_for_file<'a>(
+    workspace: &'a Workspace,
+    file_path: &Path,
+) -> Option<&'a Package> {
+    workspace.members.iter().filter(|package| file_path.starts_with(&package.root_dir)).next()
 }
 
 pub(crate) fn prepare_package<'file_manager, 'parsed_files>(

tooling/lsp/src/notifications/mod.rs

@@ -14,7 +14,7 @@ use crate::types::{
 
 use crate::{
     byte_span_to_range, get_package_tests_in_crate, parse_diff, prepare_source,
-    resolve_workspace_for_source_path, LspState,
+    resolve_workspace_for_source_path, workspace_package_for_file, LspState,
 };
 
 pub(super) fn on_initialized(
@@ -55,11 +55,14 @@ pub(super) fn on_did_change_text_document(
     let text = params.content_changes.into_iter().next().unwrap().text;
     state.input_files.insert(params.text_document.uri.to_string(), text.clone());
 
+    let file_path = params.text_document.uri.to_file_path().unwrap();
+
     let (mut context, crate_id) = prepare_source(text, state);
     let _ = check_crate(&mut context, crate_id, false, false, false);
 
     let workspace = match resolve_workspace_for_source_path(
         params.text_document.uri.to_file_path().unwrap().as_path(),
+        &state.root_path,
     ) {
         Ok(workspace) => workspace,
         Err(lsp_error) => {
@@ -70,7 +73,8 @@ pub(super) fn on_did_change_text_document(
             .into()))
         }
     };
-    let package = match workspace.members.first() {
+
+    let package = match workspace_package_for_file(&workspace, &file_path) {
         Some(package) => package,
         None => {
             return ControlFlow::Break(Err(ResponseError::new(
@@ -124,9 +128,10 @@ fn process_noir_document(
         ResponseError::new(ErrorCode::REQUEST_FAILED, "URI is not a valid file path")
     })?;
 
-    let workspace = resolve_workspace_for_source_path(&file_path).map_err(|lsp_error| {
-        ResponseError::new(ErrorCode::REQUEST_FAILED, lsp_error.to_string())
-    })?;
+    let workspace =
+        resolve_workspace_for_source_path(&file_path, &state.root_path).map_err(|lsp_error| {
+            ResponseError::new(ErrorCode::REQUEST_FAILED, lsp_error.to_string())
+        })?;
 
     let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
     insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);

tooling/lsp/src/requests/code_lens_request.rs

@@ -61,8 +61,12 @@ fn on_code_lens_request_inner(
         ResponseError::new(ErrorCode::REQUEST_FAILED, "Could not read file from disk")
     })?;
 
-    let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
-    let package = workspace.members.first().unwrap();
+    let workspace =
+        resolve_workspace_for_source_path(file_path.as_path(), &state.root_path).unwrap();
+
+    let package = crate::workspace_package_for_file(&workspace, &file_path).ok_or_else(|| {
+        ResponseError::new(ErrorCode::REQUEST_FAILED, "Could not find package for file")
+    })?;
 
     let (mut context, crate_id) = prepare_source(source_string, state);
     // We ignore the warnings and errors produced by compilation for producing code lenses

tooling/lsp/src/requests/goto_declaration.rs

@@ -20,7 +20,7 @@ fn on_goto_definition_inner(
     state: &mut LspState,
     params: GotoDeclarationParams,
 ) -> Result<GotoDeclarationResult, ResponseError> {
-    process_request(state, params.text_document_position_params, |location, interner, files| {
+    process_request(state, params.text_document_position_params, |location, interner, files, _| {
         interner.get_declaration_location_from(location).and_then(|found_location| {
             let file_id = found_location.file;
             let definition_position = to_lsp_location(files, file_id, found_location.span)?;

tooling/lsp/src/requests/goto_definition.rs

@@ -29,7 +29,7 @@ fn on_goto_definition_inner(
     params: GotoDefinitionParams,
     return_type_location_instead: bool,
 ) -> Result<GotoDefinitionResult, ResponseError> {
-    process_request(state, params.text_document_position_params, |location, interner, files| {
+    process_request(state, params.text_document_position_params, |location, interner, files, _| {
         interner.get_definition_location_from(location, return_type_location_instead).and_then(
             |found_location| {
                 let file_id = found_location.file;

tooling/lsp/src/requests/mod.rs

@@ -1,4 +1,4 @@
-use std::future::Future;
+use std::{collections::HashMap, future::Future};
 
 use crate::{
     parse_diff, resolve_workspace_for_source_path,
@@ -270,15 +270,18 @@ pub(crate) fn process_request<F, T>(
     callback: F,
 ) -> Result<T, ResponseError>
 where
-    F: FnOnce(noirc_errors::Location, &NodeInterner, &FileMap) -> T,
+    F: FnOnce(noirc_errors::Location, &NodeInterner, &FileMap, &HashMap<String, NodeInterner>) -> T,
 {
     let file_path =
         text_document_position_params.text_document.uri.to_file_path().map_err(|_| {
             ResponseError::new(ErrorCode::REQUEST_FAILED, "URI is not a valid file path")
         })?;
 
-    let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
-    let package = workspace.members.first().unwrap();
+    let workspace =
+        resolve_workspace_for_source_path(file_path.as_path(), &state.root_path).unwrap();
+    let package = crate::workspace_package_for_file(&workspace, &file_path).ok_or_else(|| {
+        ResponseError::new(ErrorCode::REQUEST_FAILED, "Could not find package for file")
+    })?;
 
     let package_root_path: String = package.root_dir.as_os_str().to_string_lossy().into();
 
@@ -306,7 +309,81 @@ where
         &text_document_position_params.position,
     )?;
 
-    Ok(callback(location, interner, files))
+    Ok(callback(location, interner, files, &state.cached_definitions))
+}
+pub(crate) fn find_all_references_in_workspace(
+    location: noirc_errors::Location,
+    interner: &NodeInterner,
+    cached_interners: &HashMap<String, NodeInterner>,
+    files: &FileMap,
+    include_declaration: bool,
+    include_self_type_name: bool,
+) -> Option<Vec<Location>> {
+    // First find the node that's referenced by the given location, if any
+    let referenced = interner.find_referenced(location);
+
+    if let Some(referenced) = referenced {
+        // If we found the referenced node, find its location
+        let referenced_location = interner.reference_location(referenced);
+
+        // Now we find all references that point to this location, in all interners
+        // (there's one interner per package, and all interners in a workspace rely on the
+        // same FileManager so a Location/FileId in one package is the same as in another package)
+        let mut locations = find_all_references(
+            referenced_location,
+            interner,
+            files,
+            include_declaration,
+            include_self_type_name,
+        );
+        for interner in cached_interners.values() {
+            locations.extend(find_all_references(
+                referenced_location,
+                interner,
+                files,
+                include_declaration,
+                include_self_type_name,
+            ));
+        }
+
+        // The LSP client usually removes duplicate loctions, but we do it here just in case they don't
+        locations.sort_by_key(|location| {
+            (
+                location.uri.to_string(),
+                location.range.start.line,
+                location.range.start.character,
+                location.range.end.line,
+                location.range.end.character,
+            )
+        });
+        locations.dedup();
+
+        if locations.is_empty() {
+            None
+        } else {
+            Some(locations)
+        }
+    } else {
+        None
+    }
+}
+
+pub(crate) fn find_all_references(
+    referenced_location: noirc_errors::Location,
+    interner: &NodeInterner,
+    files: &FileMap,
+    include_declaration: bool,
+    include_self_type_name: bool,
+) -> Vec<Location> {
+    interner
+        .find_all_references(referenced_location, include_declaration, include_self_type_name)
+        .map(|locations| {
+            locations
+                .iter()
+                .filter_map(|location| to_lsp_location(files, location.file, location.span))
+                .collect()
+        })
+        .unwrap_or(vec![])
 }
 
 #[cfg(test)]

tooling/lsp/src/requests/references.rs

@@ -5,23 +5,27 @@ use lsp_types::{Location, ReferenceParams};
 
 use crate::LspState;
 
-use super::{process_request, to_lsp_location};
+use super::{find_all_references_in_workspace, process_request};
 
 pub(crate) fn on_references_request(
     state: &mut LspState,
     params: ReferenceParams,
 ) -> impl Future<Output = Result<Option<Vec<Location>>, ResponseError>> {
-    let result =
-        process_request(state, params.text_document_position, |location, interner, files| {
-            interner.find_all_references(location, params.context.include_declaration, true).map(
-                |locations| {
-                    locations
-                        .iter()
-                        .filter_map(|location| to_lsp_location(files, location.file, location.span))
-                        .collect()
-                },
+    let include_declaration = params.context.include_declaration;
+    let result = process_request(
+        state,
+        params.text_document_position,
+        |location, interner, files, cached_interners| {
+            find_all_references_in_workspace(
+                location,
+                interner,
+                cached_interners,
+                files,
+                include_declaration,
+                true,
             )
-        });
+        },
+    );
     future::ready(result)
 }