feat: consider context free requests

crates/tinymist-query/src/document_symbol.rs

@@ -11,11 +11,9 @@ pub struct DocumentSymbolRequest {
 impl DocumentSymbolRequest {
     pub fn request(
         self,
-        world: &TypstSystemWorld,
+        source: Source,
         position_encoding: PositionEncoding,
     ) -> Option<DocumentSymbolResponse> {
-        let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
-
         let symbols = get_lexical_hierarchy(source.clone(), LexicalScopeKind::Symbol)?;
 
         let symbols = filter_document_symbols(&symbols, &source, position_encoding);
@@ -74,8 +72,11 @@ mod tests {
 }
 "#,
             |world, path| {
-                let request = DocumentSymbolRequest { path };
-                let result = request.request(world, PositionEncoding::Utf16);
+                let request = DocumentSymbolRequest { path: path.clone() };
+
+                let source = get_suitable_source_in_workspace(world, &path).unwrap();
+
+                let result = request.request(source, PositionEncoding::Utf16);
                 assert_snapshot!(JsonRepr::new_redacted(result.unwrap(), &REDACT_LOC), @r###"
                 [
                  {

crates/tinymist-query/src/folding_range.rs

@@ -12,13 +12,11 @@ pub struct FoldingRangeRequest {
 impl FoldingRangeRequest {
     pub fn request(
         self,
-        world: &TypstSystemWorld,
+        source: Source,
         position_encoding: PositionEncoding,
     ) -> Option<Vec<FoldingRange>> {
         let line_folding_only = self.line_folding_only;
 
-        let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
-
         let symbols = get_lexical_hierarchy(source.clone(), LexicalScopeKind::Block)?;
 
         let mut results = vec![];
@@ -123,10 +121,13 @@ mod tests {
     fn test_folding_range_request() {
         run_with_source("#let a = 1;", |world, path| {
             let request = FoldingRangeRequest {
-                path,
+                path: path.clone(),
                 line_folding_only: true,
             };
-            let result = request.request(world, PositionEncoding::Utf16);
+
+            let source = get_suitable_source_in_workspace(world, &path).unwrap();
+
+            let result = request.request(source, PositionEncoding::Utf16);
             assert_snapshot!(JsonRepr::new_pure(result.unwrap()), @"[]");
         });
         let t = r#"#let a = {
@@ -136,10 +137,13 @@ mod tests {
 }"#;
         run_with_source(t, |world, path| {
             let request = FoldingRangeRequest {
-                path,
+                path: path.clone(),
                 line_folding_only: true,
             };
-            let result = request.request(world, PositionEncoding::Utf16);
+
+            let source = get_suitable_source_in_workspace(world, &path).unwrap();
+
+            let result = request.request(source, PositionEncoding::Utf16);
             assert_snapshot!(JsonRepr::new_pure(result.unwrap()), @r###"
             [
              {

crates/tinymist-query/src/lib.rs

@@ -42,6 +42,14 @@ mod polymorphic {
         pub path: PathBuf,
     }
 
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    pub enum FoldRequestFeature {
+        PinnedFirst,
+        Unique,
+        Mergable,
+        ContextFreeUnique,
+    }
+
     #[derive(Debug, Clone)]
     pub enum CompilerQueryRequest {
         OnSaveExport(OnSaveExportRequest),
@@ -57,7 +65,26 @@ mod polymorphic {
         FoldingRange(FoldingRangeRequest),
         SelectionRange(SelectionRangeRequest),
     }
+
     impl CompilerQueryRequest {
+        pub fn fold_feature(&self) -> FoldRequestFeature {
+            use FoldRequestFeature::*;
+            match self {
+                CompilerQueryRequest::OnSaveExport(..) => Mergable,
+                CompilerQueryRequest::Hover(..) => PinnedFirst,
+                CompilerQueryRequest::GotoDefinition(..) => PinnedFirst,
+                CompilerQueryRequest::InlayHint(..) => Unique,
+                CompilerQueryRequest::Completion(..) => Mergable,
+                CompilerQueryRequest::SignatureHelp(..) => PinnedFirst,
+                CompilerQueryRequest::DocumentSymbol(..) => ContextFreeUnique,
+                CompilerQueryRequest::Symbol(..) => Mergable,
+                CompilerQueryRequest::SemanticTokensFull(..) => ContextFreeUnique,
+                CompilerQueryRequest::SemanticTokensDelta(..) => ContextFreeUnique,
+                CompilerQueryRequest::FoldingRange(..) => ContextFreeUnique,
+                CompilerQueryRequest::SelectionRange(..) => ContextFreeUnique,
+            }
+        }
+
         pub fn associated_path(&self) -> Option<&Path> {
             Some(match self {
                 CompilerQueryRequest::OnSaveExport(req) => &req.path,

crates/tinymist-query/src/selection_range.rs

@@ -9,11 +9,9 @@ pub struct SelectionRangeRequest {
 impl SelectionRangeRequest {
     pub fn request(
         self,
-        world: &TypstSystemWorld,
+        source: Source,
         position_encoding: PositionEncoding,
     ) -> Option<Vec<SelectionRange>> {
-        let source = get_suitable_source_in_workspace(world, &self.path).ok()?;
-
         let mut ranges = Vec::new();
         for position in self.positions {
             let typst_offset =

crates/tinymist/src/actor/typst.rs

@@ -8,8 +8,8 @@ use anyhow::anyhow;
 use futures::future::join_all;
 use log::{debug, error, trace, warn};
 use tinymist_query::{
-    CompilerQueryRequest, CompilerQueryResponse, DiagnosticsMap, LspDiagnostic, LspRange,
-    OnSaveExportRequest, PositionEncoding, SemanticTokenCache,
+    CompilerQueryRequest, CompilerQueryResponse, DiagnosticsMap, FoldRequestFeature, LspDiagnostic,
+    LspRange, OnSaveExportRequest, PositionEncoding, SemanticTokenCache,
 };
 use tokio::sync::{broadcast, mpsc, watch, Mutex, RwLock};
 use tower_lsp::lsp_types::{TextDocumentContentChangeEvent, Url};
@@ -345,6 +345,19 @@ macro_rules! query_world {
     }};
 }
 
+macro_rules! query_source {
+    ($self:ident, $method:ident, $req:expr) => {{
+        let path: ImmutPath = $req.path.clone().into();
+        let vfs = $self.memory_changes.read().await;
+        let snapshot = vfs.get(&path).ok_or_else(|| anyhow!("file missing"))?;
+        let source = snapshot.content.clone();
+
+        let enc = $self.position_encoding;
+        let res = $req.request(source, enc);
+        Ok(CompilerQueryResponse::$method(res))
+    }};
+}
+
 macro_rules! query_tokens_cache {
     ($self:ident, $method:ident, $req:expr) => {{
         let path: ImmutPath = $req.path.clone().into();
@@ -368,11 +381,24 @@ impl CompileCluster {
         match query {
             SemanticTokensFull(req) => query_tokens_cache!(self, SemanticTokensFull, req),
             SemanticTokensDelta(req) => query_tokens_cache!(self, SemanticTokensDelta, req),
+            FoldingRange(req) => query_source!(self, FoldingRange, req),
+            SelectionRange(req) => query_source!(self, SelectionRange, req),
+            DocumentSymbol(req) => query_source!(self, DocumentSymbol, req),
             _ => {
-                if let Some(path) = query.associated_path() {
-                    self.primary.change_entry(path.into()).await?;
-                }
-                self.primary.query(query).await
+                let main = self.main.lock().await;
+
+                let query_target = match main.as_ref() {
+                    Some(main) => main,
+                    None => {
+                        // todo: race condition, we need atomic primary query
+                        if let Some(path) = query.associated_path() {
+                            self.primary.change_entry(path.into()).await?;
+                        }
+                        &self.primary
+                    }
+                };
+
+                query_target.query(query).await
             }
         }
     }
@@ -725,6 +751,7 @@ impl<H: CompilationHandle> CompileNode<H> {
         query: CompilerQueryRequest,
     ) -> anyhow::Result<CompilerQueryResponse> {
         use CompilerQueryRequest::*;
+        assert!(query.fold_feature() != FoldRequestFeature::ContextFreeUnique);
 
         match query {
             CompilerQueryRequest::OnSaveExport(OnSaveExportRequest { path }) => {
@@ -736,12 +763,12 @@ impl<H: CompilationHandle> CompileNode<H> {
             InlayHint(req) => query_world!(self, InlayHint, req),
             Completion(req) => query_state!(self, Completion, req),
             SignatureHelp(req) => query_world!(self, SignatureHelp, req),
-            DocumentSymbol(req) => query_world!(self, DocumentSymbol, req),
             Symbol(req) => query_world!(self, Symbol, req),
-            FoldingRange(req) => query_world!(self, FoldingRange, req),
-            SelectionRange(req) => query_world!(self, SelectionRange, req),
-            CompilerQueryRequest::SemanticTokensDelta(..)
-            | CompilerQueryRequest::SemanticTokensFull(..) => unreachable!(),
+            FoldingRange(..)
+            | SelectionRange(..)
+            | SemanticTokensDelta(..)
+            | DocumentSymbol(..)
+            | SemanticTokensFull(..) => unreachable!(),
         }
     }