feat: find syntax definition

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

@@ -1,15 +1,29 @@
-use log::trace;
+use std::ops::Deref;
+use std::path::Path;
+
+use comemo::Tracked;
+use log::{debug, trace};
+use typst::syntax::ast::Ident;
+use typst::syntax::VirtualPath;
+use typst::World;
 use typst::{
     foundations::{Func, Value},
     syntax::{
         ast::{self, AstNode},
-        LinkedNode,
+        LinkedNode, Source, Span, SyntaxKind,
     },
 };
-use typst_ts_compiler::TypstSystemWorld;
+use typst_ts_core::TypstFileId;
 
 use crate::{prelude::analyze_expr, TypstSpan};
 
+#[derive(Debug, Clone)]
+pub struct VariableDefinition<'a> {
+    pub def_site: LinkedNode<'a>,
+    pub use_site: LinkedNode<'a>,
+    pub span: TypstSpan,
+}
+
 #[derive(Debug, Clone)]
 pub struct FuncDefinition<'a> {
     pub value: Func,
@@ -17,20 +31,352 @@ pub struct FuncDefinition<'a> {
     pub span: TypstSpan,
 }
 
+#[derive(Debug, Clone)]
+pub struct ModuleDefinition<'a> {
+    pub module: TypstFileId,
+    pub use_site: LinkedNode<'a>,
+    pub span: TypstSpan,
+}
+
+#[derive(Debug, Clone)]
+pub struct ExternalDefinition<'a> {
+    pub use_site: LinkedNode<'a>,
+    pub span: TypstSpan,
+}
+
 #[derive(Debug, Clone)]
 pub enum Definition<'a> {
     Func(FuncDefinition<'a>),
+    Var(VariableDefinition<'a>),
+    Module(ModuleDefinition<'a>),
+    External(ExternalDefinition<'a>),
+}
+
+impl Definition<'_> {
+    pub fn span(&self) -> TypstSpan {
+        match self {
+            Definition::Func(f) => f.span,
+            Definition::Var(v) => v.span,
+            Definition::Module(m) => m.span,
+            Definition::External(s) => s.span,
+        }
+    }
+
+    pub fn use_site(&self) -> &LinkedNode {
+        match self {
+            Definition::Func(f) => &f.use_site,
+            Definition::Var(v) => &v.use_site,
+            Definition::Module(m) => &m.use_site,
+            Definition::External(s) => &s.use_site,
+        }
+    }
+}
+
+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> {
+    match node.prev_sibling() {
+        Some(prev) => Some(prev),
+        None => {
+            let parent = node.parent()?;
+            debug!("no prev sibling parent: {parent:?}");
+            advance_prev_adjacent(parent.clone())
+        }
+    }
+}
+
+fn find_source_by_import(
+    world: Tracked<'_, dyn World>,
+    current: TypstFileId,
+    import_node: ast::ModuleImport,
+) -> Option<Source> {
+    let v = import_node.source();
+    match v {
+        ast::Expr::Str(s) => {
+            let s = s.get();
+
+            if s.starts_with('@') {
+                // todo: import from package
+                return None;
+            }
+
+            let path = Path::new(s.as_str());
+            let vpath = if path.is_relative() {
+                current.vpath().join(path)
+            } else {
+                VirtualPath::new(path)
+            };
+
+            let id = TypstFileId::new(current.package().cloned(), vpath);
+            world.source(id).ok()
+        }
+        _ => None,
+    }
+}
+
+#[comemo::memoize]
+fn find_definition_in_module<'a>(
+    world: Tracked<'a, dyn World>,
+    current: TypstFileId,
+    source: Source,
+    name: &'a str,
+) -> Option<Span> {
+    // todo: cyclic import
+    let root = source.root();
+    let node = LinkedNode::new(root);
+    let last_expr = if let Some(m) = root.cast::<ast::Markup>() {
+        m.exprs().last()?
+    } else {
+        debug!("unexpected root kind {:?}", root.kind());
+        return None;
+    };
+    let last = node.find(last_expr.span())?;
+    let e = find_syntax_definition(world, current, last, name)?;
+    Some(e.span())
+}
+
+enum ImportRef<'a> {
+    /// `import "foo" as bar;`
+    ///                  ^^^
+    ModuleAs(Ident<'a>),
+    /// `import "foo.typ"`
+    ///          ^^^
+    Path(ast::Expr<'a>),
+    /// `import "foo": bar`
+    ///                ^^^
+    Ident(Ident<'a>),
+    /// `import "foo": bar as baz`
+    ///                       ^^^
+    IdentAs(ast::RenamedImportItem<'a>),
+    /// `import "foo": *`
+    ExternalResolved(Span),
+}
+
+fn find_ref_in_import<'a>(
+    world: Tracked<'_, dyn World>,
+    current: TypstFileId,
+    import_node: ast::ModuleImport<'a>,
+    name: &str,
+) -> Option<ImportRef<'a>> {
+    if let Some(import_node) = import_node.new_name() {
+        if import_node.get() == name {
+            return Some(ImportRef::ModuleAs(import_node));
+        }
+    }
+
+    let Some(imports) = import_node.imports() else {
+        let v = import_node.source();
+        match v {
+            ast::Expr::Str(e) => {
+                let e = e.get();
+                let e = Path::new(e.as_ref());
+                let Some(e) = e.file_name() else {
+                    return None;
+                };
+                let e = e.to_string_lossy();
+                let e = e.as_ref();
+                let Some(e) = e.strip_suffix(".typ") else {
+                    return None;
+                };
+                return (e == name).then_some(ImportRef::Path(v));
+            }
+            _ => return None,
+        }
+    };
+
+    match imports {
+        ast::Imports::Wildcard => {
+            let dep = find_source_by_import(world, current, import_node)?;
+            let res = find_definition_in_module(world, current, dep, name)?;
+            return Some(ImportRef::ExternalResolved(res));
+        }
+        ast::Imports::Items(items) => {
+            for handle in items.iter() {
+                match handle {
+                    ast::ImportItem::Simple(e) => {
+                        if e.get() == name {
+                            return Some(ImportRef::Ident(e));
+                        }
+                    }
+                    ast::ImportItem::Renamed(e) => {
+                        let o = e.new_name();
+                        if o.get() == name {
+                            return Some(ImportRef::IdentAs(e));
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    None
+}
+
+fn find_syntax_definition<'a>(
+    world: Tracked<'a, dyn World>,
+    current: TypstFileId,
+    node: LinkedNode<'a>,
+    name: &str,
+) -> Option<Definition<'a>> {
+    struct SyntaxDefinitionWorker<'a, 'b> {
+        world: Tracked<'a, dyn World>,
+        current: TypstFileId,
+        name: &'b str,
+        use_site: LinkedNode<'a>,
+    }
+
+    impl<'a, 'b> SyntaxDefinitionWorker<'a, 'b> {
+        fn find(&self, mut node: LinkedNode<'a>) -> Option<Definition<'a>> {
+            loop {
+                if let Some(def) = self.check(node.clone()) {
+                    return Some(def);
+                }
+
+                let Some(prev) = advance_prev_adjacent(node) else {
+                    debug!("no prev sibling parent");
+                    return None;
+                };
+
+                node = prev;
+            }
+        }
+
+        fn resolve_as_var(&self, node: LinkedNode<'a>, name: ast::Ident) -> Option<Definition<'a>> {
+            if name.get() != self.name {
+                return None;
+            }
+
+            let def_site = node.find(name.span())?;
+            Some(Definition::Var(VariableDefinition {
+                def_site,
+                use_site: self.use_site.clone(),
+                span: node.span(),
+            }))
+        }
+
+        fn check(&self, node: LinkedNode<'a>) -> Option<Definition<'a>> {
+            let node = deref_lvalue(node)?;
+            match node.kind() {
+                SyntaxKind::LetBinding => {
+                    let binding = node.cast::<ast::LetBinding>()?;
+                    match binding.kind() {
+                        ast::LetBindingKind::Closure(name) => {
+                            if name.get() == self.name {
+                                let values =
+                                    analyze_expr(self.world.deref(), &node.find(name.span())?);
+                                let func = values.into_iter().find_map(|v| match v {
+                                    Value::Func(f) => Some(f),
+                                    _ => None,
+                                });
+                                let Some(func) = func else {
+                                    debug!("no func found... {name:?}");
+                                    return None;
+                                };
+
+                                return Some(Definition::Func(FuncDefinition {
+                                    value: func,
+                                    use_site: self.use_site.clone(),
+                                    span: node.span(),
+                                }));
+                            }
+                            None
+                        }
+                        ast::LetBindingKind::Normal(ast::Pattern::Normal(ast::Expr::Ident(
+                            name,
+                        ))) => {
+                            return self.resolve_as_var(node.clone(), name);
+                        }
+                        ast::LetBindingKind::Normal(ast::Pattern::Normal(e)) => {
+                            let e = deref_lvalue(node.find(e.span())?)?;
+                            if let Some(name) = e.cast::<ast::Ident>() {
+                                return self.resolve_as_var(e.clone(), name);
+                            }
+                            None
+                        }
+                        ast::LetBindingKind::Normal(ast::Pattern::Destructuring(n)) => {
+                            for i in n.idents() {
+                                if i.get() == self.name {
+                                    return self.resolve_as_var(node.clone(), i);
+                                }
+                            }
+                            None
+                        }
+                        ast::LetBindingKind::Normal(ast::Pattern::Placeholder(..)) => None,
+                    }
+                }
+                SyntaxKind::ModuleImport => {
+                    let import_node = node.cast::<ast::ModuleImport>()?;
+
+                    match find_ref_in_import(self.world, self.current, import_node, self.name)? {
+                        ImportRef::ModuleAs(ident) => {
+                            let m = find_source_by_import(self.world, self.current, import_node)?;
+                            return Some(Definition::Module(ModuleDefinition {
+                                module: m.id(),
+                                use_site: self.use_site.clone(),
+                                span: ident.span(),
+                            }));
+                        }
+                        ImportRef::Path(s) => {
+                            let m = find_source_by_import(self.world, self.current, import_node)?;
+                            return Some(Definition::Module(ModuleDefinition {
+                                module: m.id(),
+                                use_site: self.use_site.clone(),
+                                span: s.span(),
+                            }));
+                        }
+                        ImportRef::Ident(ident) => {
+                            return Some(Definition::Var(VariableDefinition {
+                                def_site: node.find(ident.span())?,
+                                use_site: self.use_site.clone(),
+                                span: ident.span(),
+                            }));
+                        }
+                        ImportRef::IdentAs(item) => {
+                            let ident = item.new_name();
+                            return Some(Definition::Var(VariableDefinition {
+                                def_site: node.find(ident.span())?,
+                                use_site: self.use_site.clone(),
+                                span: ident.span(),
+                            }));
+                        }
+                        ImportRef::ExternalResolved(def_span) => {
+                            return Some(Definition::External(ExternalDefinition {
+                                use_site: self.use_site.clone(),
+                                span: def_span,
+                            }));
+                        }
+                    }
+                }
+                _ => None,
+            }
+        }
+    }
+
+    let worker = SyntaxDefinitionWorker {
+        world,
+        current,
+        name,
+        use_site: node.clone(),
+    };
+    worker.find(node)
 }
 
 // todo: field definition
 pub(crate) fn find_definition<'a>(
-    world: &TypstSystemWorld,
+    world: Tracked<'a, dyn World>,
+    current: TypstFileId,
     node: LinkedNode<'a>,
 ) -> Option<Definition<'a>> {
-    let mut ancestor = &node;
+    let mut ancestor = node;
     while !ancestor.is::<ast::Expr>() {
-        ancestor = ancestor.parent()?;
+        ancestor = ancestor.parent()?.clone();
     }
+    let ancestor = deref_lvalue(ancestor)?;
 
     let may_ident = ancestor.cast::<ast::Expr>()?;
     if !may_ident.hash() && !matches!(may_ident, ast::Expr::MathIdent(_)) {
@@ -39,7 +385,7 @@ pub(crate) fn find_definition<'a>(
 
     let mut is_ident_only = false;
     trace!("got ast_node kind {kind:?}", kind = ancestor.kind());
-    let callee_node = match may_ident {
+    let ref_node = match may_ident {
         // todo: label, reference
         // todo: import
         // todo: include
@@ -49,23 +395,21 @@ pub(crate) fn find_definition<'a>(
             is_ident_only = true;
             may_ident
         }
-        _ => return None,
+        _ => {
+            debug!("unsupported kind {kind:?}", kind = ancestor.kind());
+            return None;
+        }
     };
-    trace!("got callee_node {callee_node:?} {is_ident_only:?}");
 
     let use_site = if is_ident_only {
         ancestor.clone()
     } else {
-        ancestor.find(callee_node.span())?
+        ancestor.find(ref_node.span())?
     };
 
-    let values = analyze_expr(world, &use_site);
+    let values = analyze_expr(world.deref(), &use_site);
 
     let func_or_module = values.into_iter().find_map(|v| match &v {
-        Value::Args(a) => {
-            trace!("got args {a:?}");
-            None
-        }
         Value::Func(..) | Value::Module(..) => Some(v),
         _ => None,
     });
@@ -82,8 +426,20 @@ pub(crate) fn find_definition<'a>(
             return None;
         }
         _ => {
-            trace!("find value by lexical result. {use_site:?}");
-            return None;
+            return match may_ident {
+                ast::Expr::Ident(e) => find_syntax_definition(world, current, use_site, e.get()),
+                ast::Expr::MathIdent(e) => {
+                    find_syntax_definition(world, current, use_site, e.get())
+                }
+                ast::Expr::FieldAccess(..) => {
+                    debug!("find field access");
+                    None
+                }
+                _ => {
+                    debug!("unsupported kind {kind:?}", kind = ancestor.kind());
+                    None
+                }
+            }
         }
     })
 }

crates/tinymist-query/src/goto_definition.rs

@@ -1,10 +1,8 @@
+use comemo::Track;
 use log::debug;
 use tower_lsp::lsp_types::LocationLink;
 
-use crate::{
-    analysis::{find_definition, Definition},
-    prelude::*,
-};
+use crate::{analysis::find_definition, prelude::*};
 
 #[derive(Debug, Clone)]
 pub struct GotoDefinitionRequest {
@@ -21,12 +19,14 @@ impl GotoDefinitionRequest {
         let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
         let typst_offset = lsp_to_typst::position(self.position, position_encoding, &source)?;
 
-        let ast_node = LinkedNode::new(source.root()).leaf_at(typst_offset)?;
+        let ast_node = LinkedNode::new(source.root()).leaf_at(typst_offset + 1)?;
 
-        let Definition::Func(func) = find_definition(world, ast_node)?;
+        let t: &dyn World = world;
 
-        let span = func.span;
-        let callee_link = func.use_site;
+        let def = find_definition(t.track(), source.id(), ast_node)?;
+        // todo: handle other definitions
+        let span = def.span();
+        let use_site = def.use_site();
 
         if span.is_detached() {
             return None;
@@ -37,7 +37,7 @@ impl GotoDefinitionRequest {
         };
 
         let origin_selection_range =
-            typst_to_lsp::range(callee_link.range(), &source, position_encoding);
+            typst_to_lsp::range(use_site.range(), &source, position_encoding);
 
         let span_path = world.path_for_id(id).ok()?;
         let span_source = world.source(id).ok()?;

crates/tinymist-query/src/prelude.rs

@@ -5,6 +5,7 @@ pub use std::{
     sync::Arc,
 };
 
+pub use comemo::Track;
 pub use itertools::{Format, Itertools};
 pub use log::{error, trace};
 pub use tower_lsp::lsp_types::{

crates/tinymist-query/src/prepare_rename.rs

@@ -1,9 +1,8 @@
-use log::debug;
-
 use crate::{
     analysis::{find_definition, Definition},
     prelude::*,
 };
+use log::debug;
 
 #[derive(Debug, Clone)]
 pub struct PrepareRenameRequest {
@@ -25,9 +24,13 @@ impl PrepareRenameRequest {
         let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
         let typst_offset = lsp_to_typst::position(self.position, position_encoding, &source)?;
 
-        let ast_node = LinkedNode::new(source.root()).leaf_at(typst_offset)?;
+        let ast_node = LinkedNode::new(source.root()).leaf_at(typst_offset + 1)?;
 
-        let Definition::Func(func) = find_definition(world, ast_node)?;
+        let t: &dyn World = world;
+        let Definition::Func(func) = find_definition(t.track(), source.id(), ast_node)? else {
+            // todo: handle other definitions
+            return None;
+        };
 
         use typst::foundations::func::Repr;
         let mut f = func.value.clone();

crates/tinymist-query/src/rename.rs

@@ -24,9 +24,13 @@ impl RenameRequest {
         let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
         let typst_offset = lsp_to_typst::position(self.position, position_encoding, &source)?;
 
-        let ast_node = LinkedNode::new(source.root()).leaf_at(typst_offset)?;
+        let ast_node = LinkedNode::new(source.root()).leaf_at(typst_offset + 1)?;
 
-        let Definition::Func(func) = find_definition(world, ast_node)?;
+        let t: &dyn World = world;
+        let Definition::Func(func) = find_definition(t.track(), source.id(), ast_node)? else {
+            // todo: handle other definitions
+            return None;
+        };
 
         // todo: unwrap parentheses
 
@@ -211,6 +215,7 @@ fn search_in_workspace(
                 let Some(import_node) = stack_store.cast::<ast::ModuleImport>() else {
                     continue;
                 };
+                // todo: don't ignore import node
                 if import_node.new_name().is_some() {
                     continue;
                 }