tinymist_query/index/
query.rs

1//! Offline query the knowledge from the index.
2//!
3//! Types of Language Server Index Format (LSIF). LSIF is a standard format
4//! for language servers or other programming tools to dump their knowledge
5//! about a workspace.
6//!
7//! Based on <https://microsoft.github.io/language-server-protocol/specifications/lsif/0.6.0/specification/>
8
9use std::io::{BufRead, Read};
10
11use tinymist_std::error::prelude::*;
12use tinymist_std::hash::{FxHashMap, FxHashSet};
13
14use crate::{
15    CompilerQueryRequest, CompilerQueryResponse, GotoDefinitionRequest, HoverRequest, path_to_url,
16};
17
18use super::protocol::*;
19use lsp_types::{GotoDefinitionResponse, Hover, LocationLink, Position, Range, Url};
20
21/// The context for querying an LSIF JSONL index.
22#[derive(Default)]
23pub struct IndexQueryCtx {
24    meta: Option<MetaData>,
25    document_by_uri: FxHashMap<Url, Id>,
26    document_ranges: FxHashMap<Id, Vec<Id>>,
27    definition_ranges: FxHashSet<Id>,
28    graph: FxHashMap<i32, Vertex>,
29    edges: FxHashMap<i32, Vec<Edge>>,
30}
31
32impl IndexQueryCtx {
33    /// Reads the index from a reader.
34    pub fn read(reader: &mut impl BufRead) -> Result<Self> {
35        let mut this = Self::default();
36
37        for event in ReaderDecoder::new(reader) {
38            this.update(event?)?;
39        }
40
41        Ok(this)
42    }
43
44    fn update(&mut self, event: Entry) -> Result<()> {
45        match event.data {
46            Element::Vertex(Vertex::MetaData(meta)) => {
47                self.meta = Some(meta);
48            }
49            Element::Vertex(Vertex::Project(..)) => {
50                bail!("project is not supported");
51            }
52            Element::Vertex(Vertex::Event(..)) => {
53                bail!("event is not supported");
54            }
55            Element::Vertex(Vertex::Document(document)) => {
56                self.document_by_uri.insert(document.uri.clone(), event.id);
57                self.graph.insert(event.id, Vertex::Document(document));
58            }
59            Element::Vertex(vertex) => {
60                self.graph.insert(event.id, vertex);
61            }
62            Element::Edge(edge) => {
63                match &edge {
64                    Edge::Contains(data) => {
65                        self.document_ranges
66                            .entry(data.out_v)
67                            .or_default()
68                            .extend(data.in_vs.iter().copied());
69                    }
70                    Edge::Item(item) if self.is_definition_result(item.edge_data.out_v) => {
71                        self.definition_ranges
72                            .extend(item.edge_data.in_vs.iter().copied());
73                    }
74                    _ => {}
75                }
76                self.edges.entry(edge.out_v()).or_default().push(edge);
77            }
78        }
79
80        Ok(())
81    }
82
83    /// Requests the index for a compiler query.
84    pub fn request(&mut self, request: CompilerQueryRequest) -> Option<CompilerQueryResponse> {
85        match request {
86            CompilerQueryRequest::Hover(request) => {
87                Some(CompilerQueryResponse::Hover(self.hover(request)))
88            }
89            CompilerQueryRequest::GotoDefinition(request) => Some(
90                CompilerQueryResponse::GotoDefinition(self.goto_definition(request)),
91            ),
92            _ => None,
93        }
94    }
95
96    fn hover(&self, request: HoverRequest) -> Option<Hover> {
97        let uri = path_to_url(&request.path).ok()?;
98        let doc_id = *self.document_by_uri.get(&uri)?;
99        let source_range_id = self.find_range(doc_id, request.position)?;
100        let source_range = self.range_of(source_range_id)?;
101        let hover_result_id = self.hover_result(source_range_id).or_else(|| {
102            let result_id = self.next_result(source_range_id)?;
103            self.hover_result(result_id)
104        })?;
105        let mut hover = self.hover_value(hover_result_id)?;
106        hover.range.get_or_insert(source_range);
107        Some(hover)
108    }
109
110    fn goto_definition(&self, request: GotoDefinitionRequest) -> Option<GotoDefinitionResponse> {
111        let uri = path_to_url(&request.path).ok()?;
112        let doc_id = *self.document_by_uri.get(&uri)?;
113        let source_range_id = self.find_range(doc_id, request.position)?;
114        let origin_selection_range = self.range_of(source_range_id)?;
115        let result_id = match self.next_result(source_range_id) {
116            Some(result_id) => result_id,
117            None => {
118                return self.definition_target_link(
119                    doc_id,
120                    source_range_id,
121                    origin_selection_range,
122                );
123            }
124        };
125        let definition_result_id = self.definition_result(result_id)?;
126        let links = self.definition_links(definition_result_id, origin_selection_range);
127
128        if links.is_empty() {
129            None
130        } else {
131            Some(GotoDefinitionResponse::Link(links))
132        }
133    }
134
135    fn is_definition_result(&self, id: Id) -> bool {
136        matches!(self.graph.get(&id), Some(Vertex::DefinitionResult))
137    }
138
139    fn find_range(&self, document_id: Id, position: Position) -> Option<Id> {
140        self.document_ranges
141            .get(&document_id)?
142            .iter()
143            .filter_map(|id| Some((*id, self.range_of(*id)?)))
144            .filter(|(_, range)| contains_position(range, position))
145            .min_by_key(|(_, range)| range_len_key(range))
146            .map(|(id, _)| id)
147    }
148
149    fn range_of(&self, id: Id) -> Option<Range> {
150        match self.graph.get(&id)? {
151            Vertex::Range { range, .. } => Some(*range),
152            _ => None,
153        }
154    }
155
156    fn document_uri(&self, id: Id) -> Option<Url> {
157        match self.graph.get(&id)? {
158            Vertex::Document(document) => Some(document.uri.clone()),
159            _ => None,
160        }
161    }
162
163    fn definition_target_link(
164        &self,
165        document_id: Id,
166        range_id: Id,
167        target_range: Range,
168    ) -> Option<GotoDefinitionResponse> {
169        if !self.definition_ranges.contains(&range_id) {
170            return None;
171        }
172
173        Some(GotoDefinitionResponse::Link(vec![LocationLink {
174            origin_selection_range: Some(target_range),
175            target_uri: self.document_uri(document_id)?,
176            target_range,
177            target_selection_range: target_range,
178        }]))
179    }
180
181    fn next_result(&self, range_id: Id) -> Option<Id> {
182        self.edges
183            .get(&range_id)?
184            .iter()
185            .find_map(|edge| match edge {
186                Edge::Next(data) => Some(data.in_v),
187                _ => None,
188            })
189    }
190
191    fn definition_result(&self, result_id: Id) -> Option<Id> {
192        self.edges
193            .get(&result_id)?
194            .iter()
195            .find_map(|edge| match edge {
196                Edge::Definition(data) => Some(data.in_v),
197                _ => None,
198            })
199    }
200
201    fn hover_result(&self, out_v: Id) -> Option<Id> {
202        self.edges.get(&out_v)?.iter().find_map(|edge| match edge {
203            Edge::Hover(data) => Some(data.in_v),
204            _ => None,
205        })
206    }
207
208    fn hover_value(&self, hover_result_id: Id) -> Option<Hover> {
209        match self.graph.get(&hover_result_id)? {
210            Vertex::HoverResult { result } => Some(result.clone()),
211            _ => None,
212        }
213    }
214
215    fn definition_links(
216        &self,
217        definition_result_id: Id,
218        origin_selection_range: Range,
219    ) -> Vec<LocationLink> {
220        self.edges
221            .get(&definition_result_id)
222            .into_iter()
223            .flatten()
224            .filter_map(|edge| match edge {
225                Edge::Item(item) => Some(item),
226                _ => None,
227            })
228            .filter_map(|item| {
229                let target_uri = self.document_uri(item.document)?;
230                Some(item.edge_data.in_vs.iter().filter_map(move |range_id| {
231                    let target_selection_range = self.range_of(*range_id)?;
232                    Some(LocationLink {
233                        origin_selection_range: Some(origin_selection_range),
234                        target_uri: target_uri.clone(),
235                        target_range: target_selection_range,
236                        target_selection_range,
237                    })
238                }))
239            })
240            .flatten()
241            .collect()
242    }
243}
244
245// MetaData(MetaData),
246// /// <https://github.com/Microsoft/language-server-protocol/blob/master/indexFormat/specification.md#the-project-vertex>
247// Project(Project),
248// Document(Document),
249// /// <https://github.com/Microsoft/language-server-protocol/blob/master/indexFormat/specification.md#ranges>
250// Range {
251//     #[serde(flatten)]
252//     range: Range,
253//     #[serde(skip_serializing_if = "Option::is_none")]
254//     tag: Option<RangeTag>,
255// },
256// /// <https://github.com/Microsoft/language-server-protocol/blob/master/indexFormat/specification.md#result-set>
257// ResultSet(ResultSet),
258// Moniker(crate::Moniker),
259// PackageInformation(PackageInformation),
260
261// #[serde(rename = "$event")]
262// Event(Event),
263
264// DefinitionResult,
265// DeclarationResult,
266// TypeDefinitionResult,
267// ReferenceResult,
268// ImplementationResult,
269// FoldingRangeResult {
270//     result: Vec<crate::FoldingRange>,
271// },
272// HoverResult {
273//     result: crate::Hover,
274// },
275// DocumentSymbolResult {
276//     result: DocumentSymbolOrRangeBasedVec,
277// },
278// DocumentLinkResult {
279//     result: Vec<crate::DocumentLink>,
280// },
281// DiagnosticResult {
282//     result: Vec<crate::Diagnostic>,
283// },
284
285struct ReaderDecoder<R: Read> {
286    reader: R,
287}
288
289impl<R: BufRead> ReaderDecoder<R> {
290    pub fn new(reader: R) -> Self {
291        Self { reader }
292    }
293}
294
295impl<R: BufRead> Iterator for ReaderDecoder<R> {
296    type Item = Result<Entry>;
297
298    fn next(&mut self) -> Option<Result<Entry>> {
299        let mut line = String::new();
300        match self.reader.read_line(&mut line) {
301            Ok(0) => return None,
302            Ok(_) => {}
303            Err(e) => return Some(Err(e).context("failed to read line")),
304        }
305        Some(serde_json::from_str(&line).context("failed to deserialize entry"))
306    }
307}
308
309fn contains_position(range: &Range, position: Position) -> bool {
310    position_after_or_eq(position, range.start) && position_before(position, range.end)
311}
312
313fn position_after_or_eq(a: Position, b: Position) -> bool {
314    (a.line, a.character) >= (b.line, b.character)
315}
316
317fn position_before(a: Position, b: Position) -> bool {
318    (a.line, a.character) < (b.line, b.character)
319}
320
321fn range_len_key(range: &Range) -> (u32, u32) {
322    (
323        range.end.line.saturating_sub(range.start.line),
324        range.end.character.saturating_sub(range.start.character),
325    )
326}
327
328#[cfg(test)]
329mod tests {
330    use std::io::BufReader;
331    use std::path::{Path, PathBuf};
332
333    use lsp_types::{HoverContents, MarkedString};
334
335    use super::*;
336
337    #[test]
338    fn hover_from_result_set() {
339        let path = fixture_path();
340        let mut index = read_index([
341            document_line(&path),
342            r#"{"id":2,"type":"vertex","label":"range","start":{"line":0,"character":0},"end":{"line":0,"character":4}}"#.to_owned(),
343            r#"{"id":3,"type":"vertex","label":"resultSet"}"#.to_owned(),
344            r#"{"id":4,"type":"vertex","label":"hoverResult","result":{"contents":"hello"}}"#.to_owned(),
345            r#"{"id":5,"type":"edge","label":"contains","outV":1,"inVs":[2]}"#.to_owned(),
346            r#"{"id":6,"type":"edge","label":"next","outV":2,"inV":3}"#.to_owned(),
347            r#"{"id":7,"type":"edge","label":"textDocument/hover","outV":3,"inV":4}"#.to_owned(),
348        ]);
349
350        assert_hover(&mut index, path);
351    }
352
353    #[test]
354    fn hover_from_range() {
355        let path = fixture_path();
356        let mut index = read_index([
357            document_line(&path),
358            r#"{"id":2,"type":"vertex","label":"range","start":{"line":0,"character":0},"end":{"line":0,"character":4}}"#.to_owned(),
359            r#"{"id":3,"type":"vertex","label":"hoverResult","result":{"contents":"hello"}}"#.to_owned(),
360            r#"{"id":4,"type":"edge","label":"contains","outV":1,"inVs":[2]}"#.to_owned(),
361            r#"{"id":5,"type":"edge","label":"textDocument/hover","outV":2,"inV":3}"#.to_owned(),
362        ]);
363
364        assert_hover(&mut index, path);
365    }
366
367    fn fixture_path() -> PathBuf {
368        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("main.typ")
369    }
370
371    fn document_line(path: &Path) -> String {
372        let uri = path_to_url(path).unwrap();
373        format!(
374            r#"{{"id":1,"type":"vertex","label":"document","uri":{},"languageId":"typst"}}"#,
375            serde_json::to_string(uri.as_str()).unwrap()
376        )
377    }
378
379    fn read_index(lines: impl IntoIterator<Item = String>) -> IndexQueryCtx {
380        let mut input = lines.into_iter().collect::<Vec<_>>().join("\n");
381        input.push('\n');
382        IndexQueryCtx::read(&mut BufReader::new(input.as_bytes())).unwrap()
383    }
384
385    fn assert_hover(index: &mut IndexQueryCtx, path: PathBuf) {
386        let response = index.request(CompilerQueryRequest::Hover(HoverRequest {
387            path,
388            position: Position {
389                line: 0,
390                character: 1,
391            },
392        }));
393
394        let Some(CompilerQueryResponse::Hover(Some(hover))) = response else {
395            panic!("expected hover response");
396        };
397
398        assert_eq!(
399            hover.range,
400            Some(Range {
401                start: Position {
402                    line: 0,
403                    character: 0,
404                },
405                end: Position {
406                    line: 0,
407                    character: 4,
408                },
409            })
410        );
411        assert_eq!(
412            hover.contents,
413            HoverContents::Scalar(MarkedString::String("hello".to_owned()))
414        );
415    }
416}