1use crate::ast::{Assignment, Expr, Literal};
29use crate::plan::PlanNode;
30use rustc_hash::FxHashMap;
31
32pub struct PlanCache {
38 cache: FxHashMap<u64, PlanNode>,
39 capacity: usize,
40 pub hits: u64,
41 pub misses: u64,
42}
43
44impl PlanCache {
45 pub fn new(capacity: usize) -> Self {
46 PlanCache {
47 cache: FxHashMap::default(),
48 capacity,
49 hits: 0,
50 misses: 0,
51 }
52 }
53
54 pub fn insert(&mut self, hash: u64, plan: PlanNode) {
57 if self.cache.len() >= self.capacity && !self.cache.contains_key(&hash) {
58 self.cache.clear();
63 }
64 self.cache.insert(hash, plan);
65 }
66
67 pub fn get_with_substitution(&mut self, hash: u64, literals: &[Literal]) -> Option<PlanNode> {
79 match self.cache.get(&hash) {
80 Some(template) => {
81 self.hits += 1;
82 let mut plan = template.clone();
83 let mut idx = 0usize;
84 substitute_plan(&mut plan, literals, &mut idx);
85 debug_assert_eq!(
86 idx,
87 literals.len(),
88 "plan substitution consumed {idx} literals but query had {}",
89 literals.len(),
90 );
91 Some(plan)
92 }
93 None => {
94 self.misses += 1;
95 None
96 }
97 }
98 }
99
100 pub fn len(&self) -> usize {
101 self.cache.len()
102 }
103
104 pub fn is_empty(&self) -> bool {
105 self.cache.is_empty()
106 }
107
108 pub fn clear(&mut self) {
109 self.cache.clear();
110 }
111}
112
113pub(crate) fn substitute_plan(plan: &mut PlanNode, literals: &[Literal], idx: &mut usize) {
132 match plan {
133 PlanNode::SeqScan { .. } => {}
134 PlanNode::AliasScan { .. } => {}
135 PlanNode::IndexScan { key, .. } => {
136 substitute_expr(key, literals, idx);
137 }
138 PlanNode::RangeScan { start, end, .. } => {
139 if let Some((expr, _)) = start {
140 substitute_expr(expr, literals, idx);
141 }
142 if let Some((expr, _)) = end {
143 substitute_expr(expr, literals, idx);
144 }
145 }
146 PlanNode::Filter { input, predicate } => {
147 substitute_plan(input, literals, idx);
148 substitute_expr(predicate, literals, idx);
149 }
150 PlanNode::Project { input, fields } => {
151 substitute_plan(input, literals, idx);
152 for f in fields {
153 substitute_expr(&mut f.expr, literals, idx);
154 }
155 }
156 PlanNode::Sort { input, .. } => substitute_plan(input, literals, idx),
157 PlanNode::AlterTable { .. } => {}
158 PlanNode::DropTable { .. } => {}
159 PlanNode::Limit { input, count } => {
160 if let PlanNode::Offset {
169 input: inner,
170 count: off_count,
171 } = input.as_mut()
172 {
173 substitute_plan(inner, literals, idx);
174 substitute_expr(count, literals, idx);
175 substitute_expr(off_count, literals, idx);
176 } else {
177 substitute_plan(input, literals, idx);
178 substitute_expr(count, literals, idx);
179 }
180 }
181 PlanNode::Offset { input, count } => {
182 substitute_plan(input, literals, idx);
185 substitute_expr(count, literals, idx);
186 }
187 PlanNode::Aggregate { input, .. } => {
188 substitute_plan(input, literals, idx);
189 }
190 PlanNode::NestedLoopJoin {
191 left, right, on, ..
192 } => {
193 substitute_plan(left, literals, idx);
198 substitute_plan(right, literals, idx);
199 if let Some(pred) = on {
200 substitute_expr(pred, literals, idx);
201 }
202 }
203 PlanNode::Distinct { input } => {
204 substitute_plan(input, literals, idx);
205 }
206 PlanNode::GroupBy { input, having, .. } => {
207 substitute_plan(input, literals, idx);
208 if let Some(pred) = having {
209 substitute_expr(pred, literals, idx);
210 }
211 }
212 PlanNode::Insert { assignments, .. } => {
213 substitute_assignments(assignments, literals, idx);
214 }
215 PlanNode::Upsert {
216 assignments,
217 on_conflict,
218 ..
219 } => {
220 substitute_assignments(assignments, literals, idx);
221 substitute_assignments(on_conflict, literals, idx);
222 }
223 PlanNode::Update {
224 input, assignments, ..
225 } => {
226 substitute_plan(input, literals, idx);
227 substitute_assignments(assignments, literals, idx);
228 }
229 PlanNode::Delete { input, .. } => {
230 substitute_plan(input, literals, idx);
231 }
232 PlanNode::CreateTable { .. } => {}
233 PlanNode::CreateView { .. } => {}
234 PlanNode::RefreshView { .. } => {}
235 PlanNode::DropView { .. } => {}
236 PlanNode::Window { input, windows } => {
237 substitute_plan(input, literals, idx);
238 for w in windows {
239 for arg in &mut w.args {
240 substitute_expr(arg, literals, idx);
241 }
242 }
243 }
244 PlanNode::Union { left, right, .. } => {
245 substitute_plan(left, literals, idx);
246 substitute_plan(right, literals, idx);
247 }
248 PlanNode::Explain { input } => {
249 substitute_plan(input, literals, idx);
250 }
251 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
252 }
253}
254
255fn substitute_assignments(assignments: &mut [Assignment], literals: &[Literal], idx: &mut usize) {
256 for a in assignments {
257 substitute_expr(&mut a.value, literals, idx);
258 }
259}
260
261pub(crate) fn count_literal_slots(plan: &PlanNode) -> usize {
267 let mut n = 0usize;
268 count_plan(plan, &mut n);
269 n
270}
271
272fn count_plan(plan: &PlanNode, n: &mut usize) {
273 match plan {
274 PlanNode::SeqScan { .. } => {}
275 PlanNode::AliasScan { .. } => {}
276 PlanNode::IndexScan { key, .. } => count_expr(key, n),
277 PlanNode::RangeScan { start, end, .. } => {
278 if let Some((expr, _)) = start {
279 count_expr(expr, n);
280 }
281 if let Some((expr, _)) = end {
282 count_expr(expr, n);
283 }
284 }
285 PlanNode::Filter { input, predicate } => {
286 count_plan(input, n);
287 count_expr(predicate, n);
288 }
289 PlanNode::Project { input, fields } => {
290 count_plan(input, n);
291 for f in fields {
292 count_expr(&f.expr, n);
293 }
294 }
295 PlanNode::Sort { input, .. } => count_plan(input, n),
296 PlanNode::Limit { input, count } => {
297 if let PlanNode::Offset {
302 input: inner,
303 count: off_count,
304 } = input.as_ref()
305 {
306 count_plan(inner, n);
307 count_expr(count, n);
308 count_expr(off_count, n);
309 } else {
310 count_plan(input, n);
311 count_expr(count, n);
312 }
313 }
314 PlanNode::Offset { input, count } => {
315 count_plan(input, n);
316 count_expr(count, n);
317 }
318 PlanNode::Aggregate { input, .. } => count_plan(input, n),
319 PlanNode::NestedLoopJoin {
320 left, right, on, ..
321 } => {
322 count_plan(left, n);
323 count_plan(right, n);
324 if let Some(pred) = on {
325 count_expr(pred, n);
326 }
327 }
328 PlanNode::Distinct { input } => count_plan(input, n),
329 PlanNode::GroupBy { input, having, .. } => {
330 count_plan(input, n);
331 if let Some(pred) = having {
332 count_expr(pred, n);
333 }
334 }
335 PlanNode::Insert { assignments, .. } => {
336 for a in assignments {
337 count_expr(&a.value, n);
338 }
339 }
340 PlanNode::Upsert {
341 assignments,
342 on_conflict,
343 ..
344 } => {
345 for a in assignments {
346 count_expr(&a.value, n);
347 }
348 for a in on_conflict {
349 count_expr(&a.value, n);
350 }
351 }
352 PlanNode::Update {
353 input, assignments, ..
354 } => {
355 count_plan(input, n);
356 for a in assignments {
357 count_expr(&a.value, n);
358 }
359 }
360 PlanNode::Delete { input, .. } => count_plan(input, n),
361 PlanNode::CreateTable { .. } => {}
362 PlanNode::AlterTable { .. } => {}
363 PlanNode::DropTable { .. } => {}
364 PlanNode::CreateView { .. } => {}
365 PlanNode::RefreshView { .. } => {}
366 PlanNode::DropView { .. } => {}
367 PlanNode::Window { input, windows } => {
368 count_plan(input, n);
369 for w in windows {
370 for arg in &w.args {
371 count_expr(arg, n);
372 }
373 }
374 }
375 PlanNode::Union { left, right, .. } => {
376 count_plan(left, n);
377 count_plan(right, n);
378 }
379 PlanNode::Explain { input } => {
380 count_plan(input, n);
381 }
382 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
383 }
384}
385
386fn count_expr(expr: &Expr, n: &mut usize) {
387 match expr {
388 Expr::Literal(_) => *n += 1,
389 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
390 Expr::BinaryOp(l, _, r) => {
391 count_expr(l, n);
392 count_expr(r, n);
393 }
394 Expr::UnaryOp(_, inner) => count_expr(inner, n),
395 Expr::FunctionCall(_, inner) => count_expr(inner, n),
396 Expr::Coalesce(l, r) => {
397 count_expr(l, n);
398 count_expr(r, n);
399 }
400 Expr::InList { expr, list, .. } => {
401 count_expr(expr, n);
402 for item in list {
403 count_expr(item, n);
404 }
405 }
406 Expr::ScalarFunc(_, args) => {
407 for a in args {
408 count_expr(a, n);
409 }
410 }
411 Expr::Cast(inner, _) => count_expr(inner, n),
412 Expr::Case { whens, else_expr } => {
413 for (cond, result) in whens {
414 count_expr(cond, n);
415 count_expr(result, n);
416 }
417 if let Some(e) = else_expr {
418 count_expr(e, n);
419 }
420 }
421 Expr::InSubquery { expr, .. } => {
422 count_expr(expr, n);
423 }
426 Expr::ExistsSubquery { .. } => {
427 }
430 Expr::Window { args, .. } => {
431 for a in args {
432 count_expr(a, n);
433 }
434 }
435 Expr::Null => {}
436 }
437}
438
439fn substitute_expr(expr: &mut Expr, literals: &[Literal], idx: &mut usize) {
440 match expr {
441 Expr::Literal(_) => {
442 *expr = Expr::Literal(literals[*idx].clone());
446 *idx += 1;
447 }
448 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
449 Expr::BinaryOp(l, _, r) => {
450 substitute_expr(l, literals, idx);
451 substitute_expr(r, literals, idx);
452 }
453 Expr::UnaryOp(_, inner) => {
454 substitute_expr(inner, literals, idx);
455 }
456 Expr::FunctionCall(_, inner) => {
457 substitute_expr(inner, literals, idx);
458 }
459 Expr::Coalesce(l, r) => {
460 substitute_expr(l, literals, idx);
461 substitute_expr(r, literals, idx);
462 }
463 Expr::InList { expr, list, .. } => {
464 substitute_expr(expr, literals, idx);
465 for item in list {
466 substitute_expr(item, literals, idx);
467 }
468 }
469 Expr::ScalarFunc(_, args) => {
470 for a in args {
471 substitute_expr(a, literals, idx);
472 }
473 }
474 Expr::Cast(inner, _) => substitute_expr(inner, literals, idx),
475 Expr::Case { whens, else_expr } => {
476 for (cond, result) in whens {
477 substitute_expr(cond, literals, idx);
478 substitute_expr(result, literals, idx);
479 }
480 if let Some(e) = else_expr {
481 substitute_expr(e, literals, idx);
482 }
483 }
484 Expr::InSubquery { expr, .. } => {
485 substitute_expr(expr, literals, idx);
486 }
487 Expr::ExistsSubquery { .. } => {
488 }
491 Expr::Window { args, .. } => {
492 for a in args {
493 substitute_expr(a, literals, idx);
494 }
495 }
496 Expr::Null => {}
497 }
498}
499
500#[cfg(test)]
501mod tests {
502 use super::*;
503 use crate::canonicalize::canonicalize;
504 use crate::planner;
505
506 #[test]
507 fn test_cache_hit_substitutes_literal() {
508 let mut cache = PlanCache::new(100);
509
510 let q1 = "User filter .id = 42";
512 let (h1, lits1) = canonicalize(q1).unwrap();
513 let p1 = planner::plan(q1).unwrap();
514 cache.insert(h1, p1);
515
516 let q2 = "User filter .id = 99";
519 let (h2, lits2) = canonicalize(q2).unwrap();
520 assert_eq!(h1, h2, "different literals must hash the same");
521
522 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
523
524 match plan {
526 PlanNode::IndexScan { key, .. } => {
527 assert_eq!(key, Expr::Literal(Literal::Int(99)));
528 }
529 other => panic!("expected IndexScan, got {other:?}"),
530 }
531
532 assert_eq!(lits1, vec![Literal::Int(42)]);
535 assert_eq!(cache.hits, 1);
536 assert_eq!(cache.misses, 0);
537 }
538
539 #[test]
540 fn test_cache_miss_returns_none_and_bumps_counter() {
541 let mut cache = PlanCache::new(100);
542 assert!(cache.get_with_substitution(99999, &[]).is_none());
543 assert_eq!(cache.misses, 1);
544 assert_eq!(cache.hits, 0);
545 }
546
547 #[test]
548 fn test_multi_literal_filter_substitution() {
549 let mut cache = PlanCache::new(100);
550 let q1 = r#"User filter .age > 30 and .status = "active" { .name }"#;
551 let (h1, _) = canonicalize(q1).unwrap();
552 cache.insert(h1, planner::plan(q1).unwrap());
553
554 let q2 = r#"User filter .age > 50 and .status = "pending" { .name }"#;
555 let (h2, lits2) = canonicalize(q2).unwrap();
556 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
557
558 let mut found = Vec::new();
560 collect_literals_for_test(&plan, &mut found);
561 assert_eq!(
562 found,
563 vec![Literal::Int(50), Literal::String("pending".into()),]
564 );
565 }
566
567 #[test]
568 fn test_update_by_pk_substitution() {
569 let mut cache = PlanCache::new(100);
570 let q1 = "User filter .id = 1 update { age := 100 }";
571 let (h1, _) = canonicalize(q1).unwrap();
572 cache.insert(h1, planner::plan(q1).unwrap());
573
574 let q2 = "User filter .id = 7 update { age := 200 }";
575 let (h2, lits2) = canonicalize(q2).unwrap();
576 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
577
578 let mut found = Vec::new();
579 collect_literals_for_test(&plan, &mut found);
580 assert_eq!(found, vec![Literal::Int(7), Literal::Int(200)]);
581 }
582
583 #[test]
584 fn test_insert_substitution() {
585 let mut cache = PlanCache::new(100);
586 let q1 = r#"insert User { id := 1, name := "Alice", age := 20 }"#;
587 let (h1, _) = canonicalize(q1).unwrap();
588 cache.insert(h1, planner::plan(q1).unwrap());
589
590 let q2 = r#"insert User { id := 2, name := "Bob", age := 30 }"#;
591 let (h2, lits2) = canonicalize(q2).unwrap();
592 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
593
594 let mut found = Vec::new();
595 collect_literals_for_test(&plan, &mut found);
596 assert_eq!(
597 found,
598 vec![
599 Literal::Int(2),
600 Literal::String("Bob".into()),
601 Literal::Int(30),
602 ]
603 );
604 }
605
606 #[test]
607 fn test_eviction_on_capacity() {
608 let mut cache = PlanCache::new(2);
609 let q1 = "User";
610 let q2 = "User filter .age > 1";
611 let _q3 = "User filter .age > 2";
612 let q3_distinct = "User filter .id = 5";
615
616 let (h1, _) = canonicalize(q1).unwrap();
617 let (h2, _) = canonicalize(q2).unwrap();
618 let (h3, _) = canonicalize(q3_distinct).unwrap();
619 cache.insert(h1, planner::plan(q1).unwrap());
620 cache.insert(h2, planner::plan(q2).unwrap());
621 cache.insert(h3, planner::plan(q3_distinct).unwrap());
623 assert!(cache.cache.contains_key(&h3));
624 assert_eq!(cache.cache.len(), 1);
625 }
626
627 fn collect_literals_for_test(plan: &PlanNode, out: &mut Vec<Literal>) {
631 match plan {
632 PlanNode::SeqScan { .. } => {}
633 PlanNode::AliasScan { .. } => {}
634 PlanNode::IndexScan { key, .. } => collect_expr_literals(key, out),
635 PlanNode::RangeScan { start, end, .. } => {
636 if let Some((expr, _)) = start {
637 collect_expr_literals(expr, out);
638 }
639 if let Some((expr, _)) = end {
640 collect_expr_literals(expr, out);
641 }
642 }
643 PlanNode::Filter { input, predicate } => {
644 collect_literals_for_test(input, out);
645 collect_expr_literals(predicate, out);
646 }
647 PlanNode::Project { input, fields } => {
648 collect_literals_for_test(input, out);
649 for f in fields {
650 collect_expr_literals(&f.expr, out);
651 }
652 }
653 PlanNode::Sort { input, .. } => collect_literals_for_test(input, out),
654 PlanNode::Limit { input, count } => {
655 collect_literals_for_test(input, out);
656 collect_expr_literals(count, out);
657 }
658 PlanNode::Offset { input, count } => {
659 collect_literals_for_test(input, out);
660 collect_expr_literals(count, out);
661 }
662 PlanNode::Aggregate { input, .. } => collect_literals_for_test(input, out),
663 PlanNode::NestedLoopJoin {
664 left, right, on, ..
665 } => {
666 collect_literals_for_test(left, out);
667 collect_literals_for_test(right, out);
668 if let Some(pred) = on {
669 collect_expr_literals(pred, out);
670 }
671 }
672 PlanNode::Insert { assignments, .. } => {
673 for a in assignments {
674 collect_expr_literals(&a.value, out);
675 }
676 }
677 PlanNode::Upsert {
678 assignments,
679 on_conflict,
680 ..
681 } => {
682 for a in assignments {
683 collect_expr_literals(&a.value, out);
684 }
685 for a in on_conflict {
686 collect_expr_literals(&a.value, out);
687 }
688 }
689 PlanNode::Update {
690 input, assignments, ..
691 } => {
692 collect_literals_for_test(input, out);
693 for a in assignments {
694 collect_expr_literals(&a.value, out);
695 }
696 }
697 PlanNode::Distinct { input } => collect_literals_for_test(input, out),
698 PlanNode::GroupBy { input, having, .. } => {
699 collect_literals_for_test(input, out);
700 if let Some(pred) = having {
701 collect_expr_literals(pred, out);
702 }
703 }
704 PlanNode::Delete { input, .. } => collect_literals_for_test(input, out),
705 PlanNode::CreateTable { .. } => {}
706 PlanNode::AlterTable { .. } => {}
707 PlanNode::DropTable { .. } => {}
708 PlanNode::CreateView { .. } => {}
709 PlanNode::RefreshView { .. } => {}
710 PlanNode::DropView { .. } => {}
711 PlanNode::Window { input, windows } => {
712 collect_literals_for_test(input, out);
713 for w in windows {
714 for arg in &w.args {
715 collect_expr_literals(arg, out);
716 }
717 }
718 }
719 PlanNode::Union { left, right, .. } => {
720 collect_literals_for_test(left, out);
721 collect_literals_for_test(right, out);
722 }
723 PlanNode::Explain { input } => {
724 collect_literals_for_test(input, out);
725 }
726 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
727 }
728 }
729
730 fn collect_expr_literals(expr: &Expr, out: &mut Vec<Literal>) {
731 match expr {
732 Expr::Literal(l) => out.push(l.clone()),
733 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
734 Expr::BinaryOp(l, _, r) => {
735 collect_expr_literals(l, out);
736 collect_expr_literals(r, out);
737 }
738 Expr::UnaryOp(_, inner) => collect_expr_literals(inner, out),
739 Expr::FunctionCall(_, inner) => collect_expr_literals(inner, out),
740 Expr::Coalesce(l, r) => {
741 collect_expr_literals(l, out);
742 collect_expr_literals(r, out);
743 }
744 Expr::InList { expr, list, .. } => {
745 collect_expr_literals(expr, out);
746 for item in list {
747 collect_expr_literals(item, out);
748 }
749 }
750 Expr::ScalarFunc(_, args) => {
751 for a in args {
752 collect_expr_literals(a, out);
753 }
754 }
755 Expr::Cast(inner, _) => collect_expr_literals(inner, out),
756 Expr::Case { whens, else_expr } => {
757 for (cond, result) in whens {
758 collect_expr_literals(cond, out);
759 collect_expr_literals(result, out);
760 }
761 if let Some(e) = else_expr {
762 collect_expr_literals(e, out);
763 }
764 }
765 Expr::InSubquery { expr, .. } => {
766 collect_expr_literals(expr, out);
767 }
768 Expr::ExistsSubquery { .. } => {}
769 Expr::Window { args, .. } => {
770 for a in args {
771 collect_expr_literals(a, out);
772 }
773 }
774 Expr::Null => {}
775 }
776 }
777}