Add "Go to type definition" capability

CHANGELOG.md

@@ -19,6 +19,10 @@
 
 ### Language server
 
+- The language server now has the ability to jump to the type definition of any
+  hovered value.
+  ([Giacomo Cavalieri](https://github.com/giacomocavalieri))
+
 - The language server now offers a code action to convert the first step of a
   pipeline to a regular function call. For example, this code:
 

compiler-core/src/analyse.rs

@@ -74,7 +74,7 @@ impl<T> Inferred<T> {
 }
 
 impl Inferred<PatternConstructor> {
-    pub fn definition_location(&self) -> Option<DefinitionLocation<'_>> {
+    pub fn definition_location(&self) -> Option<DefinitionLocation> {
         match self {
             Inferred::Known(value) => value.definition_location(),
             Inferred::Unknown => None,

compiler-core/src/ast.rs

@@ -1731,8 +1731,8 @@ impl SrcSpan {
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
-pub struct DefinitionLocation<'module> {
-    pub module: Option<&'module str>,
+pub struct DefinitionLocation {
+    pub module: Option<EcoString>,
     pub span: SrcSpan,
 }
 
@@ -1912,7 +1912,7 @@ impl<A> Pattern<A> {
 }
 
 impl TypedPattern {
-    pub fn definition_location(&self) -> Option<DefinitionLocation<'_>> {
+    pub fn definition_location(&self) -> Option<DefinitionLocation> {
         match self {
             Pattern::Int { .. }
             | Pattern::Float { .. }
@@ -2484,7 +2484,7 @@ impl TypedStatement {
         }
     }
 
-    pub fn definition_location(&self) -> Option<DefinitionLocation<'_>> {
+    pub fn definition_location(&self) -> Option<DefinitionLocation> {
         match self {
             Statement::Expression(expression) => expression.definition_location(),
             Statement::Assignment(_) => None,

compiler-core/src/ast/typed.rs

@@ -550,7 +550,7 @@ impl TypedExpr {
         }
     }
 
-    pub fn definition_location(&self) -> Option<DefinitionLocation<'_>> {
+    pub fn definition_location(&self) -> Option<DefinitionLocation> {
         match self {
             TypedExpr::Fn { .. }
             | TypedExpr::Int { .. }
@@ -582,7 +582,7 @@ impl TypedExpr {
                 constructor,
                 ..
             } => Some(DefinitionLocation {
-                module: Some(module_name.as_str()),
+                module: Some(module_name.clone()),
                 span: constructor.location(),
             }),
 

compiler-core/src/build.rs

@@ -349,25 +349,25 @@ impl<'a> Located<'a> {
         &self,
         importable_modules: &'a im::HashMap<EcoString, type_::ModuleInterface>,
         type_: std::sync::Arc<Type>,
-    ) -> Option<DefinitionLocation<'_>> {
+    ) -> Option<DefinitionLocation> {
         type_constructor_from_modules(importable_modules, type_).map(|t| DefinitionLocation {
-            module: Some(&t.module),
+            module: Some(t.module.clone()),
             span: t.origin,
         })
     }
 
     pub fn definition_location(
         &self,
         importable_modules: &'a im::HashMap<EcoString, type_::ModuleInterface>,
-    ) -> Option<DefinitionLocation<'_>> {
+    ) -> Option<DefinitionLocation> {
         match self {
             Self::PatternSpread { .. } => None,
             Self::Pattern(pattern) => pattern.definition_location(),
             Self::Statement(statement) => statement.definition_location(),
             Self::FunctionBody(statement) => None,
             Self::Expression(expression) => expression.definition_location(),
             Self::ModuleStatement(Definition::Import(import)) => Some(DefinitionLocation {
-                module: Some(import.module.as_str()),
+                module: Some(import.module.clone()),
                 span: SrcSpan { start: 0, end: 0 },
             }),
             Self::ModuleStatement(statement) => Some(DefinitionLocation {
@@ -382,12 +382,12 @@ impl<'a> Located<'a> {
             }) => importable_modules.get(*module).and_then(|m| {
                 if *is_type {
                     m.types.get(*name).map(|t| DefinitionLocation {
-                        module: Some(&module),
+                        module: Some((*module).clone()),
                         span: t.origin,
                     })
                 } else {
                     m.values.get(*name).map(|v| DefinitionLocation {
-                        module: Some(&module),
+                        module: Some((*module).clone()),
                         span: v.definition_location().span,
                     })
                 }
@@ -397,6 +397,106 @@ impl<'a> Located<'a> {
             Self::Label(_, _) => None,
         }
     }
+
+    pub(crate) fn type_(&self) -> Option<Arc<Type>> {
+        match self {
+            Located::Pattern(pattern) => Some(pattern.type_()),
+            Located::Statement(statement) => Some(statement.type_()),
+            Located::Expression(typed_expr) => Some(typed_expr.type_()),
+            Located::Arg(arg) => Some(arg.type_.clone()),
+            Located::Label(_, type_) | Located::Annotation(_, type_) => Some(type_.clone()),
+
+            Located::PatternSpread { .. } => None,
+            Located::ModuleStatement(definition) => None,
+            Located::FunctionBody(function) => None,
+            Located::UnqualifiedImport(unqualified_import) => None,
+        }
+    }
+
+    pub(crate) fn type_definition_locations(
+        &self,
+        importable_modules: &im::HashMap<EcoString, type_::ModuleInterface>,
+    ) -> Option<Vec<DefinitionLocation>> {
+        let type_ = self.type_()?;
+        Some(type_to_definition_locations(type_, importable_modules))
+    }
+}
+
+/// Returns the locations of all the types that one could reach starting from
+/// the given type (included). This includes all types that are part of a
+/// tuple/function type or that are used as args in a named type.
+///
+/// For example, given this type `Dict(Int, #(Wibble, Wobble))` all the
+/// "reachable" include: `Dict`, `Int`, `Wibble` and `Wobble`.
+///
+/// This is what powers the "go to type definition" capability of the language
+/// server.
+///
+fn type_to_definition_locations<'a>(
+    type_: Arc<Type>,
+    importable_modules: &'a im::HashMap<EcoString, type_::ModuleInterface>,
+) -> Vec<DefinitionLocation> {
+    match type_.as_ref() {
+        // For named types we start with the location of the named type itself
+        // followed by the locations of all types they reference in their args.
+        //
+        // For example with a `Dict(Wibble, Wobble)` we'd start with the
+        // definition of `Dict`, followed by the definition of `Wibble` and
+        // `Wobble`.
+        //
+        Type::Named {
+            module, name, args, ..
+        } => {
+            let Some(module) = importable_modules.get(module) else {
+                return vec![];
+            };
+
+            let Some(type_) = module.get_public_type(&name) else {
+                return vec![];
+            };
+
+            let mut locations = vec![DefinitionLocation {
+                module: Some(module.name.clone()),
+                span: type_.origin,
+            }];
+            for arg in args {
+                locations.extend(type_to_definition_locations(
+                    arg.clone(),
+                    importable_modules,
+                ));
+            }
+            locations
+        }
+
+        // For fn types we just get the locations of their arguments and return
+        // type.
+        //
+        Type::Fn { args, retrn } => args
+            .iter()
+            .flat_map(|arg| type_to_definition_locations(arg.clone(), importable_modules))
+            .chain(type_to_definition_locations(
+                retrn.clone(),
+                importable_modules,
+            ))
+            .collect_vec(),
+
+        // In case of a var we just follow it and get the locations of the type
+        // it points to.
+        //
+        Type::Var { type_ } => match type_.borrow().clone() {
+            type_::TypeVar::Unbound { .. } | type_::TypeVar::Generic { .. } => vec![],
+            type_::TypeVar::Link { type_ } => {
+                type_to_definition_locations(type_, importable_modules)
+            }
+        },
+
+        // In case of tuples we get the locations of the wrapped types.
+        //
+        Type::Tuple { elems } => elems
+            .iter()
+            .flat_map(|elem| type_to_definition_locations(elem.clone(), importable_modules))
+            .collect_vec(),
+    }
 }
 
 // Looks up the type constructor for the given type

compiler-core/src/language_server/engine.rs

@@ -1,8 +1,8 @@
 use crate::{
     analyse::name::correct_name_case,
     ast::{
-        ArgNames, CustomType, Definition, ModuleConstant, Pattern, SrcSpan, TypedArg, TypedExpr,
-        TypedFunction, TypedModule, TypedPattern,
+        ArgNames, CustomType, Definition, DefinitionLocation, ModuleConstant, Pattern, SrcSpan,
+        TypedArg, TypedExpr, TypedFunction, TypedModule, TypedPattern,
     },
     build::{type_constructor_from_modules, Located, Module, UnqualifiedImport},
     config::PackageConfig,
@@ -182,31 +182,64 @@ where
                 None => return Ok(None),
             };
 
-            let location = match node
-                .definition_location(this.compiler.project_compiler.get_importable_modules())
-            {
+            let Some(location) =
+                node.definition_location(this.compiler.project_compiler.get_importable_modules())
+            else {
+                return Ok(None);
+            };
+
+            Ok(this.definition_location_to_lsp_location(&line_numbers, &params, location))
+        })
+    }
+
+    pub(crate) fn goto_type_definition(
+        &mut self,
+        params: lsp_types::GotoDefinitionParams,
+    ) -> Response<Vec<lsp::Location>> {
+        self.respond(|this| {
+            let params = params.text_document_position_params;
+            let (line_numbers, node) = match this.node_at_position(&params) {
                 Some(location) => location,
-                None => return Ok(None),
+                None => return Ok(vec![]),
             };
 
-            let (uri, line_numbers) = match location.module {
-                None => (params.text_document.uri, &line_numbers),
-                Some(name) => {
-                    let module = match this.compiler.get_source(name) {
-                        Some(module) => module,
-                        _ => return Ok(None),
-                    };
-                    let url = Url::parse(&format!("file:///{}", &module.path))
-                        .expect("goto definition URL parse");
-                    (url, &module.line_numbers)
-                }
+            let Some(locations) = node
+                .type_definition_locations(this.compiler.project_compiler.get_importable_modules())
+            else {
+                return Ok(vec![]);
             };
-            let range = src_span_to_lsp_range(location.span, line_numbers);
 
-            Ok(Some(lsp::Location { uri, range }))
+            let locations = locations
+                .into_iter()
+                .filter_map(|location| {
+                    this.definition_location_to_lsp_location(&line_numbers, &params, location)
+                })
+                .collect_vec();
+
+            Ok(locations)
         })
     }
 
+    fn definition_location_to_lsp_location(
+        &self,
+        line_numbers: &LineNumbers,
+        params: &lsp_types::TextDocumentPositionParams,
+        location: DefinitionLocation,
+    ) -> Option<lsp::Location> {
+        let (uri, line_numbers) = match location.module {
+            None => (params.text_document.uri.clone(), line_numbers),
+            Some(name) => {
+                let module = self.compiler.get_source(&name)?;
+                let url = Url::parse(&format!("file:///{}", &module.path))
+                    .expect("goto definition URL parse");
+                (url, &module.line_numbers)
+            }
+        };
+        let range = src_span_to_lsp_range(location.span, line_numbers);
+
+        Some(lsp::Location { uri, range })
+    }
+
     pub fn completion(
         &mut self,
         params: lsp::TextDocumentPositionParams,

compiler-core/src/language_server/messages.rs

@@ -7,8 +7,8 @@ use lsp_types::{
     self as lsp,
     notification::{DidChangeTextDocument, DidCloseTextDocument, DidSaveTextDocument},
     request::{
-        CodeActionRequest, Completion, DocumentSymbolRequest, Formatting, HoverRequest,
-        PrepareRenameRequest, Rename, SignatureHelpRequest,
+        CodeActionRequest, Completion, DocumentSymbolRequest, Formatting, GotoTypeDefinition,
+        HoverRequest, PrepareRenameRequest, Rename, SignatureHelpRequest,
     },
 };
 use std::time::Duration;
@@ -24,6 +24,7 @@ pub enum Request {
     Format(lsp::DocumentFormattingParams),
     Hover(lsp::HoverParams),
     GoToDefinition(lsp::GotoDefinitionParams),
+    GoToTypeDefinition(lsp::GotoDefinitionParams),
     Completion(lsp::CompletionParams),
     CodeAction(lsp::CodeActionParams),
     SignatureHelp(lsp::SignatureHelpParams),
@@ -72,6 +73,10 @@ impl Request {
                 let params = cast_request::<PrepareRenameRequest>(request);
                 Some(Message::Request(id, Request::PrepareRename(params)))
             }
+            "textDocument/typeDefinition" => {
+                let params = cast_request::<GotoTypeDefinition>(request);
+                Some(Message::Request(id, Request::GoToTypeDefinition(params)))
+            }
             _ => None,
         }
     }

compiler-core/src/language_server/server.rs

@@ -107,6 +107,7 @@ where
             Request::DocumentSymbol(param) => self.document_symbol(param),
             Request::PrepareRename(param) => self.prepare_rename(param),
             Request::Rename(param) => self.rename(param),
+            Request::GoToTypeDefinition(param) => self.goto_type_definition(param),
         };
 
         self.publish_feedback(feedback);
@@ -308,6 +309,14 @@ where
         self.respond_with_engine(path, |engine| engine.goto_definition(params))
     }
 
+    fn goto_type_definition(
+        &mut self,
+        params: lsp_types::GotoDefinitionParams,
+    ) -> (Json, Feedback) {
+        let path = super::path(&params.text_document_position_params.text_document.uri);
+        self.respond_with_engine(path, |engine| engine.goto_type_definition(params))
+    }
+
     fn completion(&mut self, params: lsp::CompletionParams) -> (Json, Feedback) {
         let path = super::path(&params.text_document_position.text_document.uri);
 
@@ -418,7 +427,7 @@ fn initialisation_handshake(connection: &lsp_server::Connection) -> InitializePa
             },
         }),
         definition_provider: Some(lsp::OneOf::Left(true)),
-        type_definition_provider: None,
+        type_definition_provider: Some(lsp::TypeDefinitionProviderCapability::Simple(true)),
         implementation_provider: None,
         references_provider: None,
         document_highlight_provider: None,