feat: add search index to quickly filter unrelated files (#864)

crates/tinymist-query/src/references.rs

@@ -1,10 +1,12 @@
+use std::sync::OnceLock;
+
 use log::debug;
 use typst::syntax::Span;
 
 use crate::{
     analysis::{Definition, SearchCtx},
     prelude::*,
-    syntax::{DerefTarget, RefExpr},
+    syntax::{get_index_info, DerefTarget, RefExpr},
     ty::Interned,
 };
 
@@ -59,6 +61,7 @@ pub(crate) fn find_references(
         ctx: ctx.fork_for_search(),
         references: vec![],
         def,
+        module_path: OnceLock::new(),
     };
 
     if finding_label {
@@ -73,6 +76,7 @@ struct ReferencesWorker<'a> {
     ctx: SearchCtx<'a>,
     references: Vec<LspLocation>,
     def: Definition,
+    module_path: OnceLock<Interned<str>>,
 }
 
 impl<'a> ReferencesWorker<'a> {
@@ -103,7 +107,25 @@ impl<'a> ReferencesWorker<'a> {
 
     fn file(&mut self, ref_fid: TypstFileId) -> Option<()> {
         log::debug!("references: file: {ref_fid:?}");
-        let ei = self.ctx.ctx.expr_stage_by_id(ref_fid)?;
+        let src = self.ctx.ctx.source_by_id(ref_fid).ok()?;
+        let index = get_index_info(&src);
+        match self.def.decl.kind() {
+            DefKind::Constant | DefKind::Function | DefKind::Struct | DefKind::Variable => {
+                if !index.identifiers.contains(self.def.decl.name()) {
+                    return Some(());
+                }
+            }
+            DefKind::Module => {
+                let ref_by_ident = index.identifiers.contains(self.def.decl.name());
+                let ref_by_path = index.paths.contains(self.module_path());
+                if !(ref_by_ident || ref_by_path) {
+                    return Some(());
+                }
+            }
+            DefKind::Reference => {}
+        }
+
+        let ei = self.ctx.ctx.expr_stage(&src);
         let uri = self.ctx.ctx.uri_for_id(ref_fid).ok()?;
 
         let t = ei.get_refs(self.def.decl.clone());
@@ -135,6 +157,23 @@ impl<'a> ReferencesWorker<'a> {
             })
         }));
     }
+
+    // todo: references of package
+    fn module_path(&self) -> &Interned<str> {
+        self.module_path.get_or_init(|| {
+            self.def
+                .decl
+                .file_id()
+                .and_then(|fid| {
+                    fid.vpath()
+                        .as_rooted_path()
+                        .file_name()?
+                        .to_str()
+                        .map(From::from)
+                })
+                .unwrap_or_default()
+        })
+    }
 }
 
 #[cfg(test)]

crates/tinymist-query/src/syntax/index.rs

@@ -0,0 +1,66 @@
+use std::str::FromStr;
+
+use reflexo_typst::package::PackageSpec;
+use rustc_hash::FxHashSet;
+
+use crate::{adt::interner::Interned, prelude::*};
+
+#[derive(Default)]
+pub struct IndexInfo {
+    pub(crate) paths: FxHashSet<Interned<str>>,
+    pub(crate) packages: FxHashSet<PackageSpec>,
+    pub(crate) identifiers: FxHashSet<Interned<str>>,
+}
+
+#[comemo::memoize]
+pub fn get_index_info(src: &Source) -> Arc<IndexInfo> {
+    let root = src.root();
+    let mut worker = IndexWorker {
+        info: IndexInfo::default(),
+    };
+    worker.visit(root);
+    Arc::new(worker.info)
+}
+
+struct IndexWorker {
+    info: IndexInfo,
+}
+
+impl IndexWorker {
+    fn visit(&mut self, node: &SyntaxNode) {
+        match node.cast::<ast::Expr>() {
+            Some(ast::Expr::Str(s)) => {
+                if s.to_untyped().text().len() > 65536 {
+                    // skip long strings
+                    return;
+                }
+                let s = s.get();
+
+                if s.starts_with('@') {
+                    let pkg_spec = PackageSpec::from_str(&s).ok();
+                    if let Some(pkg_spec) = pkg_spec {
+                        self.info.identifiers.insert(pkg_spec.name.clone().into());
+                        self.info.packages.insert(pkg_spec);
+                    }
+                    return;
+                }
+                let p = Path::new(s.as_str());
+                let name = p.file_name().unwrap_or_default().to_str();
+                if let Some(name) = name {
+                    self.info.paths.insert(name.into());
+                }
+            }
+            Some(ast::Expr::MathIdent(i)) => {
+                self.info.identifiers.insert(i.get().into());
+            }
+            Some(ast::Expr::Ident(i)) => {
+                self.info.identifiers.insert(i.get().into());
+            }
+            _ => {}
+        }
+
+        for child in node.children() {
+            self.visit(child);
+        }
+    }
+}

crates/tinymist-query/src/syntax/mod.rs

@@ -22,3 +22,5 @@ pub(crate) mod def;
 pub use def::*;
 pub(crate) mod repr;
 use repr::*;
+pub(crate) mod index;
+pub use index::*;