LSP: add support for finding module fn references

compiler-core/src/language_server.rs

@@ -5,6 +5,7 @@ mod feedback;
 mod files;
 mod messages;
 mod progress;
+mod references;
 mod router;
 mod server;
 

compiler-core/src/language_server/engine.rs

@@ -19,13 +19,15 @@ use crate::{
 };
 use camino::Utf8PathBuf;
 use ecow::EcoString;
+use itertools::Itertools;
 use lsp::CodeAction;
 use lsp_types::{self as lsp, Hover, HoverContents, MarkedString, Url};
 use std::sync::Arc;
 use strum::IntoEnumIterator;
 
 use super::{
     code_action::{CodeActionBuilder, RedundantTupleInCaseSubject},
+    references::ReferenceSearcher,
     src_span_to_lsp_range, DownloadDependencies, MakeLocker,
 };
 
@@ -266,6 +268,81 @@ where
         })
     }
 
+    pub fn references(
+        &mut self,
+        params: lsp::ReferenceParams,
+    ) -> Response<Option<Vec<lsp::Location>>> {
+        self.respond(|this| {
+            let Some(module) =
+                this.module_for_uri(&params.text_document_position.text_document.uri)
+            else {
+                return Ok(None);
+            };
+
+            let Some((_, node)) =
+                this.module_node_at_position(&params.text_document_position, module)
+            else {
+                return Ok(None);
+            };
+
+            let Some((def_module, def_location)) = ({
+                // Convert to have a matching type
+                let modules: im::HashMap<_, _> = this
+                    .compiler
+                    .modules
+                    .iter()
+                    .map(|(name, module)| (name.clone(), module.ast.type_info.clone()))
+                    .collect();
+
+                node.definition_location(&modules).map(|def| {
+                    (
+                        def.module
+                            // Defaults to current module
+                            .unwrap_or(&module.name)
+                            .to_string(),
+                        def.span,
+                    )
+                })
+            }) else {
+                return Ok(None);
+            };
+
+            let searcher = ReferenceSearcher::new(
+                &this.compiler.modules,
+                &def_module,
+                def_location,
+                params.context.include_declaration,
+            );
+            let references = searcher.references();
+
+            let mut locations = vec![];
+
+            for (module_name, spans) in references
+                .into_iter()
+                // Sorted for deterministic ordering
+                .sorted_by_key(|(k, _)| (*k).to_string())
+            {
+                let Some((uri, line_numbers)) = this.module_uri_and_line_numbers(module_name)
+                else {
+                    // Best effort
+                    continue;
+                };
+
+                locations.extend(
+                    spans
+                        .into_iter()
+                        // Convert to LSP locations
+                        .map(|span| lsp::Location {
+                            uri: uri.clone(),
+                            range: src_span_to_lsp_range(span, line_numbers),
+                        }),
+                );
+            }
+
+            Ok(Some(locations))
+        })
+    }
+
     fn respond<T>(&mut self, handler: impl FnOnce(&mut Self) -> Result<T>) -> Response<T> {
         let result = handler(self);
         let warnings = self.take_warnings();
@@ -694,6 +771,12 @@ where
     fn root_package_name(&self) -> &str {
         self.compiler.project_compiler.config.name.as_str()
     }
+
+    fn module_uri_and_line_numbers(&self, name: &str) -> Option<(Url, &LineNumbers)> {
+        let module = self.compiler.get_source(name)?;
+        let uri = Url::parse(&format!("file:///{}", &module.path)).ok()?;
+        Some((uri, &module.line_numbers))
+    }
 }
 
 fn type_completion(

compiler-core/src/language_server/messages.rs

@@ -1,7 +1,7 @@
 use camino::Utf8PathBuf;
 use lsp::{
     notification::{DidChangeWatchedFiles, DidOpenTextDocument},
-    request::GotoDefinition,
+    request::{GotoDefinition, References},
 };
 use lsp_types::{
     self as lsp,
@@ -23,6 +23,7 @@ pub enum Request {
     GoToDefinition(lsp::GotoDefinitionParams),
     Completion(lsp::CompletionParams),
     CodeAction(lsp::CodeActionParams),
+    References(lsp::ReferenceParams),
 }
 
 impl Request {
@@ -49,6 +50,10 @@ impl Request {
                 let params = cast_request::<CodeActionRequest>(request);
                 Some(Message::Request(id, Request::CodeAction(params)))
             }
+            "textDocument/references" => {
+                let params = cast_request::<References>(request);
+                Some(Message::Request(id, Request::References(params)))
+            }
             _ => None,
         }
     }

compiler-core/src/language_server/references.rs

@@ -0,0 +1,112 @@
+use std::{collections::HashMap, sync::Arc};
+
+use ecow::EcoString;
+
+use crate::{
+    ast::{self, visit::Visit as _, SrcSpan},
+    build,
+    type_::{ModuleValueConstructor, Type, ValueConstructor, ValueConstructorVariant},
+};
+
+pub struct ReferenceSearcher<'a> {
+    modules: &'a HashMap<EcoString, build::Module>,
+    def_module: &'a str,
+    def_location: SrcSpan,
+    include_declaration: bool,
+    current_module: &'a str,
+    references: HashMap<&'a str, Vec<SrcSpan>>,
+}
+
+impl<'a> ReferenceSearcher<'a> {
+    pub fn new(
+        modules: &'a HashMap<EcoString, build::Module>,
+        def_module: &'a str,
+        def_location: SrcSpan,
+        include_declaration: bool,
+    ) -> ReferenceSearcher<'a> {
+        Self {
+            modules,
+            def_module,
+            def_location,
+            include_declaration,
+            current_module: "",
+            references: Default::default(),
+        }
+    }
+
+    pub fn references(mut self) -> HashMap<&'a str, Vec<SrcSpan>> {
+        if self.include_declaration {
+            self.current_module = self.def_module;
+            self.add_reference(self.def_location);
+        }
+
+        // Collect all references from all modules
+        for module in self.modules.values() {
+            self.current_module = &module.name;
+            self.visit_typed_module(&module.ast);
+        }
+
+        self.references
+    }
+
+    /// Adds location on the current module to references
+    fn add_reference(&mut self, location: SrcSpan) {
+        self.references
+            .entry(self.current_module)
+            .or_default()
+            .push(location);
+    }
+}
+
+impl<'ast> ast::visit::Visit<'ast> for ReferenceSearcher<'_> {
+    fn visit_typed_expr_var(
+        &mut self,
+        location: &'ast SrcSpan,
+        constructor: &'ast ValueConstructor,
+        name: &'ast EcoString,
+    ) {
+        match &constructor.variant {
+            ValueConstructorVariant::ModuleFn {
+                module: def_module,
+                location: def_location,
+                ..
+            } if def_module == self.def_module && *def_location == self.def_location => {
+                self.add_reference(*location);
+            }
+            _ => {}
+        }
+
+        ast::visit::visit_typed_expr_var(self, location, constructor, name);
+    }
+
+    fn visit_typed_expr_module_select(
+        &mut self,
+        location: &'ast SrcSpan,
+        typ: &'ast Arc<Type>,
+        label: &'ast EcoString,
+        module_name: &'ast EcoString,
+        module_alias: &'ast EcoString,
+        constructor: &'ast ModuleValueConstructor,
+    ) {
+        match constructor {
+            ModuleValueConstructor::Fn {
+                location: def_location,
+                module: def_module,
+                ..
+            } if def_module == self.def_module && *def_location == self.def_location => {
+                self.add_reference(*location);
+            }
+            _ => {}
+        }
+
+        ast::visit::visit_typed_expr_module_select(
+            self,
+            location,
+            typ,
+            label,
+            module_name,
+            module_alias,
+            constructor,
+        );
+    }
+}

compiler-core/src/language_server/server.rs

@@ -101,6 +101,7 @@ where
             Request::GoToDefinition(param) => self.goto_definition(param),
             Request::Completion(param) => self.completion(param),
             Request::CodeAction(param) => self.code_action(param),
+            Request::References(param) => self.references(param),
         };
 
         self.publish_feedback(feedback);
@@ -319,6 +320,11 @@ where
         self.respond_with_engine(path, |engine| engine.action(params))
     }
 
+    fn references(&mut self, params: lsp::ReferenceParams) -> (Json, Feedback) {
+        let path = super::path(&params.text_document_position.text_document.uri);
+        self.respond_with_engine(path, |engine| engine.references(params))
+    }
+
     fn cache_file_in_memory(&mut self, path: Utf8PathBuf, text: String) -> Feedback {
         self.project_changed(&path);
         if let Err(error) = self.io.write_mem_cache(&path, &text) {
@@ -388,7 +394,7 @@ fn initialisation_handshake(connection: &lsp_server::Connection) -> InitializePa
         definition_provider: Some(lsp::OneOf::Left(true)),
         type_definition_provider: None,
         implementation_provider: None,
-        references_provider: None,
+        references_provider: Some(lsp::OneOf::Left(true)),
         document_highlight_provider: None,
         document_symbol_provider: None,
         workspace_symbol_provider: None,

compiler-core/src/language_server/tests/mod.rs

@@ -3,6 +3,7 @@ mod compilation;
 mod completion;
 mod definition;
 mod hover;
+mod references;
 
 use std::{
     collections::HashMap,