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::AggregateOrdered { call, order_by } => {
255 collect_expr(call, out);
256 for o in order_by {
257 collect_expr(&o.expr, out);
258 }
259 }
260 Expr::ScalarSubquery(inner) => collect_from_select(inner, out),
261 Expr::Exists { subquery, .. } => collect_from_select(subquery, out),
262 Expr::InSubquery { expr, subquery, .. } => {
263 collect_expr(expr, out);
264 collect_from_select(subquery, out);
265 }
266 Expr::Binary { lhs, rhs, .. } => {
267 collect_expr(lhs, out);
268 collect_expr(rhs, out);
269 }
270 Expr::Unary { expr, .. } => collect_expr(expr, out),
271 Expr::Cast { expr, .. } => collect_expr(expr, out),
272 Expr::IsNull { expr, .. } => collect_expr(expr, out),
273 Expr::Like { expr, pattern, .. } => {
274 collect_expr(expr, out);
275 collect_expr(pattern, out);
276 }
277 Expr::FunctionCall { args, .. } => {
278 for a in args {
279 collect_expr(a, out);
280 }
281 }
282 Expr::WindowFunction {
283 args,
284 partition_by,
285 order_by,
286 ..
287 } => {
288 for a in args {
289 collect_expr(a, out);
290 }
291 for p in partition_by {
292 collect_expr(p, out);
293 }
294 for (o, _) in order_by {
295 collect_expr(o, out);
296 }
297 }
298 Expr::Extract { source, .. } => collect_expr(source, out),
299 Expr::Array(items) => {
300 for elem in items {
301 collect_expr(elem, out);
302 }
303 }
304 Expr::ArraySubscript { target, index } => {
305 collect_expr(target, out);
306 collect_expr(index, out);
307 }
308 Expr::AnyAll { expr, array, .. } => {
309 collect_expr(expr, out);
310 collect_expr(array, out);
311 }
312 Expr::Case {
313 operand,
314 branches,
315 else_branch,
316 } => {
317 if let Some(o) = operand {
318 collect_expr(o, out);
319 }
320 for (w, t) in branches {
321 collect_expr(w, out);
322 collect_expr(t, out);
323 }
324 if let Some(e) = else_branch {
325 collect_expr(e, out);
326 }
327 }
328 Expr::Literal(_) | Expr::Column(_) | Expr::Placeholder(_) => {}
329 }
330}
331
332fn push_unique(out: &mut Vec<String>, s: &str) {
333 if !out.iter().any(|x| x == s) {
334 out.push(String::from(s));
335 }
336}
337
338#[cfg(test)]
341mod tests {
342 use super::*;
343 use alloc::string::ToString;
344 use spg_sql::parser::parse_statement;
345
346 fn dummy_plan(version: u64, tables: &[&str]) -> PreparedPlan {
347 let stmt = parse_statement("SELECT 1").expect("trivial SELECT parses");
348 PreparedPlan {
349 stmt,
350 statistics_version: version,
351 source_tables: tables.iter().map(|s| s.to_string()).collect(),
352 describe_columns: Vec::new(),
353 }
354 }
355
356 #[test]
357 fn new_cache_is_empty() {
358 let cache = PlanCache::new();
359 assert!(cache.is_empty());
360 assert_eq!(cache.len(), 0);
361 }
362
363 #[test]
364 fn insert_then_get_returns_the_plan() {
365 let mut cache = PlanCache::new();
366 cache.insert("SELECT 1".into(), dummy_plan(0, &["t"]));
367 assert_eq!(cache.len(), 1);
368 let plan = cache.get("SELECT 1").expect("hit");
369 assert_eq!(plan.source_tables, alloc::vec!["t".to_string()]);
370 }
371
372 #[test]
373 fn miss_returns_none() {
374 let mut cache = PlanCache::new();
375 cache.insert("SELECT 1".into(), dummy_plan(0, &[]));
376 assert!(cache.get("SELECT 2").is_none());
377 }
378
379 #[test]
380 fn replace_overwrites_existing_entry() {
381 let mut cache = PlanCache::new();
382 cache.insert("SELECT 1".into(), dummy_plan(1, &["a"]));
383 cache.insert("SELECT 1".into(), dummy_plan(2, &["b"]));
384 assert_eq!(cache.len(), 1);
385 let plan = cache.get("SELECT 1").expect("hit");
386 assert_eq!(plan.statistics_version, 2);
387 }
388
389 #[test]
390 fn lru_evicts_oldest_at_cap() {
391 let mut cache = PlanCache::new();
392 for i in 0..PLAN_CACHE_MAX_ENTRIES {
393 cache.insert(alloc::format!("SELECT {i}"), dummy_plan(i as u64, &[]));
394 }
395 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
396 cache.insert("SELECT new".into(), dummy_plan(999, &[]));
397 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
398 assert!(cache.get("SELECT 0").is_none());
399 assert!(cache.get("SELECT new").is_some());
400 }
401
402 #[test]
403 fn get_promotes_lru_position() {
404 let mut cache = PlanCache::new();
405 cache.insert("a".into(), dummy_plan(0, &[]));
406 cache.insert("b".into(), dummy_plan(0, &[]));
407 cache.insert("c".into(), dummy_plan(0, &[]));
408 let _ = cache.get("a");
410 for i in 0..(PLAN_CACHE_MAX_ENTRIES - 3) {
413 cache.insert(alloc::format!("filler{i}"), dummy_plan(0, &[]));
414 }
415 cache.insert("trigger".into(), dummy_plan(0, &[]));
416 assert!(
417 cache.get("a").is_some(),
418 "a was MRU after get(); should survive"
419 );
420 assert!(cache.get("b").is_none(), "b should be evicted");
421 }
422
423 #[test]
424 fn clear_drops_everything() {
425 let mut cache = PlanCache::new();
426 cache.insert("a".into(), dummy_plan(0, &[]));
427 cache.insert("b".into(), dummy_plan(0, &[]));
428 cache.clear();
429 assert!(cache.is_empty());
430 assert!(cache.get("a").is_none());
431 }
432
433 #[test]
434 fn evict_referencing_drops_only_matching_plans() {
435 let mut cache = PlanCache::new();
436 cache.insert("a".into(), dummy_plan(0, &["users"]));
437 cache.insert("b".into(), dummy_plan(0, &["orders"]));
438 cache.insert("c".into(), dummy_plan(0, &["users", "orders"]));
439 let n = cache.evict_referencing("users");
440 assert_eq!(n, 2);
441 assert!(cache.get("a").is_none());
442 assert!(cache.get("b").is_some());
443 assert!(cache.get("c").is_none());
444 }
445
446 #[test]
447 fn collect_source_tables_from_simple_select() {
448 let stmt = parse_statement("SELECT a, b FROM t1 WHERE x = 1").expect("parses");
449 let tables = collect_source_tables(&stmt);
450 assert_eq!(tables, alloc::vec!["t1".to_string()]);
451 }
452
453 #[test]
454 fn collect_source_tables_from_join() {
455 let stmt =
456 parse_statement("SELECT * FROM t1 JOIN t2 ON t1.a = t2.b JOIN t3 ON t2.c = t3.d")
457 .expect("parses");
458 let tables = collect_source_tables(&stmt);
459 assert_eq!(
460 tables,
461 alloc::vec!["t1".to_string(), "t2".to_string(), "t3".to_string()]
462 );
463 }
464
465 #[test]
466 fn collect_source_tables_dedupes_self_join() {
467 let stmt = parse_statement("SELECT * FROM t1 a JOIN t1 b ON a.x = b.y").expect("parses");
468 let tables = collect_source_tables(&stmt);
469 assert_eq!(tables, alloc::vec!["t1".to_string()]);
470 }
471}