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