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::Array(items) => {
289 for elem in items {
290 collect_expr(elem, out);
291 }
292 }
293 Expr::ArraySubscript { target, index } => {
294 collect_expr(target, out);
295 collect_expr(index, out);
296 }
297 Expr::AnyAll { expr, array, .. } => {
298 collect_expr(expr, out);
299 collect_expr(array, out);
300 }
301 Expr::Literal(_) | Expr::Column(_) | Expr::Placeholder(_) => {}
302 }
303}
304
305fn push_unique(out: &mut Vec<String>, s: &str) {
306 if !out.iter().any(|x| x == s) {
307 out.push(String::from(s));
308 }
309}
310
311#[cfg(test)]
314mod tests {
315 use super::*;
316 use alloc::string::ToString;
317 use spg_sql::parser::parse_statement;
318
319 fn dummy_plan(version: u64, tables: &[&str]) -> PreparedPlan {
320 let stmt = parse_statement("SELECT 1").expect("trivial SELECT parses");
321 PreparedPlan {
322 stmt,
323 statistics_version: version,
324 source_tables: tables.iter().map(|s| s.to_string()).collect(),
325 describe_columns: Vec::new(),
326 }
327 }
328
329 #[test]
330 fn new_cache_is_empty() {
331 let cache = PlanCache::new();
332 assert!(cache.is_empty());
333 assert_eq!(cache.len(), 0);
334 }
335
336 #[test]
337 fn insert_then_get_returns_the_plan() {
338 let mut cache = PlanCache::new();
339 cache.insert("SELECT 1".into(), dummy_plan(0, &["t"]));
340 assert_eq!(cache.len(), 1);
341 let plan = cache.get("SELECT 1").expect("hit");
342 assert_eq!(plan.source_tables, alloc::vec!["t".to_string()]);
343 }
344
345 #[test]
346 fn miss_returns_none() {
347 let mut cache = PlanCache::new();
348 cache.insert("SELECT 1".into(), dummy_plan(0, &[]));
349 assert!(cache.get("SELECT 2").is_none());
350 }
351
352 #[test]
353 fn replace_overwrites_existing_entry() {
354 let mut cache = PlanCache::new();
355 cache.insert("SELECT 1".into(), dummy_plan(1, &["a"]));
356 cache.insert("SELECT 1".into(), dummy_plan(2, &["b"]));
357 assert_eq!(cache.len(), 1);
358 let plan = cache.get("SELECT 1").expect("hit");
359 assert_eq!(plan.statistics_version, 2);
360 }
361
362 #[test]
363 fn lru_evicts_oldest_at_cap() {
364 let mut cache = PlanCache::new();
365 for i in 0..PLAN_CACHE_MAX_ENTRIES {
366 cache.insert(alloc::format!("SELECT {i}"), dummy_plan(i as u64, &[]));
367 }
368 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
369 cache.insert("SELECT new".into(), dummy_plan(999, &[]));
370 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
371 assert!(cache.get("SELECT 0").is_none());
372 assert!(cache.get("SELECT new").is_some());
373 }
374
375 #[test]
376 fn get_promotes_lru_position() {
377 let mut cache = PlanCache::new();
378 cache.insert("a".into(), dummy_plan(0, &[]));
379 cache.insert("b".into(), dummy_plan(0, &[]));
380 cache.insert("c".into(), dummy_plan(0, &[]));
381 let _ = cache.get("a");
383 for i in 0..(PLAN_CACHE_MAX_ENTRIES - 3) {
386 cache.insert(alloc::format!("filler{i}"), dummy_plan(0, &[]));
387 }
388 cache.insert("trigger".into(), dummy_plan(0, &[]));
389 assert!(cache.get("a").is_some(), "a was MRU after get(); should survive");
390 assert!(cache.get("b").is_none(), "b should be evicted");
391 }
392
393 #[test]
394 fn clear_drops_everything() {
395 let mut cache = PlanCache::new();
396 cache.insert("a".into(), dummy_plan(0, &[]));
397 cache.insert("b".into(), dummy_plan(0, &[]));
398 cache.clear();
399 assert!(cache.is_empty());
400 assert!(cache.get("a").is_none());
401 }
402
403 #[test]
404 fn evict_referencing_drops_only_matching_plans() {
405 let mut cache = PlanCache::new();
406 cache.insert("a".into(), dummy_plan(0, &["users"]));
407 cache.insert("b".into(), dummy_plan(0, &["orders"]));
408 cache.insert("c".into(), dummy_plan(0, &["users", "orders"]));
409 let n = cache.evict_referencing("users");
410 assert_eq!(n, 2);
411 assert!(cache.get("a").is_none());
412 assert!(cache.get("b").is_some());
413 assert!(cache.get("c").is_none());
414 }
415
416 #[test]
417 fn collect_source_tables_from_simple_select() {
418 let stmt = parse_statement("SELECT a, b FROM t1 WHERE x = 1").expect("parses");
419 let tables = collect_source_tables(&stmt);
420 assert_eq!(tables, alloc::vec!["t1".to_string()]);
421 }
422
423 #[test]
424 fn collect_source_tables_from_join() {
425 let stmt = parse_statement(
426 "SELECT * FROM t1 JOIN t2 ON t1.a = t2.b JOIN t3 ON t2.c = t3.d",
427 )
428 .expect("parses");
429 let tables = collect_source_tables(&stmt);
430 assert_eq!(
431 tables,
432 alloc::vec!["t1".to_string(), "t2".to_string(), "t3".to_string()]
433 );
434 }
435
436 #[test]
437 fn collect_source_tables_dedupes_self_join() {
438 let stmt = parse_statement("SELECT * FROM t1 a JOIN t1 b ON a.x = b.y").expect("parses");
439 let tables = collect_source_tables(&stmt);
440 assert_eq!(tables, alloc::vec!["t1".to_string()]);
441 }
442}