Refactor for less laziness (right now)

Cargo.lock

@@ -2639,6 +2639,7 @@ dependencies = [
  "parking_lot",
  "roc_can",
  "roc_load",
+ "roc_module",
  "roc_packaging",
  "roc_problem",
  "roc_region",

crates/compiler/can/src/traverse.rs

@@ -703,6 +703,16 @@ pub enum FoundSymbol {
     Symbol(Symbol),
 }
 
+impl FoundSymbol {
+    pub fn implementation_symbol(&self) -> Symbol {
+        match self {
+            FoundSymbol::Specialization(_, sym)
+            | FoundSymbol::AbilityMember(_, sym)
+            | FoundSymbol::Symbol(sym) => *sym,
+        }
+    }
+}
+
 /// Given an ability Foo implements foo : ..., returns (T, foo1) if the symbol at the given region is a
 /// Like [find_type_at], but descends into the narrowest node containing [position].
 pub fn find_closest_type_at(

crates/lang_srv/Cargo.toml

@@ -10,6 +10,7 @@ path = "src/server.rs"
 [dependencies]
 roc_can = { path = "../compiler/can" }
 roc_load = { path = "../compiler/load" }
+roc_module = { path = "../compiler/module" }
 roc_problem = { path = "../compiler/problem" }
 roc_region = { path = "../compiler/region" }
 roc_reporting = { path = "../reporting" }

crates/lang_srv/src/convert.rs

@@ -101,7 +101,7 @@ pub(crate) mod diag {
         fn into_lsp_diagnostic(self, feed: &'a Self::Feed) -> Option<Diagnostic>;
     }
 
-    impl IntoLspDiagnostic<'_> for LoadingProblem<'_> {
+    impl IntoLspDiagnostic<'_> for &LoadingProblem<'_> {
         type Feed = ();
 
         fn into_lsp_diagnostic(self, _feed: &()) -> Option<Diagnostic> {
@@ -132,22 +132,22 @@ pub(crate) mod diag {
                     msg = format!("Unexpected header: {}", header);
                 }
                 LoadingProblem::ChannelProblem(_) => {
-                    msg = format!("Internal error: message channel died");
+                    msg = "Internal error: message channel died".to_string();
                 }
                 LoadingProblem::ErrJoiningWorkerThreads => {
-                    msg = format!("Internal error: analysis worker threads died");
+                    msg = "Internal error: analysis worker threads died".to_string();
                 }
                 LoadingProblem::TriedToImportAppModule => {
-                    msg = format!("Attempted to import app module");
+                    msg = "Attempted to import app module".to_string();
                 }
                 LoadingProblem::FormattedReport(report) => {
-                    msg = report;
+                    msg = report.clone();
                 }
                 LoadingProblem::ImportCycle(_, _) => {
-                    msg = format!("Circular dependency between modules");
+                    msg = "Circular dependency between modules".to_string();
                 }
                 LoadingProblem::IncorrectModuleName(_) => {
-                    msg = format!("Incorrect module name");
+                    msg = "Incorrect module name".to_string();
                 }
                 LoadingProblem::CouldNotFindCacheDir => {
                     msg = format!(
@@ -187,8 +187,8 @@ pub(crate) mod diag {
                 .to_range(fmt.line_info);
 
             let report = roc_reporting::report::can_problem(
-                &fmt.alloc,
-                &fmt.line_info,
+                fmt.alloc,
+                fmt.line_info,
                 fmt.path.to_path_buf(),
                 self,
             );
@@ -222,8 +222,8 @@ pub(crate) mod diag {
                 .to_range(fmt.line_info);
 
             let report = roc_reporting::report::type_problem(
-                &fmt.alloc,
-                &fmt.line_info,
+                fmt.alloc,
+                fmt.line_info,
                 fmt.path.to_path_buf(),
                 self,
             )?;

crates/lang_srv/src/registry.rs

@@ -1,11 +1,16 @@
 use std::collections::HashMap;
 
 use bumpalo::Bump;
-use roc_load::{LoadedModule, LoadingProblem};
+use roc_can::{abilities::AbilitiesStore, expr::Declarations};
+use roc_load::LoadedModule;
+use roc_module::symbol::{Interns, ModuleId, Symbol};
 use roc_packaging::cache::{self, RocCacheDir};
-use roc_region::all::LineInfo;
+use roc_region::all::{LineInfo, Region};
 use roc_reporting::report::RocDocAllocator;
-use tower_lsp::lsp_types::{Diagnostic, Hover, HoverContents, MarkedString, Position, Range, Url};
+use roc_types::subs::Subs;
+use tower_lsp::lsp_types::{
+    Diagnostic, GotoDefinitionResponse, Hover, HoverContents, MarkedString, Position, Range, Url,
+};
 
 use crate::convert::{
     diag::{IntoLspDiagnostic, ProblemFmt},
@@ -18,60 +23,21 @@ pub(crate) enum DocumentChange {
 }
 
 #[derive(Debug)]
-struct Document {
-    url: Url,
-    source: String,
-
-    arena: Bump,
-
-    // Incrementally updated module, diagnostis, etc.
-    line_info: Option<LineInfo>,
-    module: Option<Result<LoadedModule, ()>>,
-    diagnostics: Option<Vec<Diagnostic>>,
+struct Analysis {
+    line_info: LineInfo,
+    module: Option<LoadedModule>,
+    diagnostics: Vec<Diagnostic>,
 }
 
-impl Document {
-    fn new(url: Url, source: String) -> Self {
-        Self {
-            url,
-            source,
-            arena: Bump::new(),
-
-            line_info: None,
-            module: None,
-            diagnostics: None,
-        }
-    }
-
-    fn prime(&mut self, source: String) {
-        self.source = source;
-        self.module = None;
-        self.diagnostics = None;
-    }
-
-    fn prime_line_info(&mut self) {
-        if self.line_info.is_none() {
-            self.line_info = Some(LineInfo::new(&self.source));
-        }
-    }
-
-    fn line_info(&self) -> &LineInfo {
-        self.line_info.as_ref().unwrap()
-    }
-
-    fn module(&mut self) -> Result<&mut LoadedModule, LoadingProblem<'_>> {
-        if let Some(Ok(module)) = &mut self.module {
-            // Safety: returning for time self is alive
-            return Ok(unsafe { std::mem::transmute(module) });
-        }
-
-        let fi = self.url.to_file_path().unwrap();
+impl Analysis {
+    fn new(url: &Url, source: &str, arena: &Bump) -> Self {
+        let fi = url.to_file_path().unwrap();
         let src_dir = fi.parent().unwrap().to_path_buf();
 
-        let loaded = roc_load::load_and_typecheck_str(
-            &self.arena,
+        let mut loaded = roc_load::load_and_typecheck_str(
+            arena,
             fi,
-            &self.source,
+            source,
             src_dir,
             roc_target::TargetInfo::default_x86_64(),
             roc_load::FunctionKind::LambdaSet,
@@ -80,41 +46,19 @@ impl Document {
             roc_reporting::report::DEFAULT_PALETTE,
         );
 
-        match loaded {
+        let line_info = LineInfo::new(source);
+
+        let diagnostics = match loaded.as_mut() {
             Ok(module) => {
-                self.module = Some(Ok(module));
-                Ok(self.module.as_mut().unwrap().as_mut().unwrap())
-            }
-            Err(problem) => {
-                self.module = Some(Err(()));
-                Err(problem)
-            }
-        }
-    }
-
-    fn diagnostics(&mut self) -> Vec<Diagnostic> {
-        if let Some(diagnostics) = &self.diagnostics {
-            return diagnostics.clone();
-        }
-
-        let loaded: Result<&'static mut LoadedModule, LoadingProblem> =
-            unsafe { std::mem::transmute(self.module()) };
-
-        let diagnostics = match loaded {
-            Ok(module) => {
-                let line_info = {
-                    self.prime_line_info();
-                    self.line_info()
-                };
-                let lines: Vec<_> = self.source.lines().collect();
+                let lines: Vec<_> = source.lines().collect();
 
                 let alloc = RocDocAllocator::new(&lines, module.module_id, &module.interns);
 
                 let mut all_problems = Vec::new();
-                let module_path = self.url.to_file_path().unwrap();
+                let module_path = url.to_file_path().unwrap();
                 let fmt = ProblemFmt {
                     alloc: &alloc,
-                    line_info,
+                    line_info: &line_info,
                     path: &module_path,
                 };
 
@@ -147,51 +91,95 @@ impl Document {
             }
         };
 
-        self.diagnostics = Some(diagnostics);
-        self.diagnostics.as_ref().unwrap().clone()
+        Self {
+            line_info: LineInfo::new(source),
+            module: loaded.ok(),
+            diagnostics,
+        }
+    }
+}
+
+#[derive(Debug)]
+struct Document {
+    url: Url,
+    source: String,
+
+    arena: Bump,
+
+    analysis: Analysis,
+}
+
+impl Document {
+    fn new(url: Url, source: String) -> Self {
+        let arena = Bump::new();
+        let analysis = Analysis::new(&url, &source, &arena);
+
+        Self {
+            url,
+            source,
+            arena,
+            analysis,
+        }
+    }
+
+    fn refresh(&mut self, source: String) {
+        self.source = source;
+        self.analysis = Analysis::new(&self.url, &self.source, &self.arena);
+    }
+
+    fn line_info(&self) -> &LineInfo {
+        &self.analysis.line_info
+    }
+
+    fn module(&mut self) -> Option<&mut LoadedModule> {
+        self.analysis.module.as_mut()
+    }
+
+    fn diagnostics(&mut self) -> Vec<Diagnostic> {
+        self.analysis.diagnostics.clone()
+    }
+
+    fn split_module(&mut self) -> Option<SplitModule<'_>> {
+        self.module()?.try_into().ok()
+    }
+
+    fn symbol_at(&mut self, position: Position) -> Option<Symbol> {
+        let line_info = self.line_info();
+
+        let region = Region::from_pos(position.to_roc_position(line_info));
+
+        let SplitModule {
+            decls,
+            abilities_store,
+            ..
+        } = self.split_module()?;
+
+        let found_symbol = roc_can::traverse::find_symbol_at(region, decls, abilities_store)?;
+
+        Some(found_symbol.implementation_symbol())
     }
 
     fn hover(&mut self, position: Position) -> Option<Hover> {
-        let line_info = {
-            self.prime_line_info();
-            self.line_info()
-        };
+        let line_info = self.line_info();
 
-        let region = position.to_roc_position(line_info);
+        let pos = position.to_roc_position(line_info);
 
-        let module = match self.module() {
-            Ok(module) => module,
-            Err(_) => {
-                return Some(Hover {
-                    contents: HoverContents::Scalar(MarkedString::String("bar".to_owned())),
-                    range: Some(Range::new(
-                        position,
-                        Position {
-                            line: position.line,
-                            character: position.character + 1,
-                        },
-                    )),
-                })
-            }
-        };
+        let SplitModule {
+            subs,
+            decls,
+            module_id,
+            interns,
+            ..
+        } = self.split_module()?;
 
-        let (subs, decls) = match module.typechecked.get_mut(&module.module_id) {
-            Some(m) => (&mut m.solved_subs, &m.decls),
-            None => match module.declarations_by_id.get(&module.module_id) {
-                Some(decls) => (&mut module.solved, decls),
-                None => return missing_hover(module, position),
-            },
-        };
+        let (region, var) = roc_can::traverse::find_closest_type_at(pos, decls)?;
 
-        let (region, var) = roc_can::traverse::find_closest_type_at(region, decls)?;
-
-        let subs = subs.inner_mut();
         let snapshot = subs.snapshot();
         let type_str = roc_types::pretty_print::name_and_print_var(
             var,
             subs,
-            module.module_id,
-            &module.interns,
+            module_id,
+            interns,
             roc_types::pretty_print::DebugPrint::NOTHING,
         );
         subs.rollback_to(snapshot);
@@ -205,6 +193,45 @@ impl Document {
     }
 }
 
+struct SplitModule<'a> {
+    subs: &'a mut Subs,
+    abilities_store: &'a AbilitiesStore,
+    decls: &'a Declarations,
+    module_id: ModuleId,
+    interns: &'a Interns,
+}
+
+impl<'a> TryFrom<&'a mut LoadedModule> for SplitModule<'a> {
+    type Error = ();
+
+    fn try_from(module: &'a mut LoadedModule) -> Result<Self, Self::Error> {
+        let module_id = module.module_id;
+        let interns = &module.interns;
+        let subs;
+        let abilities_store;
+        let decls;
+        if let Some(m) = module.typechecked.get_mut(&module.module_id) {
+            subs = &mut m.solved_subs;
+            abilities_store = &m.abilities_store;
+            decls = &m.decls;
+        } else if let Some(d) = module.declarations_by_id.get(&module.module_id) {
+            subs = &mut module.solved;
+            abilities_store = &module.abilities_store;
+            decls = d;
+        } else {
+            return Err(());
+        }
+
+        Ok(Self {
+            subs: subs.inner_mut(),
+            abilities_store,
+            decls,
+            module_id,
+            interns,
+        })
+    }
+}
+
 fn missing_hover(module: &mut LoadedModule, position: Position) -> Option<Hover> {
     Some(Hover {
         contents: HoverContents::Scalar(MarkedString::String(format!(
@@ -230,7 +257,7 @@ impl Registry {
     pub fn apply_change(&mut self, change: DocumentChange) {
         match change {
             DocumentChange::Modified(url, source) => match self.documents.get_mut(&url) {
-                Some(document) => document.prime(source),
+                Some(document) => document.refresh(source),
                 None => {
                     self.documents
                         .insert(url.clone(), Document::new(url, source));
@@ -249,4 +276,17 @@ impl Registry {
     pub fn hover(&mut self, document: &Url, position: Position) -> Option<Hover> {
         self.documents.get_mut(document).unwrap().hover(position)
     }
+
+    pub fn goto_definition(
+        &mut self,
+        document: &Url,
+        position: Position,
+    ) -> Option<GotoDefinitionResponse> {
+        let symbol = self
+            .documents
+            .get_mut(document)
+            .unwrap()
+            .symbol_at(position)?;
+        None
+    }
 }

crates/lang_srv/src/server.rs

@@ -28,19 +28,25 @@ impl RocLs {
     }
 
     pub fn capabilities() -> ServerCapabilities {
-        let text_document_sync = Some(TextDocumentSyncCapability::Options(
+        let text_document_sync = TextDocumentSyncCapability::Options(
             // TODO: later on make this incremental
             TextDocumentSyncOptions {
                 open_close: Some(true),
                 change: Some(TextDocumentSyncKind::FULL),
                 ..TextDocumentSyncOptions::default()
             },
-        ));
-        let hover_provider = Some(HoverProviderCapability::Simple(true));
+        );
+        let hover_provider = HoverProviderCapability::Simple(true);
+        let definition_provider = DefinitionOptions {
+            work_done_progress_options: WorkDoneProgressOptions {
+                work_done_progress: None,
+            },
+        };
 
         ServerCapabilities {
-            text_document_sync,
-            hover_provider,
+            text_document_sync: Some(text_document_sync),
+            hover_provider: Some(hover_provider),
+            definition_provider: Some(OneOf::Right(definition_provider)),
             ..ServerCapabilities::default()
         }
     }
@@ -117,10 +123,34 @@ impl LanguageServer for RocLs {
             work_done_progress_params: _,
         } = params;
 
-        match std::panic::catch_unwind(|| self.registry().hover(&text_document.uri, position)) {
-            Ok(h) => Ok(h),
-            Err(_) => Ok(None),
-        }
+        panic_wrapper(|| self.registry().hover(&text_document.uri, position))
+    }
+
+    async fn goto_definition(
+        &self,
+        params: GotoDefinitionParams,
+    ) -> Result<Option<GotoDefinitionResponse>> {
+        let GotoDefinitionParams {
+            text_document_position_params:
+                TextDocumentPositionParams {
+                    text_document,
+                    position,
+                },
+            work_done_progress_params: _,
+            partial_result_params: _,
+        } = params;
+
+        panic_wrapper(|| {
+            self.registry()
+                .goto_definition(&text_document.uri, position)
+        })
+    }
+}
+
+fn panic_wrapper<T>(f: impl FnOnce() -> Option<T> + std::panic::UnwindSafe) -> Result<Option<T>> {
+    match std::panic::catch_unwind(f) {
+        Ok(r) => Ok(r),
+        Err(_) => Err(tower_lsp::jsonrpc::Error::internal_error()),
     }
 }