fix goto definition for locations in other file

kls/Cargo.toml

@@ -16,7 +16,7 @@ crate-type = ["cdylib"]
 bench = false
 
 [dependencies]
-kanata-parser = { path = "../kanata-local/parser", features = ["cmd"]}
+kanata-parser = { path = "../kanata-local/parser", features = ["cmd", "gui"]}
 
 console_error_panic_hook = "0.1.7"
 js-sys = "0.3.53"

kls/src/helpers.rs

@@ -105,7 +105,6 @@ pub struct LocationInfo {
 pub struct DefinitionLocations(pub kanata_parser::lsp_hints::DefinitionLocations);
 
 impl DefinitionLocations {
-    // Returns kind of reference and it's name at given position.
     pub fn search_references_at_position(&self, pos: &lsp_types::Position) -> Option<LocationInfo> {
         log!("looking for references @ {:?}", pos);
         for ((name, span), ref_kind) in chain!(
@@ -137,7 +136,6 @@ impl DefinitionLocations {
 pub struct ReferenceLocations(pub kanata_parser::lsp_hints::ReferenceLocations);
 
 impl ReferenceLocations {
-    // Returns kind of definiiton and it's name at given position.
     pub fn search_definitions_at_position(
         &self,
         pos: &lsp_types::Position,

kls/src/lib.rs

@@ -516,21 +516,22 @@ impl KanataLanguageServer {
         params: &GotoDefinitionParams,
     ) -> Option<GotoDefinitionResponse> {
         log!("========= on_go_to_definition_impl ========");
-        let (_, definition_locations_storage, reference_locations_storage) = self.parse();
+        let (_, definition_locations_per_doc, reference_locations_per_doc) = self.parse();
         let source_doc_uri = &params.text_document_position_params.text_document.uri;
 
-        let definition_link = match navigation::goto_definition(
+        let definition_link = match navigation::definition_location(
             &params.text_document_position_params.position,
-            definition_locations_storage.get(source_doc_uri)?,
-            reference_locations_storage.get(source_doc_uri)?,
+            source_doc_uri,
+            &definition_locations_per_doc,
+            &reference_locations_per_doc,
         ) {
             Some(x) => x,
             None => {
                 return Some(GotoDefinitionResponse::Link(self.on_references_impl(
                     &params.text_document_position_params.position,
                     source_doc_uri,
-                    definition_locations_storage.get(source_doc_uri)?,
-                    reference_locations_storage.get(source_doc_uri)?,
+                    &definition_locations_per_doc,
+                    &reference_locations_per_doc,
                 )?))
             }
         };
@@ -560,11 +561,15 @@ impl KanataLanguageServer {
         &mut self,
         position: &Position,
         source_doc_uri: &Url,
-        identifier_locations: &DefinitionLocations,
-        reference_locations: &ReferenceLocations,
+        definition_locations_by_doc: &HashMap<Url, DefinitionLocations>,
+        reference_locations_by_doc: &HashMap<Url, ReferenceLocations>,
     ) -> Option<Vec<LocationLink>> {
-        let references =
-            navigation::references(position, identifier_locations, reference_locations)?;
+        let references = navigation::references(
+            position,
+            source_doc_uri,
+            definition_locations_by_doc,
+            reference_locations_by_doc,
+        )?;
         log!("matching reference(s) found: {:#?}", references);
         references
             .iter()

kls/src/navigation/mod.rs

@@ -1,4 +1,6 @@
-use lsp_types::{Position, Range};
+use std::collections::HashMap;
+
+use lsp_types::{Position, Range, Url};
 
 use crate::{
     helpers::{lsp_range_from_span, DefinitionLocations, ReferenceKind, ReferenceLocations},
@@ -12,74 +14,99 @@ pub struct GotoDefinitionLink {
     pub target_filename: String,
 }
 
-pub fn goto_definition(
+pub fn definition_location(
     pos: &Position,
-    identifier_locations: &DefinitionLocations,
-    reference_locations: &ReferenceLocations,
+    source_doc: &Url,
+    definition_locations_by_doc: &HashMap<Url, DefinitionLocations>,
+    reference_locations_by_doc: &HashMap<Url, ReferenceLocations>,
 ) -> Option<GotoDefinitionLink> {
-    let location_info = match reference_locations.search_definitions_at_position(pos) {
+    let source_doc_reference_locations = reference_locations_by_doc.get(source_doc)?;
+
+    let location_info = match source_doc_reference_locations.search_definitions_at_position(pos) {
         Some(x) => x,
         None => return None,
     };
     log!("{:?}", &location_info);
 
-    use ReferenceKind::*;
-    let location_map = match location_info.ref_kind {
-        Alias => &identifier_locations.0.alias,
-        Variable => &identifier_locations.0.variable,
-        VirtualKey => &identifier_locations.0.virtual_key,
-        ChordGroup => &identifier_locations.0.chord_group,
-        Layer => &identifier_locations.0.layer,
-        Include => {
-            return {
-                // (ref_name, Range::default())
-                // ref_name here is included file name
-                Some(GotoDefinitionLink {
-                    source_range: location_info.source_range,
-                    target_range: Range::default(),
-                    target_filename: location_info.ref_name,
-                })
-            };
+    for (_, definition_locations) in definition_locations_by_doc.iter() {
+        use ReferenceKind::*;
+        let location_map = match location_info.ref_kind {
+            Alias => &definition_locations.0.alias,
+            Variable => &definition_locations.0.variable,
+            VirtualKey => &definition_locations.0.virtual_key,
+            ChordGroup => &definition_locations.0.chord_group,
+            Layer => &definition_locations.0.layer,
+            Include => {
+                return {
+                    Some(GotoDefinitionLink {
+                        source_range: location_info.source_range,
+                        target_range: Range::default(),
+                        target_filename: location_info.ref_name,
+                    })
+                };
+            }
+        };
+
+        let loc = location_map
+            .get(&location_info.ref_name)
+            .map(|span| GotoDefinitionLink {
+                source_range: location_info.source_range,
+                target_range: lsp_range_from_span(span),
+                target_filename: span.file_name(),
+            });
+        if loc.is_none() {
+            continue;
         }
-    };
-    location_map
-        .get(&location_info.ref_name)
-        .map(|span| GotoDefinitionLink {
-            source_range: location_info.source_range,
-            target_range: lsp_range_from_span(span),
-            target_filename: span.file_name(),
-        })
+
+        // We're not checking definition locations in other files, because
+        // there should be only 1 definition for an identifier anyway.
+        return loc;
+    }
+    None
 }
 
 pub fn references(
     pos: &Position,
-    identifier_locations: &DefinitionLocations,
-    reference_locations: &ReferenceLocations,
+    source_doc: &Url,
+    definition_locations_by_doc: &HashMap<Url, DefinitionLocations>,
+    reference_locations_by_doc: &HashMap<Url, ReferenceLocations>,
 ) -> Option<Vec<GotoDefinitionLink>> {
-    let location_info = match identifier_locations.search_references_at_position(pos) {
+    let source_doc_definition_locations = definition_locations_by_doc.get(source_doc)?;
+
+    let location_info = match source_doc_definition_locations.search_references_at_position(pos) {
         Some(x) => x,
         None => return None,
     };
     log!("{:?}", &location_info);
 
-    use ReferenceKind::*;
-    let location_map = match location_info.ref_kind {
-        Alias => &reference_locations.0.alias,
-        Variable => &reference_locations.0.variable,
-        VirtualKey => &reference_locations.0.virtual_key,
-        ChordGroup => &reference_locations.0.chord_group,
-        Layer => &reference_locations.0.layer,
-        Include => unreachable!("includes can't be backreferenced"),
-    };
+    let mut reference_links: Vec<GotoDefinitionLink> = Vec::new();
 
-    location_map.0.get(&location_info.ref_name).map(|spans| {
-        spans
-            .iter()
-            .map(|span| GotoDefinitionLink {
-                source_range: location_info.source_range,
-                target_range: lsp_range_from_span(span),
-                target_filename: span.file_name(),
-            })
-            .collect()
-    })
+    for (_, reference_locations) in reference_locations_by_doc.iter() {
+        use ReferenceKind::*;
+        let location_map = match location_info.ref_kind {
+            Alias => &reference_locations.0.alias,
+            Variable => &reference_locations.0.variable,
+            VirtualKey => &reference_locations.0.virtual_key,
+            ChordGroup => &reference_locations.0.chord_group,
+            Layer => &reference_locations.0.layer,
+            Include => unreachable!("includes can't be backreferenced"),
+        };
+        let locations: Option<Vec<_>> = location_map.0.get(&location_info.ref_name).map(|spans| {
+            spans
+                .iter()
+                .map(|span| GotoDefinitionLink {
+                    source_range: location_info.source_range,
+                    target_range: lsp_range_from_span(span),
+                    target_filename: span.file_name(),
+                })
+                .collect()
+        });
+        if let Some(locations) = locations {
+            reference_links.extend(locations)
+        }
+    }
+    if reference_links.is_empty() {
+        return None;
+    }
+    Some(reference_links)
 }