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 { rows, .. } => {
213 for assignments in rows {
214 substitute_assignments(assignments, literals, idx);
215 }
216 }
217 PlanNode::Upsert {
218 assignments,
219 on_conflict,
220 ..
221 } => {
222 substitute_assignments(assignments, literals, idx);
223 substitute_assignments(on_conflict, literals, idx);
224 }
225 PlanNode::Update {
226 input, assignments, ..
227 } => {
228 substitute_plan(input, literals, idx);
229 substitute_assignments(assignments, literals, idx);
230 }
231 PlanNode::Delete { input, .. } => {
232 substitute_plan(input, literals, idx);
233 }
234 PlanNode::CreateTable { .. } => {}
235 PlanNode::CreateView { .. } => {}
236 PlanNode::RefreshView { .. } => {}
237 PlanNode::DropView { .. } => {}
238 PlanNode::Window { input, windows } => {
239 substitute_plan(input, literals, idx);
240 for w in windows {
241 for arg in &mut w.args {
242 substitute_expr(arg, literals, idx);
243 }
244 }
245 }
246 PlanNode::Union { left, right, .. } => {
247 substitute_plan(left, literals, idx);
248 substitute_plan(right, literals, idx);
249 }
250 PlanNode::Explain { input } => {
251 substitute_plan(input, literals, idx);
252 }
253 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
254 }
255}
256
257fn substitute_assignments(assignments: &mut [Assignment], literals: &[Literal], idx: &mut usize) {
258 for a in assignments {
259 substitute_expr(&mut a.value, literals, idx);
260 }
261}
262
263pub(crate) fn count_literal_slots(plan: &PlanNode) -> usize {
269 let mut n = 0usize;
270 count_plan(plan, &mut n);
271 n
272}
273
274fn count_plan(plan: &PlanNode, n: &mut usize) {
275 match plan {
276 PlanNode::SeqScan { .. } => {}
277 PlanNode::AliasScan { .. } => {}
278 PlanNode::IndexScan { key, .. } => count_expr(key, n),
279 PlanNode::RangeScan { start, end, .. } => {
280 if let Some((expr, _)) = start {
281 count_expr(expr, n);
282 }
283 if let Some((expr, _)) = end {
284 count_expr(expr, n);
285 }
286 }
287 PlanNode::Filter { input, predicate } => {
288 count_plan(input, n);
289 count_expr(predicate, n);
290 }
291 PlanNode::Project { input, fields } => {
292 count_plan(input, n);
293 for f in fields {
294 count_expr(&f.expr, n);
295 }
296 }
297 PlanNode::Sort { input, .. } => count_plan(input, n),
298 PlanNode::Limit { input, count } => {
299 if let PlanNode::Offset {
304 input: inner,
305 count: off_count,
306 } = input.as_ref()
307 {
308 count_plan(inner, n);
309 count_expr(count, n);
310 count_expr(off_count, n);
311 } else {
312 count_plan(input, n);
313 count_expr(count, n);
314 }
315 }
316 PlanNode::Offset { input, count } => {
317 count_plan(input, n);
318 count_expr(count, n);
319 }
320 PlanNode::Aggregate { input, .. } => count_plan(input, n),
321 PlanNode::NestedLoopJoin {
322 left, right, on, ..
323 } => {
324 count_plan(left, n);
325 count_plan(right, n);
326 if let Some(pred) = on {
327 count_expr(pred, n);
328 }
329 }
330 PlanNode::Distinct { input } => count_plan(input, n),
331 PlanNode::GroupBy { input, having, .. } => {
332 count_plan(input, n);
333 if let Some(pred) = having {
334 count_expr(pred, n);
335 }
336 }
337 PlanNode::Insert { rows, .. } => {
338 for assignments in rows {
339 for a in assignments {
340 count_expr(&a.value, n);
341 }
342 }
343 }
344 PlanNode::Upsert {
345 assignments,
346 on_conflict,
347 ..
348 } => {
349 for a in assignments {
350 count_expr(&a.value, n);
351 }
352 for a in on_conflict {
353 count_expr(&a.value, n);
354 }
355 }
356 PlanNode::Update {
357 input, assignments, ..
358 } => {
359 count_plan(input, n);
360 for a in assignments {
361 count_expr(&a.value, n);
362 }
363 }
364 PlanNode::Delete { input, .. } => count_plan(input, n),
365 PlanNode::CreateTable { .. } => {}
366 PlanNode::AlterTable { .. } => {}
367 PlanNode::DropTable { .. } => {}
368 PlanNode::CreateView { .. } => {}
369 PlanNode::RefreshView { .. } => {}
370 PlanNode::DropView { .. } => {}
371 PlanNode::Window { input, windows } => {
372 count_plan(input, n);
373 for w in windows {
374 for arg in &w.args {
375 count_expr(arg, n);
376 }
377 }
378 }
379 PlanNode::Union { left, right, .. } => {
380 count_plan(left, n);
381 count_plan(right, n);
382 }
383 PlanNode::Explain { input } => {
384 count_plan(input, n);
385 }
386 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
387 }
388}
389
390fn count_expr(expr: &Expr, n: &mut usize) {
391 match expr {
392 Expr::Literal(_) => *n += 1,
393 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
394 Expr::BinaryOp(l, _, r) => {
395 count_expr(l, n);
396 count_expr(r, n);
397 }
398 Expr::UnaryOp(_, inner) => count_expr(inner, n),
399 Expr::FunctionCall(_, inner) => count_expr(inner, n),
400 Expr::Coalesce(l, r) => {
401 count_expr(l, n);
402 count_expr(r, n);
403 }
404 Expr::InList { expr, list, .. } => {
405 count_expr(expr, n);
406 for item in list {
407 count_expr(item, n);
408 }
409 }
410 Expr::ScalarFunc(_, args) => {
411 for a in args {
412 count_expr(a, n);
413 }
414 }
415 Expr::Cast(inner, _) => count_expr(inner, n),
416 Expr::Case { whens, else_expr } => {
417 for (cond, result) in whens {
418 count_expr(cond, n);
419 count_expr(result, n);
420 }
421 if let Some(e) = else_expr {
422 count_expr(e, n);
423 }
424 }
425 Expr::InSubquery { expr, .. } => {
426 count_expr(expr, n);
427 }
430 Expr::ExistsSubquery { .. } => {
431 }
434 Expr::Window { args, .. } => {
435 for a in args {
436 count_expr(a, n);
437 }
438 }
439 Expr::ValueLit(_) => {}
442 Expr::Null => {}
443 }
444}
445
446fn substitute_expr(expr: &mut Expr, literals: &[Literal], idx: &mut usize) {
447 match expr {
448 Expr::Literal(_) => {
449 *expr = Expr::Literal(literals[*idx].clone());
453 *idx += 1;
454 }
455 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
456 Expr::BinaryOp(l, _, r) => {
457 substitute_expr(l, literals, idx);
458 substitute_expr(r, literals, idx);
459 }
460 Expr::UnaryOp(_, inner) => {
461 substitute_expr(inner, literals, idx);
462 }
463 Expr::FunctionCall(_, inner) => {
464 substitute_expr(inner, literals, idx);
465 }
466 Expr::Coalesce(l, r) => {
467 substitute_expr(l, literals, idx);
468 substitute_expr(r, literals, idx);
469 }
470 Expr::InList { expr, list, .. } => {
471 substitute_expr(expr, literals, idx);
472 for item in list {
473 substitute_expr(item, literals, idx);
474 }
475 }
476 Expr::ScalarFunc(_, args) => {
477 for a in args {
478 substitute_expr(a, literals, idx);
479 }
480 }
481 Expr::Cast(inner, _) => substitute_expr(inner, literals, idx),
482 Expr::Case { whens, else_expr } => {
483 for (cond, result) in whens {
484 substitute_expr(cond, literals, idx);
485 substitute_expr(result, literals, idx);
486 }
487 if let Some(e) = else_expr {
488 substitute_expr(e, literals, idx);
489 }
490 }
491 Expr::InSubquery { expr, .. } => {
492 substitute_expr(expr, literals, idx);
493 }
494 Expr::ExistsSubquery { .. } => {
495 }
498 Expr::Window { args, .. } => {
499 for a in args {
500 substitute_expr(a, literals, idx);
501 }
502 }
503 Expr::ValueLit(_) => {}
506 Expr::Null => {}
507 }
508}
509
510#[cfg(test)]
511mod tests {
512 use super::*;
513 use crate::canonicalize::canonicalize;
514 use crate::planner;
515
516 #[test]
517 fn test_cache_hit_substitutes_literal() {
518 let mut cache = PlanCache::new(100);
519
520 let q1 = "User filter .id = 42";
522 let (h1, lits1) = canonicalize(q1).unwrap();
523 let p1 = planner::plan(q1).unwrap();
524 cache.insert(h1, p1);
525
526 let q2 = "User filter .id = 99";
529 let (h2, lits2) = canonicalize(q2).unwrap();
530 assert_eq!(h1, h2, "different literals must hash the same");
531
532 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
533
534 match plan {
536 PlanNode::IndexScan { key, .. } => {
537 assert_eq!(key, Expr::Literal(Literal::Int(99)));
538 }
539 other => panic!("expected IndexScan, got {other:?}"),
540 }
541
542 assert_eq!(lits1, vec![Literal::Int(42)]);
545 assert_eq!(cache.hits, 1);
546 assert_eq!(cache.misses, 0);
547 }
548
549 #[test]
550 fn test_cache_miss_returns_none_and_bumps_counter() {
551 let mut cache = PlanCache::new(100);
552 assert!(cache.get_with_substitution(99999, &[]).is_none());
553 assert_eq!(cache.misses, 1);
554 assert_eq!(cache.hits, 0);
555 }
556
557 #[test]
558 fn test_multi_literal_filter_substitution() {
559 let mut cache = PlanCache::new(100);
560 let q1 = r#"User filter .age > 30 and .status = "active" { .name }"#;
561 let (h1, _) = canonicalize(q1).unwrap();
562 cache.insert(h1, planner::plan(q1).unwrap());
563
564 let q2 = r#"User filter .age > 50 and .status = "pending" { .name }"#;
565 let (h2, lits2) = canonicalize(q2).unwrap();
566 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
567
568 let mut found = Vec::new();
570 collect_literals_for_test(&plan, &mut found);
571 assert_eq!(
572 found,
573 vec![Literal::Int(50), Literal::String("pending".into()),]
574 );
575 }
576
577 #[test]
578 fn test_update_by_pk_substitution() {
579 let mut cache = PlanCache::new(100);
580 let q1 = "User filter .id = 1 update { age := 100 }";
581 let (h1, _) = canonicalize(q1).unwrap();
582 cache.insert(h1, planner::plan(q1).unwrap());
583
584 let q2 = "User filter .id = 7 update { age := 200 }";
585 let (h2, lits2) = canonicalize(q2).unwrap();
586 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
587
588 let mut found = Vec::new();
589 collect_literals_for_test(&plan, &mut found);
590 assert_eq!(found, vec![Literal::Int(7), Literal::Int(200)]);
591 }
592
593 #[test]
594 fn test_insert_substitution() {
595 let mut cache = PlanCache::new(100);
596 let q1 = r#"insert User { id := 1, name := "Alice", age := 20 }"#;
597 let (h1, _) = canonicalize(q1).unwrap();
598 cache.insert(h1, planner::plan(q1).unwrap());
599
600 let q2 = r#"insert User { id := 2, name := "Bob", age := 30 }"#;
601 let (h2, lits2) = canonicalize(q2).unwrap();
602 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
603
604 let mut found = Vec::new();
605 collect_literals_for_test(&plan, &mut found);
606 assert_eq!(
607 found,
608 vec![
609 Literal::Int(2),
610 Literal::String("Bob".into()),
611 Literal::Int(30),
612 ]
613 );
614 }
615
616 #[test]
617 fn test_eviction_on_capacity() {
618 let mut cache = PlanCache::new(2);
619 let q1 = "User";
620 let q2 = "User filter .age > 1";
621 let _q3 = "User filter .age > 2";
622 let q3_distinct = "User filter .id = 5";
625
626 let (h1, _) = canonicalize(q1).unwrap();
627 let (h2, _) = canonicalize(q2).unwrap();
628 let (h3, _) = canonicalize(q3_distinct).unwrap();
629 cache.insert(h1, planner::plan(q1).unwrap());
630 cache.insert(h2, planner::plan(q2).unwrap());
631 cache.insert(h3, planner::plan(q3_distinct).unwrap());
633 assert!(cache.cache.contains_key(&h3));
634 assert_eq!(cache.cache.len(), 1);
635 }
636
637 fn collect_literals_for_test(plan: &PlanNode, out: &mut Vec<Literal>) {
641 match plan {
642 PlanNode::SeqScan { .. } => {}
643 PlanNode::AliasScan { .. } => {}
644 PlanNode::IndexScan { key, .. } => collect_expr_literals(key, out),
645 PlanNode::RangeScan { start, end, .. } => {
646 if let Some((expr, _)) = start {
647 collect_expr_literals(expr, out);
648 }
649 if let Some((expr, _)) = end {
650 collect_expr_literals(expr, out);
651 }
652 }
653 PlanNode::Filter { input, predicate } => {
654 collect_literals_for_test(input, out);
655 collect_expr_literals(predicate, out);
656 }
657 PlanNode::Project { input, fields } => {
658 collect_literals_for_test(input, out);
659 for f in fields {
660 collect_expr_literals(&f.expr, out);
661 }
662 }
663 PlanNode::Sort { input, .. } => collect_literals_for_test(input, out),
664 PlanNode::Limit { input, count } => {
665 collect_literals_for_test(input, out);
666 collect_expr_literals(count, out);
667 }
668 PlanNode::Offset { input, count } => {
669 collect_literals_for_test(input, out);
670 collect_expr_literals(count, out);
671 }
672 PlanNode::Aggregate { input, .. } => collect_literals_for_test(input, out),
673 PlanNode::NestedLoopJoin {
674 left, right, on, ..
675 } => {
676 collect_literals_for_test(left, out);
677 collect_literals_for_test(right, out);
678 if let Some(pred) = on {
679 collect_expr_literals(pred, out);
680 }
681 }
682 PlanNode::Insert { rows, .. } => {
683 for assignments in rows {
684 for a in assignments {
685 collect_expr_literals(&a.value, out);
686 }
687 }
688 }
689 PlanNode::Upsert {
690 assignments,
691 on_conflict,
692 ..
693 } => {
694 for a in assignments {
695 collect_expr_literals(&a.value, out);
696 }
697 for a in on_conflict {
698 collect_expr_literals(&a.value, out);
699 }
700 }
701 PlanNode::Update {
702 input, assignments, ..
703 } => {
704 collect_literals_for_test(input, out);
705 for a in assignments {
706 collect_expr_literals(&a.value, out);
707 }
708 }
709 PlanNode::Distinct { input } => collect_literals_for_test(input, out),
710 PlanNode::GroupBy { input, having, .. } => {
711 collect_literals_for_test(input, out);
712 if let Some(pred) = having {
713 collect_expr_literals(pred, out);
714 }
715 }
716 PlanNode::Delete { input, .. } => collect_literals_for_test(input, out),
717 PlanNode::CreateTable { .. } => {}
718 PlanNode::AlterTable { .. } => {}
719 PlanNode::DropTable { .. } => {}
720 PlanNode::CreateView { .. } => {}
721 PlanNode::RefreshView { .. } => {}
722 PlanNode::DropView { .. } => {}
723 PlanNode::Window { input, windows } => {
724 collect_literals_for_test(input, out);
725 for w in windows {
726 for arg in &w.args {
727 collect_expr_literals(arg, out);
728 }
729 }
730 }
731 PlanNode::Union { left, right, .. } => {
732 collect_literals_for_test(left, out);
733 collect_literals_for_test(right, out);
734 }
735 PlanNode::Explain { input } => {
736 collect_literals_for_test(input, out);
737 }
738 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
739 }
740 }
741
742 fn collect_expr_literals(expr: &Expr, out: &mut Vec<Literal>) {
743 match expr {
744 Expr::Literal(l) => out.push(l.clone()),
745 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
746 Expr::BinaryOp(l, _, r) => {
747 collect_expr_literals(l, out);
748 collect_expr_literals(r, out);
749 }
750 Expr::UnaryOp(_, inner) => collect_expr_literals(inner, out),
751 Expr::FunctionCall(_, inner) => collect_expr_literals(inner, out),
752 Expr::Coalesce(l, r) => {
753 collect_expr_literals(l, out);
754 collect_expr_literals(r, out);
755 }
756 Expr::InList { expr, list, .. } => {
757 collect_expr_literals(expr, out);
758 for item in list {
759 collect_expr_literals(item, out);
760 }
761 }
762 Expr::ScalarFunc(_, args) => {
763 for a in args {
764 collect_expr_literals(a, out);
765 }
766 }
767 Expr::Cast(inner, _) => collect_expr_literals(inner, out),
768 Expr::Case { whens, else_expr } => {
769 for (cond, result) in whens {
770 collect_expr_literals(cond, out);
771 collect_expr_literals(result, out);
772 }
773 if let Some(e) = else_expr {
774 collect_expr_literals(e, out);
775 }
776 }
777 Expr::InSubquery { expr, .. } => {
778 collect_expr_literals(expr, out);
779 }
780 Expr::ExistsSubquery { .. } => {}
781 Expr::Window { args, .. } => {
782 for a in args {
783 collect_expr_literals(a, out);
784 }
785 }
786 Expr::ValueLit(_) => {}
787 Expr::Null => {}
788 }
789 }
790}