1use alloc::collections::{BTreeMap, VecDeque};
20use alloc::string::String;
21use alloc::vec::Vec;
22
23use spg_sql::ast::{Expr, FromClause, FromJoin, SelectItem, SelectStatement, Statement, TableRef};
24use spg_storage::ColumnSchema;
25
26pub(crate) const PLAN_CACHE_MAX_ENTRIES: usize = 256;
31
32#[derive(Debug, Clone)]
36pub struct PreparedPlan {
37 pub stmt: Statement,
38 pub statistics_version: u64,
42 pub source_tables: Vec<String>,
45 pub describe_columns: Vec<ColumnSchema>,
48}
49
50#[derive(Debug, Clone)]
51pub struct PlanCache {
52 entries: BTreeMap<String, PreparedPlan>,
56 lru: VecDeque<String>,
60 max_entries: usize,
65}
66
67impl Default for PlanCache {
68 fn default() -> Self {
69 Self {
70 entries: BTreeMap::new(),
71 lru: VecDeque::new(),
72 max_entries: PLAN_CACHE_MAX_ENTRIES,
73 }
74 }
75}
76
77impl PlanCache {
78 pub fn new() -> Self {
79 Self::default()
80 }
81
82 pub fn set_max_entries(&mut self, n: usize) {
88 self.max_entries = n.max(1).min(PLAN_CACHE_MAX_ENTRIES);
89 }
90
91 pub fn max_entries(&self) -> usize {
92 self.max_entries
93 }
94
95 pub fn len(&self) -> usize {
96 self.entries.len()
97 }
98
99 pub fn is_empty(&self) -> bool {
100 self.entries.is_empty()
101 }
102
103 pub fn get_snapshot(&self, sql: &str) -> Option<&PreparedPlan> {
107 self.entries.get(sql)
108 }
109
110 pub fn get(&mut self, sql: &str) -> Option<&PreparedPlan> {
113 if !self.entries.contains_key(sql) {
114 return None;
115 }
116 if let Some(idx) = self.lru.iter().position(|k| k == sql) {
117 let key = self.lru.remove(idx).expect("idx came from position()");
118 self.lru.push_back(key);
119 }
120 self.entries.get(sql)
121 }
122
123 pub fn insert(&mut self, sql: String, plan: PreparedPlan) {
126 if self.entries.contains_key(&sql) {
127 if let Some(idx) = self.lru.iter().position(|k| k == &sql) {
128 let key = self.lru.remove(idx).expect("idx came from position()");
129 self.lru.push_back(key);
130 }
131 self.entries.insert(sql, plan);
132 return;
133 }
134 if self.entries.len() >= self.max_entries {
135 if let Some(oldest) = self.lru.pop_front() {
136 self.entries.remove(&oldest);
137 }
138 }
139 self.lru.push_back(sql.clone());
140 self.entries.insert(sql, plan);
141 }
142
143 pub fn clear(&mut self) {
144 self.entries.clear();
145 self.lru.clear();
146 }
147
148 pub fn evict(&mut self, sql: &str) -> Option<PreparedPlan> {
151 let plan = self.entries.remove(sql)?;
152 if let Some(idx) = self.lru.iter().position(|k| k == sql) {
153 self.lru.remove(idx);
154 }
155 Some(plan)
156 }
157
158 pub fn evict_referencing(&mut self, table: &str) -> usize {
161 let to_evict: Vec<String> = self
162 .entries
163 .iter()
164 .filter_map(|(k, p)| {
165 if p.source_tables.iter().any(|t| t == table) {
166 Some(k.clone())
167 } else {
168 None
169 }
170 })
171 .collect();
172 let n = to_evict.len();
173 for k in to_evict {
174 self.entries.remove(&k);
175 if let Some(idx) = self.lru.iter().position(|x| x == &k) {
176 self.lru.remove(idx);
177 }
178 }
179 n
180 }
181}
182
183pub fn collect_source_tables(stmt: &Statement) -> Vec<String> {
188 let mut out: Vec<String> = Vec::new();
189 match stmt {
190 Statement::Select(s) => collect_from_select(s, &mut out),
191 Statement::Insert(s) => push_unique(&mut out, &s.table),
192 Statement::Update(s) => {
193 push_unique(&mut out, &s.table);
194 if let Some(w) = &s.where_ {
195 collect_expr(w, &mut out);
196 }
197 }
198 Statement::Delete(s) => {
199 push_unique(&mut out, &s.table);
200 if let Some(w) = &s.where_ {
201 collect_expr(w, &mut out);
202 }
203 }
204 Statement::Explain(inner) => {
205 collect_from_select(&inner.inner, &mut out);
206 }
207 _ => {}
208 }
209 out.sort();
210 out.dedup();
211 out
212}
213
214fn collect_from_select(s: &SelectStatement, out: &mut Vec<String>) {
215 if let Some(from) = &s.from {
216 collect_from_clause(from, out);
217 }
218 if let Some(w) = &s.where_ {
219 collect_expr(w, out);
220 }
221 if let Some(h) = &s.having {
222 collect_expr(h, out);
223 }
224 for item in &s.items {
225 if let SelectItem::Expr { expr, .. } = item {
226 collect_expr(expr, out);
227 }
228 }
229 for (_, peer) in &s.unions {
230 collect_from_select(peer, out);
231 }
232}
233
234fn collect_from_clause(from: &FromClause, out: &mut Vec<String>) {
235 collect_table_ref(&from.primary, out);
236 for j in &from.joins {
237 collect_from_join(j, out);
238 }
239}
240
241fn collect_from_join(j: &FromJoin, out: &mut Vec<String>) {
242 collect_table_ref(&j.table, out);
243 if let Some(on) = &j.on {
244 collect_expr(on, out);
245 }
246}
247
248fn collect_table_ref(t: &TableRef, out: &mut Vec<String>) {
249 push_unique(out, &t.name);
250}
251
252fn collect_expr(e: &Expr, out: &mut Vec<String>) {
253 match e {
254 Expr::ScalarSubquery(inner) => collect_from_select(inner, out),
255 Expr::Exists { subquery, .. } => collect_from_select(subquery, out),
256 Expr::InSubquery { expr, subquery, .. } => {
257 collect_expr(expr, out);
258 collect_from_select(subquery, out);
259 }
260 Expr::Binary { lhs, rhs, .. } => {
261 collect_expr(lhs, out);
262 collect_expr(rhs, out);
263 }
264 Expr::Unary { expr, .. } => collect_expr(expr, out),
265 Expr::Cast { expr, .. } => collect_expr(expr, out),
266 Expr::IsNull { expr, .. } => collect_expr(expr, out),
267 Expr::Like { expr, pattern, .. } => {
268 collect_expr(expr, out);
269 collect_expr(pattern, out);
270 }
271 Expr::FunctionCall { args, .. } => {
272 for a in args {
273 collect_expr(a, out);
274 }
275 }
276 Expr::WindowFunction { args, partition_by, order_by, .. } => {
277 for a in args {
278 collect_expr(a, out);
279 }
280 for p in partition_by {
281 collect_expr(p, out);
282 }
283 for (o, _) in order_by {
284 collect_expr(o, out);
285 }
286 }
287 Expr::Extract { source, .. } => collect_expr(source, out),
288 Expr::Literal(_) | Expr::Column(_) | Expr::Placeholder(_) => {}
289 }
290}
291
292fn push_unique(out: &mut Vec<String>, s: &str) {
293 if !out.iter().any(|x| x == s) {
294 out.push(String::from(s));
295 }
296}
297
298#[cfg(test)]
301mod tests {
302 use super::*;
303 use alloc::string::ToString;
304 use spg_sql::parser::parse_statement;
305
306 fn dummy_plan(version: u64, tables: &[&str]) -> PreparedPlan {
307 let stmt = parse_statement("SELECT 1").expect("trivial SELECT parses");
308 PreparedPlan {
309 stmt,
310 statistics_version: version,
311 source_tables: tables.iter().map(|s| s.to_string()).collect(),
312 describe_columns: Vec::new(),
313 }
314 }
315
316 #[test]
317 fn new_cache_is_empty() {
318 let cache = PlanCache::new();
319 assert!(cache.is_empty());
320 assert_eq!(cache.len(), 0);
321 }
322
323 #[test]
324 fn insert_then_get_returns_the_plan() {
325 let mut cache = PlanCache::new();
326 cache.insert("SELECT 1".into(), dummy_plan(0, &["t"]));
327 assert_eq!(cache.len(), 1);
328 let plan = cache.get("SELECT 1").expect("hit");
329 assert_eq!(plan.source_tables, alloc::vec!["t".to_string()]);
330 }
331
332 #[test]
333 fn miss_returns_none() {
334 let mut cache = PlanCache::new();
335 cache.insert("SELECT 1".into(), dummy_plan(0, &[]));
336 assert!(cache.get("SELECT 2").is_none());
337 }
338
339 #[test]
340 fn replace_overwrites_existing_entry() {
341 let mut cache = PlanCache::new();
342 cache.insert("SELECT 1".into(), dummy_plan(1, &["a"]));
343 cache.insert("SELECT 1".into(), dummy_plan(2, &["b"]));
344 assert_eq!(cache.len(), 1);
345 let plan = cache.get("SELECT 1").expect("hit");
346 assert_eq!(plan.statistics_version, 2);
347 }
348
349 #[test]
350 fn lru_evicts_oldest_at_cap() {
351 let mut cache = PlanCache::new();
352 for i in 0..PLAN_CACHE_MAX_ENTRIES {
353 cache.insert(alloc::format!("SELECT {i}"), dummy_plan(i as u64, &[]));
354 }
355 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
356 cache.insert("SELECT new".into(), dummy_plan(999, &[]));
357 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
358 assert!(cache.get("SELECT 0").is_none());
359 assert!(cache.get("SELECT new").is_some());
360 }
361
362 #[test]
363 fn get_promotes_lru_position() {
364 let mut cache = PlanCache::new();
365 cache.insert("a".into(), dummy_plan(0, &[]));
366 cache.insert("b".into(), dummy_plan(0, &[]));
367 cache.insert("c".into(), dummy_plan(0, &[]));
368 let _ = cache.get("a");
370 for i in 0..(PLAN_CACHE_MAX_ENTRIES - 3) {
373 cache.insert(alloc::format!("filler{i}"), dummy_plan(0, &[]));
374 }
375 cache.insert("trigger".into(), dummy_plan(0, &[]));
376 assert!(cache.get("a").is_some(), "a was MRU after get(); should survive");
377 assert!(cache.get("b").is_none(), "b should be evicted");
378 }
379
380 #[test]
381 fn clear_drops_everything() {
382 let mut cache = PlanCache::new();
383 cache.insert("a".into(), dummy_plan(0, &[]));
384 cache.insert("b".into(), dummy_plan(0, &[]));
385 cache.clear();
386 assert!(cache.is_empty());
387 assert!(cache.get("a").is_none());
388 }
389
390 #[test]
391 fn evict_referencing_drops_only_matching_plans() {
392 let mut cache = PlanCache::new();
393 cache.insert("a".into(), dummy_plan(0, &["users"]));
394 cache.insert("b".into(), dummy_plan(0, &["orders"]));
395 cache.insert("c".into(), dummy_plan(0, &["users", "orders"]));
396 let n = cache.evict_referencing("users");
397 assert_eq!(n, 2);
398 assert!(cache.get("a").is_none());
399 assert!(cache.get("b").is_some());
400 assert!(cache.get("c").is_none());
401 }
402
403 #[test]
404 fn collect_source_tables_from_simple_select() {
405 let stmt = parse_statement("SELECT a, b FROM t1 WHERE x = 1").expect("parses");
406 let tables = collect_source_tables(&stmt);
407 assert_eq!(tables, alloc::vec!["t1".to_string()]);
408 }
409
410 #[test]
411 fn collect_source_tables_from_join() {
412 let stmt = parse_statement(
413 "SELECT * FROM t1 JOIN t2 ON t1.a = t2.b JOIN t3 ON t2.c = t3.d",
414 )
415 .expect("parses");
416 let tables = collect_source_tables(&stmt);
417 assert_eq!(
418 tables,
419 alloc::vec!["t1".to_string(), "t2".to_string(), "t3".to_string()]
420 );
421 }
422
423 #[test]
424 fn collect_source_tables_dedupes_self_join() {
425 let stmt = parse_statement("SELECT * FROM t1 a JOIN t1 b ON a.x = b.y").expect("parses");
426 let tables = collect_source_tables(&stmt);
427 assert_eq!(tables, alloc::vec!["t1".to_string()]);
428 }
429}