Skip to main content

hearth_graph/
parse.rs

1use std::collections::hash_map::Entry;
2
3use rustc_hash::FxHashMap;
4use tree_sitter::{Parser, Query};
5
6use crate::{ImportSpec, LanguageId, LanguageRegistry};
7
8/// Lazily initialized parser and query cache for a language registry.
9pub struct ParserPool<'r> {
10    registry: &'r LanguageRegistry,
11    parsers: FxHashMap<LanguageId, Option<Parser>>,
12    tags_queries: FxHashMap<LanguageId, Option<Query>>,
13    imports_queries: FxHashMap<LanguageId, Option<Query>>,
14    injections_queries: FxHashMap<LanguageId, Option<Query>>,
15}
16
17impl<'r> ParserPool<'r> {
18    /// Creates an empty cache backed by `registry`.
19    #[must_use]
20    pub fn new(registry: &'r LanguageRegistry) -> Self {
21        Self {
22            registry,
23            parsers: FxHashMap::default(),
24            tags_queries: FxHashMap::default(),
25            imports_queries: FxHashMap::default(),
26            injections_queries: FxHashMap::default(),
27        }
28    }
29
30    /// Returns the registry whose identifiers this pool accepts.
31    pub(crate) fn registry(&self) -> &'r LanguageRegistry {
32        self.registry
33    }
34
35    /// Returns a parser for `id`, creating and configuring it on first use.
36    ///
37    /// A grammar rejected by tree-sitter is cached as a permanent miss.
38    pub fn parser(&mut self, id: LanguageId) -> Option<&mut Parser> {
39        if let Entry::Vacant(entry) = self.parsers.entry(id) {
40            let parser = self.registry.get(id).and_then(|spec| {
41                let mut parser = Parser::new();
42                parser.set_language(&spec.language).ok()?;
43                Some(parser)
44            });
45            entry.insert(parser);
46        }
47
48        self.parsers.get_mut(&id).and_then(Option::as_mut)
49    }
50
51    /// Returns the compiled tags query for `id`.
52    ///
53    /// Missing query sources and compilation failures are cached permanently.
54    pub fn tags_query(&mut self, id: LanguageId) -> Option<&Query> {
55        if let Entry::Vacant(entry) = self.tags_queries.entry(id) {
56            let query = self.registry.get(id).and_then(|spec| {
57                let source = spec.tags_query.as_deref()?;
58                Query::new(&spec.language, source).ok()
59            });
60            entry.insert(query);
61        }
62
63        self.tags_queries.get(&id).and_then(Option::as_ref)
64    }
65
66    /// Returns the compiled query-based import extractor for `id`.
67    ///
68    /// Custom extractors, missing queries, and compilation failures are cached
69    /// as permanent misses.
70    pub fn imports_query(&mut self, id: LanguageId) -> Option<&Query> {
71        if let Entry::Vacant(entry) = self.imports_queries.entry(id) {
72            let query = self
73                .registry
74                .get(id)
75                .and_then(|spec| match spec.imports.as_ref()? {
76                    ImportSpec::Query { source, .. } => {
77                        Query::new(&spec.language, source.as_ref()).ok()
78                    }
79                    ImportSpec::Custom(_) => None,
80                });
81            entry.insert(query);
82        }
83
84        self.imports_queries.get(&id).and_then(Option::as_ref)
85    }
86
87    /// Returns the compiled injection query for `id`.
88    ///
89    /// Missing query sources and compilation failures are cached permanently.
90    pub fn injections_query(&mut self, id: LanguageId) -> Option<&Query> {
91        if let Entry::Vacant(entry) = self.injections_queries.entry(id) {
92            let query = self.registry.get(id).and_then(|spec| {
93                let source = spec.injections_query.as_deref()?;
94                Query::new(&spec.language, source).ok()
95            });
96            entry.insert(query);
97        }
98
99        self.injections_queries.get(&id).and_then(Option::as_ref)
100    }
101}