tinymist_query/index/
mod.rs

1//! Dumps typst knowledge from workspace.
2//!
3//! Reference Impls:
4//! - <https://github.com/sourcegraph/lsif-jsonnet/blob/e186f9fde623efa8735261e9cb059ad3a58b535f/dumper/dumper.go>
5//! - <https://github.com/rust-lang/rust-analyzer/blob/5c0b555a65cadc14a6a16865c3e065c9d30b0bef/crates/ide/src/static_index.rs>
6//! - <https://github.com/rust-lang/rust-analyzer/blob/5c0b555a65cadc14a6a16865c3e065c9d30b0bef/crates/rust-analyzer/src/cli/lsif.rs>
7
8use core::fmt;
9use std::sync::Arc;
10
11use crate::analysis::{SemanticTokens, SharedContext};
12use crate::index::protocol::ResultSet;
13use crate::prelude::Definition;
14use crate::{LocalContext, path_to_url};
15use ecow::EcoString;
16use lsp_types::{Hover, Url};
17use tinymist_analysis::docs::DefDocs;
18use tinymist_analysis::syntax::classify_syntax;
19use tinymist_std::error::WithContextUntyped;
20use tinymist_std::hash::FxHashMap;
21use tinymist_world::EntryReader;
22use typst::syntax::{FileId, LinkedNode, Source, Span};
23use typst_shim::syntax::source_range;
24
25pub mod protocol;
26pub mod query;
27mod scip_utils;
28use protocol as p;
29
30/// The dumped knowledge.
31pub struct Knowledge {
32    /// The meta data.
33    pub meta: p::MetaData,
34    /// The files.
35    pub files: Vec<FileIndex>,
36}
37
38impl Knowledge {
39    /// Creates a new empty knowledge.
40    pub fn bind<'a>(&'a self, ctx: &'a Arc<SharedContext>) -> KnowledgeWithContext<'a> {
41        KnowledgeWithContext {
42            knowledge: self,
43            ctx,
44        }
45    }
46}
47
48/// A view of knowledge with context for dumping.
49pub struct KnowledgeWithContext<'a> {
50    knowledge: &'a Knowledge,
51    ctx: &'a Arc<SharedContext>,
52}
53
54pub mod scip;
55pub mod scip_query;
56pub use scip::{ScipPublicApi, ScipPublicModule};
57
58impl fmt::Display for KnowledgeWithContext<'_> {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        let mut files = FxHashMap::default();
61        let mut encoder = LsifEncoder {
62            ctx: self.ctx,
63            writer: f,
64            id: IdCounter::new(),
65            files: &mut files,
66            results: FxHashMap::default(),
67            hover_results: FxHashMap::default(),
68        };
69
70        encoder
71            .emit_meta(self.knowledge.meta.clone())
72            .map_err(|err| {
73                log::error!("cannot write meta data: {err}");
74                fmt::Error
75            })?;
76        encoder.emit_files(&self.knowledge.files).map_err(|err| {
77            log::error!("cannot write files: {err}");
78            fmt::Error
79        })
80    }
81}
82
83struct IdCounter {
84    next: usize,
85}
86
87impl IdCounter {
88    fn new() -> Self {
89        Self { next: 0 }
90    }
91
92    fn next(&mut self) -> usize {
93        let id = self.next;
94        self.next += 1;
95        id
96    }
97}
98
99trait LsifWrite {
100    fn write_element(&mut self, id: i32, element: p::Element) -> fmt::Result;
101}
102
103impl<T: fmt::Write> LsifWrite for T {
104    fn write_element(&mut self, id: i32, element: p::Element) -> fmt::Result {
105        let entry = p::Entry { id, data: element };
106        self.write_str(&serde_json::to_string(&entry).unwrap())?;
107        self.write_char('\n')
108    }
109}
110
111struct LsifEncoder<'a, W: fmt::Write> {
112    ctx: &'a Arc<SharedContext>,
113    writer: &'a mut W,
114    id: IdCounter,
115    files: &'a mut FxHashMap<FileId, i32>,
116    results: FxHashMap<Span, i32>,
117    hover_results: FxHashMap<String, i32>,
118}
119
120impl<'a, W: fmt::Write> LsifEncoder<'a, W> {
121    fn alloc_file_id(&mut self, fid: FileId) -> i32 {
122        *self.files.entry(fid).or_insert_with(|| {
123            let id = self.id.next() as i32;
124            self.writer
125                .write_element(
126                    id,
127                    p::Element::Vertex(p::Vertex::Document(p::Document {
128                        uri: self.ctx.uri_for_id(fid).unwrap_or_else(|err| {
129                            log::error!("cannot get uri for {fid:?}: {err}");
130                            Url::parse("file:///unknown").unwrap()
131                        }),
132                        language_id: EcoString::inline("typst"),
133                    })),
134                )
135                .unwrap();
136
137            id
138        })
139    }
140
141    fn alloc_result_id(&mut self, span: Span) -> tinymist_std::Result<i32> {
142        if let Some(id) = self.results.get(&span) {
143            return Ok(*id);
144        }
145
146        let id = self.emit_element(p::Element::Vertex(p::Vertex::ResultSet(ResultSet {
147            key: None,
148        })))?;
149
150        self.results.insert(span, id);
151        Ok(id)
152    }
153
154    fn emit_element(&mut self, element: p::Element) -> tinymist_std::Result<i32> {
155        let id = self.id.next() as i32;
156        self.writer
157            .write_element(id, element)
158            .context_ut("cannot write element")?;
159        Ok(id)
160    }
161
162    fn emit_meta(&mut self, meta: p::MetaData) -> tinymist_std::Result<()> {
163        let obj = p::Element::Vertex(p::Vertex::MetaData(meta));
164        self.emit_element(obj).map(|_| ())
165    }
166
167    fn emit_files(&mut self, files: &[FileIndex]) -> tinymist_std::Result<()> {
168        for (idx, file) in files.iter().enumerate() {
169            eprintln!("emit file: {:?}, {idx} of {}", file.fid, files.len());
170            let source = self
171                .ctx
172                .source_by_id(file.fid)
173                .context_ut("cannot get source")?;
174            let fid = self.alloc_file_id(file.fid);
175            let semantic_tokens_id =
176                self.emit_element(p::Element::Vertex(p::Vertex::SemanticTokensResult {
177                    result: lsp_types::SemanticTokens {
178                        result_id: None,
179                        data: file.semantic_tokens.as_ref().clone(),
180                    },
181                }))?;
182            self.emit_element(p::Element::Edge(p::Edge::SemanticTokens(p::EdgeData {
183                out_v: fid,
184                in_v: semantic_tokens_id,
185            })))?;
186
187            let mut tokens_id = vec![];
188            for (s, reference) in &file.references {
189                let def = &reference.definition;
190                if let Some(range_id) = self.emit_span(*s, &source) {
191                    let res_id = self.alloc_result_id(*s)?;
192                    self.emit_element(p::Element::Edge(p::Edge::Next(p::EdgeData {
193                        out_v: range_id,
194                        in_v: res_id,
195                    })))?;
196                    if let Some(hover) = &reference.hover {
197                        self.emit_hover(res_id, hover)?;
198                    }
199                    tokens_id.push(range_id);
200                }
201
202                if let Some(def_range_id) = self.emit_def_span(def, &source, false) {
203                    if let Some(hover) = &reference.hover {
204                        self.emit_hover(def_range_id, hover)?;
205                    }
206                    tokens_id.push(def_range_id);
207                }
208            }
209            self.emit_element(p::Element::Edge(p::Edge::Contains(p::EdgeDataMultiIn {
210                out_v: fid,
211                in_vs: tokens_id,
212            })))?;
213
214            for (s, reference) in &file.references {
215                let def = &reference.definition;
216                let res_id = self.alloc_result_id(*s)?;
217                let def_id = self.emit_element(p::Element::Vertex(p::Vertex::DefinitionResult))?;
218                let Some(def_range) = self.emit_def_span(def, &source, true) else {
219                    continue;
220                };
221                let Some(file_id) = def.file_id() else {
222                    continue;
223                };
224                let file_vertex_id = self.alloc_file_id(file_id);
225
226                self.emit_element(p::Element::Edge(p::Edge::Item(p::Item {
227                    document: file_vertex_id,
228                    property: None,
229                    edge_data: p::EdgeDataMultiIn {
230                        in_vs: vec![def_range],
231                        out_v: def_id,
232                    },
233                })))?;
234                self.emit_element(p::Element::Edge(p::Edge::Definition(p::EdgeData {
235                    in_v: def_id,
236                    out_v: res_id,
237                })))?;
238            }
239        }
240        Ok(())
241    }
242
243    fn emit_hover(&mut self, out_v: i32, hover: &Hover) -> tinymist_std::Result<()> {
244        let hover_id = self.alloc_hover_id(hover)?;
245        self.emit_element(p::Element::Edge(p::Edge::Hover(p::EdgeData {
246            out_v,
247            in_v: hover_id,
248        })))?;
249        Ok(())
250    }
251
252    fn alloc_hover_id(&mut self, hover: &Hover) -> tinymist_std::Result<i32> {
253        let key = serde_json::to_string(hover).context_ut("cannot serialize hover")?;
254        if let Some(id) = self.hover_results.get(&key) {
255            return Ok(*id);
256        }
257
258        let id = self.emit_element(p::Element::Vertex(p::Vertex::HoverResult {
259            result: hover.clone(),
260        }))?;
261        self.hover_results.insert(key, id);
262        Ok(id)
263    }
264
265    fn emit_span(&mut self, span: Span, source: &Source) -> Option<i32> {
266        let range = source_range(source, span)?;
267        self.emit_element(p::Element::Vertex(p::Vertex::Range {
268            range: self.ctx.to_lsp_range(range, source),
269            tag: None,
270        }))
271        .ok()
272    }
273
274    fn emit_def_span(&mut self, def: &Definition, source: &Source, external: bool) -> Option<i32> {
275        let s = def.decl.span();
276        if !s.is_detached() && s.id() == Some(source.id()) {
277            self.emit_span(s, source)
278        } else if let Some(fid) = def.file_id()
279            && fid == source.id()
280        {
281            // todo: module it self
282            None
283        } else if external && !s.is_detached() {
284            let external_src = self.ctx.source_by_id(def.file_id()?).ok()?;
285            self.emit_span(s, &external_src)
286        } else {
287            None
288        }
289    }
290}
291
292/// The index of a file.
293pub struct FileIndex {
294    /// The file id.
295    pub fid: FileId,
296    /// The semantic tokens of the file.
297    pub semantic_tokens: SemanticTokens,
298    /// The documentation of the file.
299    pub documentation: Option<EcoString>,
300    /// The references in the file.
301    pub references: FxHashMap<Span, ReferenceIndex>,
302}
303
304/// The index of a reference.
305pub struct ReferenceIndex {
306    /// The definition referenced by the source span.
307    pub definition: Definition,
308    /// The static hover result for the referenced definition.
309    pub hover: Option<Hover>,
310    /// The typed documentation for the referenced definition.
311    pub def_docs: Option<DefDocs>,
312}
313
314/// Dumps typst knowledge in [LSIF] format from workspace.
315///
316/// [LSIF]: https://microsoft.github.io/language-server-protocol/specifications/lsif/0.6.0/specification/
317pub fn knowledge(ctx: &mut LocalContext) -> tinymist_std::Result<Knowledge> {
318    let root = ctx
319        .world()
320        .entry_state()
321        .workspace_root()
322        .ok_or_else(|| tinymist_std::error_once!("workspace root is not set"))?;
323
324    let mut files = ctx.source_files().clone();
325    if let Some(main) = ctx.world().entry_state().main()
326        && !files.contains(&main)
327    {
328        files.push(main);
329    }
330    for fid in ctx.depended_source_files() {
331        if !files.contains(&fid) {
332            files.push(fid);
333        }
334    }
335
336    let mut worker = DumpWorker {
337        ctx,
338        strings: FxHashMap::default(),
339        references: FxHashMap::default(),
340        hovers: FxHashMap::default(),
341        def_docs: FxHashMap::default(),
342    };
343    let files = files
344        .iter()
345        .map(move |fid| worker.file(fid))
346        .collect::<tinymist_std::Result<Vec<FileIndex>>>()?;
347
348    Ok(Knowledge {
349        meta: p::MetaData {
350            version: "0.6.0".to_string(),
351            project_root: path_to_url(&root)?,
352            position_encoding: p::Encoding::Utf16,
353            tool_info: Some(p::ToolInfo {
354                name: "tinymist".to_string(),
355                args: vec![],
356                version: Some(env!("CARGO_PKG_VERSION").to_string()),
357            }),
358        },
359        files,
360    })
361}
362
363struct DumpWorker<'a> {
364    /// The context.
365    ctx: &'a mut LocalContext,
366    /// A string interner.
367    strings: FxHashMap<EcoString, EcoString>,
368    /// The references collected so far.
369    references: FxHashMap<Span, ReferenceIndex>,
370    /// The static hover results collected so far.
371    hovers: FxHashMap<Definition, Option<Hover>>,
372    /// The typed documentation collected so far.
373    def_docs: FxHashMap<Definition, Option<DefDocs>>,
374}
375
376impl DumpWorker<'_> {
377    fn file(&mut self, fid: &FileId) -> tinymist_std::Result<FileIndex> {
378        let source = self.ctx.source_by_id(*fid).context_ut("cannot parse")?;
379        let semantic_tokens = crate::SemanticTokensFullRequest::compute(self.ctx, &source);
380
381        let root = LinkedNode::new(source.root());
382        self.walk(&source, &root);
383        let references = std::mem::take(&mut self.references);
384
385        Ok(FileIndex {
386            fid: *fid,
387            semantic_tokens,
388            documentation: Some(self.intern("File documentation.")), // todo
389            references,
390        })
391    }
392
393    fn intern(&mut self, s: &str) -> EcoString {
394        if let Some(v) = self.strings.get(s) {
395            return v.clone();
396        }
397        let v = EcoString::from(s);
398        self.strings.insert(v.clone(), v.clone());
399        v
400    }
401
402    fn walk(&mut self, source: &Source, node: &LinkedNode) {
403        if node.get().children().len() == 0 {
404            let Some(syntax) = classify_syntax(node.clone(), node.offset()) else {
405                return;
406            };
407            let span = syntax.node().span();
408            if self.references.contains_key(&span) {
409                return;
410            }
411
412            let Some(definition) = self.ctx.def_of_syntax(source, syntax) else {
413                return;
414            };
415            let hover = self.hover(&definition);
416            let def_docs = self.def_docs(&definition);
417            self.references.insert(
418                span,
419                ReferenceIndex {
420                    definition,
421                    hover,
422                    def_docs,
423                },
424            );
425
426            return;
427        }
428
429        for child in node.children() {
430            self.walk(source, &child);
431        }
432    }
433
434    fn hover(&mut self, definition: &Definition) -> Option<Hover> {
435        if let Some(hover) = self.hovers.get(definition) {
436            return hover.clone();
437        }
438
439        let hover = crate::hover::hover_from_definition_shared(self.ctx.shared(), definition, None);
440        self.hovers.insert(definition.clone(), hover.clone());
441        hover
442    }
443
444    fn def_docs(&mut self, definition: &Definition) -> Option<DefDocs> {
445        if let Some(docs) = self.def_docs.get(definition) {
446            return docs.clone();
447        }
448
449        let docs = self.ctx.def_docs(definition);
450        self.def_docs.insert(definition.clone(), docs.clone());
451        docs
452    }
453}
tinymist_query/index/mod.rs

tinymist_query/index/
mod.rs
