feat: support goto/find references inside of modules (#34)

* feat: support goto declaration inside of modules

* fix: change goto declarations provider to references provider

* fix: redact uri in references response

crates/tinymist-query/src/analysis.rs

@@ -10,6 +10,8 @@ pub mod reference;
 pub use reference::*;
 pub mod def_use;
 pub use def_use::*;
+pub mod matcher;
+pub use matcher::*;
 
 #[cfg(test)]
 mod lexical_hierarchy_tests {

crates/tinymist-query/src/analysis/def_use.rs

@@ -6,6 +6,7 @@ use std::{
 };
 
 use comemo::Tracked;
+use log::info;
 use parking_lot::Mutex;
 use serde::Serialize;
 use typst::{syntax::Source, World};
@@ -26,8 +27,8 @@ enum Ns {
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct IdentRef {
-    name: String,
-    range: Range<usize>,
+    pub name: String,
+    pub range: Range<usize>,
 }
 
 impl PartialOrd for IdentRef {
@@ -72,6 +73,19 @@ pub struct DefUseInfo {
     exports_refs: Vec<DefId>,
 }
 
+impl DefUseInfo {
+    pub fn get_def(&self, fid: TypstFileId, ident: &IdentRef) -> Option<(DefId, &IdentDef)> {
+        let (id, _, def) = self.ident_defs.get_full(&(fid, ident.clone()))?;
+        Some((DefId(id as u64), def))
+    }
+
+    pub fn get_refs(&self, id: DefId) -> impl Iterator<Item = &IdentRef> {
+        self.ident_refs
+            .iter()
+            .filter_map(move |(k, v)| if *v == id { Some(k) } else { None })
+    }
+}
+
 pub fn get_def_use(world: Tracked<'_, dyn World>, source: Source) -> Option<Arc<DefUseInfo>> {
     let ctx = SearchCtx {
         world,
@@ -171,6 +185,7 @@ impl<'a, 'w> DefUseCollector<'a, 'w> {
                         let external_info =
                             find_source_by_import_path(self.ctx.world, self.current_id, path)
                                 .and_then(|source| {
+                                    info!("diving source for def use: {:?}", source.id());
                                     Some(source.id()).zip(get_def_use_inner(self.ctx, source))
                                 });
 

crates/tinymist-query/src/analysis/definition.rs

@@ -15,7 +15,7 @@ use typst::{
 };
 use typst_ts_core::TypstFileId;
 
-use crate::analysis::find_source_by_import;
+use crate::analysis::{deref_lvalue, find_source_by_import};
 use crate::{prelude::*, TypstSpan};
 
 #[derive(Debug, Clone)]
@@ -73,13 +73,6 @@ impl Definition<'_> {
     }
 }
 
-fn deref_lvalue(mut node: LinkedNode) -> Option<LinkedNode> {
-    while let Some(e) = node.cast::<ast::Parenthesized>() {
-        node = node.find(e.expr().span())?;
-    }
-    Some(node)
-}
-
 fn advance_prev_adjacent(node: LinkedNode) -> Option<LinkedNode> {
     // this is aworkaround for a bug in the parser
     if node.len() == 0 {

crates/tinymist-query/src/analysis/matcher.rs

@@ -0,0 +1,69 @@
+use log::debug;
+use typst::syntax::{
+    ast::{self, AstNode},
+    LinkedNode, SyntaxKind,
+};
+
+pub fn deref_lvalue(mut node: LinkedNode) -> Option<LinkedNode> {
+    while let Some(e) = node.cast::<ast::Parenthesized>() {
+        node = node.find(e.expr().span())?;
+    }
+    Some(node)
+}
+
+pub enum DerefTarget<'a> {
+    VarAccess(LinkedNode<'a>),
+    Callee(LinkedNode<'a>),
+    ImportPath(LinkedNode<'a>),
+}
+
+impl<'a> DerefTarget<'a> {
+    pub fn node(&self) -> &LinkedNode {
+        match self {
+            DerefTarget::VarAccess(node) => node,
+            DerefTarget::Callee(node) => node,
+            DerefTarget::ImportPath(node) => node,
+        }
+    }
+}
+
+pub fn get_deref_target(node: LinkedNode) -> Option<DerefTarget> {
+    let mut ancestor = node;
+    while !ancestor.is::<ast::Expr>() {
+        ancestor = ancestor.parent()?.clone();
+    }
+    debug!("deref expr: {ancestor:?}");
+    let ancestor = deref_lvalue(ancestor)?;
+    debug!("deref lvalue: {ancestor:?}");
+
+    let may_ident = ancestor.cast::<ast::Expr>()?;
+    if !may_ident.hash() && !matches!(may_ident, ast::Expr::MathIdent(_)) {
+        return None;
+    }
+
+    Some(match may_ident {
+        // todo: label, reference
+        // todo: import
+        // todo: include
+        ast::Expr::FuncCall(call) => DerefTarget::Callee(ancestor.find(call.callee().span())?),
+        ast::Expr::Set(set) => DerefTarget::Callee(ancestor.find(set.target().span())?),
+        ast::Expr::Ident(..) | ast::Expr::MathIdent(..) | ast::Expr::FieldAccess(..) => {
+            DerefTarget::VarAccess(ancestor.find(may_ident.span())?)
+        }
+        ast::Expr::Str(..) => {
+            let parent = ancestor.parent()?;
+            if parent.kind() != SyntaxKind::ModuleImport {
+                return None;
+            }
+
+            return Some(DerefTarget::ImportPath(ancestor.find(may_ident.span())?));
+        }
+        ast::Expr::Import(..) => {
+            return None;
+        }
+        _ => {
+            debug!("unsupported kind {kind:?}", kind = ancestor.kind());
+            return None;
+        }
+    })
+}

crates/tinymist-query/src/fixtures/goto_definition/at_def.typ

@@ -0,0 +1,3 @@
+#let /* ident after */ f() = 1;
+#(f());
+#(f());

crates/tinymist-query/src/fixtures/goto_definition/snaps/test@at_def.typ.snap

@@ -0,0 +1,12 @@
+---
+source: crates/tinymist-query/src/goto_definition.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/goto_definition/at_def.typ
+---
+[
+ {
+  "originSelectionRange": "0:23:0:24",
+  "targetRange": "0:24:0:26",
+  "targetSelectionRange": "0:24:0:26"
+ }
+]

crates/tinymist-query/src/fixtures/references/at_def.typ

@@ -0,0 +1,3 @@
+#let /* ident after */ f() = 1;
+#(f());
+#(f());

crates/tinymist-query/src/fixtures/references/base.typ

@@ -0,0 +1,2 @@
+#let x = 1;
+#(/* position after */ x);

crates/tinymist-query/src/fixtures/references/redefine.typ

@@ -0,0 +1,10 @@
+#let /* ident after */ x = 1;
+
+#let y = {
+  x + x;
+}
+
+#let f(y) = x + y;
+
+#let x = x;
+#let f = x;

crates/tinymist-query/src/fixtures/references/snaps/test@at_def.typ.snap

@@ -0,0 +1,13 @@
+---
+source: crates/tinymist-query/src/references.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/references/at_def.typ
+---
+[
+ {
+  "range": "2:2:2:3"
+ },
+ {
+  "range": "1:2:1:3"
+ }
+]

crates/tinymist-query/src/fixtures/references/snaps/test@base.typ.snap

@@ -0,0 +1,6 @@
+---
+source: crates/tinymist-query/src/goto_declaration.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/goto_declaration/base.typ
+---
+null

crates/tinymist-query/src/fixtures/references/snaps/test@redefine.typ.snap

@@ -0,0 +1,19 @@
+---
+source: crates/tinymist-query/src/references.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/references/redefine.typ
+---
+[
+ {
+  "range": "8:9:8:10"
+ },
+ {
+  "range": "3:2:3:3"
+ },
+ {
+  "range": "3:6:3:7"
+ },
+ {
+  "range": "6:12:6:13"
+ }
+]

crates/tinymist-query/src/goto_declaration.rs

@@ -0,0 +1,69 @@
+use std::ops::Range;
+
+use comemo::Track;
+use log::debug;
+use lsp_types::LocationLink;
+
+use crate::{
+    analysis::{get_def_use, get_deref_target, DerefTarget},
+    prelude::*,
+};
+
+#[derive(Debug, Clone)]
+pub struct GotoDeclarationRequest {
+    pub path: PathBuf,
+    pub position: LspPosition,
+}
+
+impl GotoDeclarationRequest {
+    pub fn request(
+        self,
+        world: &TypstSystemWorld,
+        position_encoding: PositionEncoding,
+    ) -> Option<GotoDeclarationResponse> {
+        let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
+        let offset = lsp_to_typst::position(self.position, position_encoding, &source)?;
+        let cursor = offset + 1;
+
+        let w: &dyn World = world;
+        let ast_node = LinkedNode::new(source.root()).leaf_at(cursor)?;
+        debug!("ast_node: {ast_node:?}", ast_node = ast_node);
+        let deref_target = get_deref_target(ast_node)?;
+
+        let use_site = deref_target.node();
+        let origin_selection_range =
+            typst_to_lsp::range(use_site.range(), &source, position_encoding);
+
+        let def_use = get_def_use(w.track(), source.clone())?;
+        let ref_spans = find_declarations(w, def_use, deref_target)?;
+
+        let mut links = vec![];
+        for ref_range in ref_spans {
+            let ref_id = source.id();
+            let ref_source = &source;
+
+            let span_path = world.path_for_id(ref_id).ok()?;
+            let range = typst_to_lsp::range(ref_range, ref_source, position_encoding);
+
+            let uri = Url::from_file_path(span_path).ok()?;
+
+            links.push(LocationLink {
+                origin_selection_range: Some(origin_selection_range),
+                target_uri: uri,
+                target_range: range,
+                target_selection_range: range,
+            });
+        }
+
+        debug!("goto_declartion: {links:?}");
+        Some(GotoDeclarationResponse::Link(links))
+    }
+}
+
+fn find_declarations(
+    _w: &dyn World,
+    _def_use: Arc<crate::analysis::DefUseInfo>,
+    _deref_target: DerefTarget<'_>,
+) -> Option<Vec<Range<usize>>> {
+    todo!()
+}

crates/tinymist-query/src/lib.rs

@@ -8,36 +8,40 @@ use std::sync::Arc;
 use typst_ts_core::TypstDocument;
 
 pub use diagnostics::*;
-pub(crate) mod signature_help;
-pub use signature_help::*;
+pub(crate) mod code_lens;
+pub use code_lens::*;
+pub(crate) mod completion;
+pub use completion::*;
 pub(crate) mod document_symbol;
 pub use document_symbol::*;
-pub(crate) mod symbol;
-pub use symbol::*;
+pub(crate) mod folding_range;
+pub use folding_range::*;
+pub(crate) mod goto_declaration;
+pub use goto_declaration::*;
+pub(crate) mod goto_definition;
+pub use goto_definition::*;
+pub(crate) mod hover;
+pub use hover::*;
+pub(crate) mod inlay_hint;
+pub use inlay_hint::*;
+pub(crate) mod rename;
+pub use rename::*;
+pub(crate) mod selection_range;
+pub use selection_range::*;
 pub(crate) mod semantic_tokens;
 pub use semantic_tokens::*;
 pub(crate) mod semantic_tokens_full;
 pub use semantic_tokens_full::*;
 pub(crate) mod semantic_tokens_delta;
 pub use semantic_tokens_delta::*;
-pub(crate) mod hover;
-pub use hover::*;
-pub(crate) mod completion;
-pub use completion::*;
-pub(crate) mod folding_range;
-pub use folding_range::*;
-pub(crate) mod selection_range;
-pub use selection_range::*;
-pub(crate) mod goto_definition;
-pub use goto_definition::*;
-pub(crate) mod inlay_hint;
-pub use inlay_hint::*;
+pub(crate) mod signature_help;
+pub use signature_help::*;
+pub(crate) mod symbol;
+pub use symbol::*;
 pub(crate) mod prepare_rename;
 pub use prepare_rename::*;
-pub(crate) mod rename;
-pub use rename::*;
-pub(crate) mod code_lens;
-pub use code_lens::*;
+pub(crate) mod references;
+pub use references::*;
 
 pub mod lsp_typst_boundary;
 pub use lsp_typst_boundary::*;
@@ -78,6 +82,8 @@ mod polymorphic {
         OnSaveExport(OnSaveExportRequest),
         Hover(HoverRequest),
         GotoDefinition(GotoDefinitionRequest),
+        GotoDeclaration(GotoDeclarationRequest),
+        References(ReferencesRequest),
         InlayHint(InlayHintRequest),
         CodeLens(CodeLensRequest),
         Completion(CompletionRequest),
@@ -100,6 +106,8 @@ mod polymorphic {
                 CompilerQueryRequest::OnSaveExport(..) => Mergable,
                 CompilerQueryRequest::Hover(..) => PinnedFirst,
                 CompilerQueryRequest::GotoDefinition(..) => PinnedFirst,
+                CompilerQueryRequest::GotoDeclaration(..) => PinnedFirst,
+                CompilerQueryRequest::References(..) => PinnedFirst,
                 CompilerQueryRequest::InlayHint(..) => Unique,
                 CompilerQueryRequest::CodeLens(..) => Unique,
                 CompilerQueryRequest::Completion(..) => Mergable,
@@ -121,6 +129,8 @@ mod polymorphic {
                 CompilerQueryRequest::OnSaveExport(req) => &req.path,
                 CompilerQueryRequest::Hover(req) => &req.path,
                 CompilerQueryRequest::GotoDefinition(req) => &req.path,
+                CompilerQueryRequest::GotoDeclaration(req) => &req.path,
+                CompilerQueryRequest::References(req) => &req.path,
                 CompilerQueryRequest::InlayHint(req) => &req.path,
                 CompilerQueryRequest::CodeLens(req) => &req.path,
                 CompilerQueryRequest::Completion(req) => &req.path,
@@ -143,6 +153,8 @@ mod polymorphic {
         OnSaveExport(()),
         Hover(Option<Hover>),
         GotoDefinition(Option<GotoDefinitionResponse>),
+        GotoDeclaration(Option<GotoDeclarationResponse>),
+        References(Option<Vec<LspLocation>>),
         InlayHint(Option<Vec<InlayHint>>),
         CodeLens(Option<Vec<CodeLens>>),
         Completion(Option<CompletionResponse>),
@@ -171,7 +183,10 @@ mod tests {
     use once_cell::sync::Lazy;
     use serde::Serialize;
     use serde_json::{ser::PrettyFormatter, Serializer, Value};
-    use typst::syntax::{LinkedNode, Source, VirtualPath};
+    use typst::syntax::{
+        ast::{self, AstNode},
+        LinkedNode, Source, SyntaxKind, VirtualPath,
+    };
     use typst_ts_compiler::{
         service::{CompileDriver, Compiler, WorkspaceProvider},
         ShadowApi,
@@ -252,25 +267,86 @@ mod tests {
     }
 
     pub fn find_test_position(s: &Source) -> LspPosition {
-        let re = s.text().find("/* position */").map(|e| (e, true));
-        let re = re.or_else(|| s.text().find("/* position after */").zip(Some(false)));
-        let (re, prev) = re
+        enum AstMatcher {
+            MatchAny { prev: bool },
+            MatchIdent { prev: bool },
+        }
+        use AstMatcher::*;
+
+        let re = s
+            .text()
+            .find("/* position */")
+            .map(|e| (e, MatchAny { prev: true }));
+        let re = re.or_else(|| {
+            s.text()
+                .find("/* position after */")
+                .zip(Some(MatchAny { prev: false }))
+        });
+        let re = re.or_else(|| {
+            s.text()
+                .find("/* ident */")
+                .zip(Some(MatchIdent { prev: true }))
+        });
+        let re = re.or_else(|| {
+            s.text()
+                .find("/* ident after */")
+                .zip(Some(MatchIdent { prev: false }))
+        });
+        let (re, m) = re
             .ok_or_else(|| panic!("No position marker found in source:\n{}", s.text()))
             .unwrap();
 
         let n = LinkedNode::new(s.root());
         let mut n = n.leaf_at(re + 1).unwrap();
 
-        while n.kind().is_trivia() {
-            let m = if prev {
-                n.prev_sibling()
-            } else {
-                n.next_sibling()
-            };
-            n = m.or_else(|| n.parent().cloned()).unwrap();
+        let match_prev = match &m {
+            MatchAny { prev } => *prev,
+            MatchIdent { prev } => *prev,
+        };
+        let match_ident = match m {
+            MatchAny { .. } => false,
+            MatchIdent { .. } => true,
+        };
+
+        'match_loop: loop {
+            if n.kind().is_trivia() {
+                let m = if match_prev {
+                    n.prev_sibling()
+                } else {
+                    n.next_sibling()
+                };
+                n = m.or_else(|| n.parent().cloned()).unwrap();
+                continue;
+            }
+            if match_ident {
+                match n.kind() {
+                    SyntaxKind::Closure => {
+                        let c = n.cast::<ast::Closure>().unwrap();
+                        if let Some(name) = c.name() {
+                            if let Some(m) = n.find(name.span()) {
+                                n = m;
+                                break 'match_loop;
+                            }
+                        }
+                    }
+                    SyntaxKind::LetBinding => {
+                        let c = n.cast::<ast::LetBinding>().unwrap();
+                        if let Some(name) = c.kind().bindings().first() {
+                            if let Some(m) = n.find(name.span()) {
+                                n = m;
+                                break 'match_loop;
+                            }
+                        }
+                    }
+                    _ => {}
+                }
+            }
+            break;
         }
 
-        typst_to_lsp::offset_to_position(n.offset() + 1, PositionEncoding::Utf16, s)
+        // eprintln!("position: {:?} -> {:?}", n.offset(), n);
+
+        typst_to_lsp::offset_to_position(n.offset(), PositionEncoding::Utf16, s)
     }
 
     // pub static REDACT_URI: Lazy<RedactFields> = Lazy::new(||
@@ -278,6 +354,7 @@ mod tests {
     pub static REDACT_LOC: Lazy<RedactFields> = Lazy::new(|| {
         RedactFields::from_iter([
             "location",
+            "uri",
             "range",
             "selectionRange",
             "targetRange",

crates/tinymist-query/src/prelude.rs

@@ -9,9 +9,9 @@ pub use comemo::{Track, Tracked};
 pub use itertools::{Format, Itertools};
 pub use log::{error, trace};
 pub use lsp_types::{
-    CodeLens, CompletionResponse, DiagnosticRelatedInformation, DocumentSymbol,
-    DocumentSymbolResponse, Documentation, FoldingRange, GotoDefinitionResponse, Hover, InlayHint,
-    Location as LspLocation, MarkupContent, MarkupKind, Position as LspPosition,
+    request::GotoDeclarationResponse, CodeLens, CompletionResponse, DiagnosticRelatedInformation,
+    DocumentSymbol, DocumentSymbolResponse, Documentation, FoldingRange, GotoDefinitionResponse,
+    Hover, InlayHint, Location as LspLocation, MarkupContent, MarkupKind, Position as LspPosition,
     PrepareRenameResponse, SelectionRange, SemanticTokens, SemanticTokensDelta,
     SemanticTokensFullDeltaResult, SemanticTokensResult, SignatureHelp, SignatureInformation,
     SymbolInformation, Url, WorkspaceEdit,

crates/tinymist-query/src/references.rs

@@ -0,0 +1,125 @@
+use std::ops::Range;
+
+use comemo::Track;
+use log::debug;
+
+use crate::{
+    analysis::{get_def_use, get_deref_target, DerefTarget, IdentRef},
+    prelude::*,
+};
+
+#[derive(Debug, Clone)]
+pub struct ReferencesRequest {
+    pub path: PathBuf,
+    pub position: LspPosition,
+}
+
+impl ReferencesRequest {
+    pub fn request(
+        self,
+        world: &TypstSystemWorld,
+        position_encoding: PositionEncoding,
+    ) -> Option<Vec<LspLocation>> {
+        let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
+        let offset = lsp_to_typst::position(self.position, position_encoding, &source)?;
+        let cursor = offset + 1;
+
+        let w: &dyn World = world;
+        let ast_node = LinkedNode::new(source.root()).leaf_at(cursor)?;
+        debug!("ast_node: {ast_node:?}", ast_node = ast_node);
+        let deref_target = get_deref_target(ast_node)?;
+
+        let def_use = get_def_use(w.track(), source.clone())?;
+        let ref_spans = find_declarations(w, def_use, deref_target)?;
+
+        let mut locations = vec![];
+        for ref_range in ref_spans {
+            let ref_id = source.id();
+            let ref_source = &source;
+
+            let span_path = world.path_for_id(ref_id).ok()?;
+            let range = typst_to_lsp::range(ref_range, ref_source, position_encoding);
+
+            let uri = Url::from_file_path(span_path).ok()?;
+
+            locations.push(LspLocation { uri, range });
+        }
+
+        debug!("references: {locations:?}");
+        Some(locations)
+    }
+}
+
+fn find_declarations(
+    _w: &dyn World,
+    def_use: Arc<crate::analysis::DefUseInfo>,
+    deref_target: DerefTarget<'_>,
+) -> Option<Vec<Range<usize>>> {
+    let node = match deref_target {
+        DerefTarget::VarAccess(node) => node,
+        DerefTarget::Callee(node) => node,
+        DerefTarget::ImportPath(..) => {
+            return None;
+        }
+    };
+
+    let mut may_ident = node.cast::<ast::Expr>()?;
+    let name;
+    loop {
+        match may_ident {
+            ast::Expr::Parenthesized(e) => {
+                may_ident = e.expr();
+            }
+            ast::Expr::FieldAccess(e) => {
+                may_ident = e.target();
+            }
+            ast::Expr::MathIdent(e) => {
+                name = e.get().to_string();
+                break;
+            }
+            ast::Expr::Ident(e) => {
+                name = e.get().to_string();
+                break;
+            }
+            _ => return None,
+        }
+    }
+
+    let ident = node.find(may_ident.span())?;
+
+    // todo: if it is exported, find all the references in the workspace
+    let ident_ref = IdentRef {
+        name,
+        range: ident.range(),
+    };
+
+    let (id, _) = def_use.get_def(ident.span().id()?, &ident_ref)?;
+    Some(
+        def_use
+            .get_refs(id)
+            .map(|r| r.range.clone())
+            .collect::<Vec<_>>(),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::*;
+
+    #[test]
+    fn test() {
+        // goto_definition
+        snapshot_testing("references", &|world, path| {
+            let source = get_suitable_source_in_workspace(world, &path).unwrap();
+
+            let request = ReferencesRequest {
+                path: path.clone(),
+                position: find_test_position(&source),
+            };
+
+            let result = request.request(world, PositionEncoding::Utf16);
+            assert_snapshot!(JsonRepr::new_redacted(result, &REDACT_LOC));
+        });
+    }
+}