reovim_plugin_treesitter/
queries.rs

1//! Query compilation and caching
2//!
3//! Queries are provided by language plugins via the LanguageSupport trait.
4
5use std::{
6    collections::HashMap,
7    sync::{Arc, RwLock},
8};
9
10use tree_sitter::Query;
11
12/// Type of query
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum QueryType {
15    /// Syntax highlighting
16    Highlights,
17    /// Semantic text objects (function, class, parameter, etc.)
18    TextObjects,
19    /// Code folding regions
20    Folds,
21    /// Visual decorations (concealment, icons, backgrounds)
22    Decorations,
23    /// Language injections (embedded code in markdown, etc.)
24    Injections,
25    /// Context/scope detection (functions, classes, headings)
26    Context,
27}
28
29/// Cache key for compiled queries
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31struct QueryKey {
32    language_id: String,
33    query_type: QueryType,
34}
35
36/// Cache for compiled queries
37///
38/// Queries are compiled lazily on first access from the source provided by language plugins.
39/// Uses interior mutability (RwLock) to allow lazy compilation through `&self` references.
40/// Uses Arc<Query> to allow cheap cloning for syntax providers.
41pub struct QueryCache {
42    queries: RwLock<HashMap<QueryKey, Arc<Query>>>,
43}
44
45impl Default for QueryCache {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl QueryCache {
52    /// Create a new empty query cache
53    #[must_use]
54    pub fn new() -> Self {
55        Self {
56            queries: RwLock::new(HashMap::new()),
57        }
58    }
59
60    /// Get a cached query (does not compile if not cached)
61    ///
62    /// Returns an Arc clone for cheap sharing with syntax providers.
63    #[must_use]
64    pub fn get(&self, language_id: &str, query_type: QueryType) -> Option<Arc<Query>> {
65        let key = QueryKey {
66            language_id: language_id.to_string(),
67            query_type,
68        };
69        self.queries.read().unwrap().get(&key).cloned()
70    }
71
72    /// Compile and cache a query from source
73    ///
74    /// Returns an Arc clone of the compiled query, or None if compilation fails.
75    /// Uses interior mutability via RwLock.
76    pub fn compile_and_cache(
77        &self,
78        language_id: &str,
79        query_type: QueryType,
80        ts_language: &tree_sitter::Language,
81        source: &str,
82    ) -> Option<Arc<Query>> {
83        let key = QueryKey {
84            language_id: language_id.to_string(),
85            query_type,
86        };
87
88        // Return cached query if available (fast path with read lock)
89        if let Some(query) = self.queries.read().unwrap().get(&key) {
90            return Some(Arc::clone(query));
91        }
92
93        // Compile the query
94        let query = match Query::new(ts_language, source) {
95            Ok(q) => q,
96            Err(e) => {
97                tracing::error!(
98                    language_id = %language_id,
99                    query_type = ?query_type,
100                    error = %e,
101                    "Failed to compile query"
102                );
103                return None;
104            }
105        };
106
107        let arc_query = Arc::new(query);
108        self.queries
109            .write()
110            .unwrap()
111            .insert(key, Arc::clone(&arc_query));
112        tracing::debug!(
113            language_id = %language_id,
114            query_type = ?query_type,
115            "Compiled and cached query"
116        );
117        Some(arc_query)
118    }
119
120    /// Get a cached query, or compile and cache it on first access (lazy compilation)
121    ///
122    /// This is the primary method for lazy query access. Checks the cache first
123    /// with a read lock, then compiles with a write lock if not cached.
124    pub fn get_or_compile(
125        &self,
126        language_id: &str,
127        query_type: QueryType,
128        ts_language: &tree_sitter::Language,
129        source: &str,
130    ) -> Option<Arc<Query>> {
131        let key = QueryKey {
132            language_id: language_id.to_string(),
133            query_type,
134        };
135
136        // Fast path: check read lock first
137        if let Some(query) = self.queries.read().unwrap().get(&key) {
138            return Some(Arc::clone(query));
139        }
140
141        // Slow path: compile and cache with write lock
142        self.compile_and_cache(language_id, query_type, ts_language, source)
143    }
144
145    /// Check if a query is cached
146    #[must_use]
147    pub fn is_cached(&self, language_id: &str, query_type: QueryType) -> bool {
148        let key = QueryKey {
149            language_id: language_id.to_string(),
150            query_type,
151        };
152        self.queries.read().unwrap().contains_key(&key)
153    }
154
155    /// Clear all cached queries
156    pub fn clear(&self) {
157        self.queries.write().unwrap().clear();
158    }
159
160    /// Clear cached queries for a specific language
161    pub fn clear_language(&self, language_id: &str) {
162        self.queries
163            .write()
164            .unwrap()
165            .retain(|k, _| k.language_id != language_id);
166    }
167}
reovim_plugin_treesitter/queries.rs

reovim_plugin_treesitter/
queries.rs
