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, source_literal_count: usize) {
72 if count_literal_slots(&plan) != source_literal_count {
73 return;
74 }
75 if self.cache.len() >= self.capacity && !self.cache.contains_key(&hash) {
76 self.cache.clear();
81 }
82 self.cache.insert(hash, plan);
83 }
84
85 pub fn get_with_substitution(&mut self, hash: u64, literals: &[Literal]) -> Option<PlanNode> {
97 match self.cache.get(&hash) {
98 Some(template) => {
99 self.hits += 1;
100 let mut plan = template.clone();
101 let mut idx = 0usize;
102 substitute_plan(&mut plan, literals, &mut idx);
103 debug_assert_eq!(
104 idx,
105 literals.len(),
106 "plan substitution consumed {idx} literals but query had {}",
107 literals.len(),
108 );
109 Some(plan)
110 }
111 None => {
112 self.misses += 1;
113 None
114 }
115 }
116 }
117
118 pub fn len(&self) -> usize {
119 self.cache.len()
120 }
121
122 pub fn is_empty(&self) -> bool {
123 self.cache.is_empty()
124 }
125
126 pub fn clear(&mut self) {
127 self.cache.clear();
128 }
129}
130
131pub(crate) fn substitute_plan(plan: &mut PlanNode, literals: &[Literal], idx: &mut usize) {
150 match plan {
151 PlanNode::SeqScan { .. } => {}
152 PlanNode::AliasScan { .. } => {}
153 PlanNode::IndexScan { key, .. } => {
154 substitute_expr(key, literals, idx);
155 }
156 PlanNode::RangeScan { start, end, .. } => {
157 if let Some((expr, _)) = start {
158 substitute_expr(expr, literals, idx);
159 }
160 if let Some((expr, _)) = end {
161 substitute_expr(expr, literals, idx);
162 }
163 }
164 PlanNode::Filter { input, predicate } => {
165 substitute_plan(input, literals, idx);
166 substitute_expr(predicate, literals, idx);
167 }
168 PlanNode::Project { input, fields } => {
169 substitute_plan(input, literals, idx);
170 for f in fields {
171 substitute_expr(&mut f.expr, literals, idx);
172 }
173 }
174 PlanNode::Sort { input, .. } => substitute_plan(input, literals, idx),
175 PlanNode::AlterTable { .. } => {}
176 PlanNode::DropTable { .. } => {}
177 PlanNode::Limit { input, count } => {
178 if let PlanNode::Offset {
187 input: inner,
188 count: off_count,
189 } = input.as_mut()
190 {
191 substitute_plan(inner, literals, idx);
192 substitute_expr(count, literals, idx);
193 substitute_expr(off_count, literals, idx);
194 } else {
195 substitute_plan(input, literals, idx);
196 substitute_expr(count, literals, idx);
197 }
198 }
199 PlanNode::Offset { input, count } => {
200 substitute_plan(input, literals, idx);
203 substitute_expr(count, literals, idx);
204 }
205 PlanNode::Aggregate { input, .. } => {
206 substitute_plan(input, literals, idx);
207 }
208 PlanNode::NestedLoopJoin {
209 left, right, on, ..
210 } => {
211 substitute_plan(left, literals, idx);
216 substitute_plan(right, literals, idx);
217 if let Some(pred) = on {
218 substitute_expr(pred, literals, idx);
219 }
220 }
221 PlanNode::Distinct { input } => {
222 substitute_plan(input, literals, idx);
223 }
224 PlanNode::GroupBy { input, having, .. } => {
225 substitute_plan(input, literals, idx);
226 if let Some(pred) = having {
227 substitute_expr(pred, literals, idx);
228 }
229 }
230 PlanNode::Insert { rows, .. } => {
231 for assignments in rows {
232 substitute_assignments(assignments, literals, idx);
233 }
234 }
235 PlanNode::Upsert {
236 assignments,
237 on_conflict,
238 ..
239 } => {
240 substitute_assignments(assignments, literals, idx);
241 substitute_assignments(on_conflict, literals, idx);
242 }
243 PlanNode::Update {
244 input, assignments, ..
245 } => {
246 substitute_plan(input, literals, idx);
247 substitute_assignments(assignments, literals, idx);
248 }
249 PlanNode::Delete { input, .. } => {
250 substitute_plan(input, literals, idx);
251 }
252 PlanNode::CreateTable { .. } => {}
253 PlanNode::CreateView { .. } => {}
254 PlanNode::RefreshView { .. } => {}
255 PlanNode::DropView { .. } => {}
256 PlanNode::Window { input, windows } => {
257 substitute_plan(input, literals, idx);
258 for w in windows {
259 for arg in &mut w.args {
260 substitute_expr(arg, literals, idx);
261 }
262 }
263 }
264 PlanNode::Union { left, right, .. } => {
265 substitute_plan(left, literals, idx);
266 substitute_plan(right, literals, idx);
267 }
268 PlanNode::Explain { input } => {
269 substitute_plan(input, literals, idx);
270 }
271 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
272 }
273}
274
275fn substitute_assignments(assignments: &mut [Assignment], literals: &[Literal], idx: &mut usize) {
276 for a in assignments {
277 substitute_expr(&mut a.value, literals, idx);
278 }
279}
280
281pub(crate) fn count_literal_slots(plan: &PlanNode) -> usize {
287 let mut n = 0usize;
288 count_plan(plan, &mut n);
289 n
290}
291
292fn count_plan(plan: &PlanNode, n: &mut usize) {
293 match plan {
294 PlanNode::SeqScan { .. } => {}
295 PlanNode::AliasScan { .. } => {}
296 PlanNode::IndexScan { key, .. } => count_expr(key, n),
297 PlanNode::RangeScan { start, end, .. } => {
298 if let Some((expr, _)) = start {
299 count_expr(expr, n);
300 }
301 if let Some((expr, _)) = end {
302 count_expr(expr, n);
303 }
304 }
305 PlanNode::Filter { input, predicate } => {
306 count_plan(input, n);
307 count_expr(predicate, n);
308 }
309 PlanNode::Project { input, fields } => {
310 count_plan(input, n);
311 for f in fields {
312 count_expr(&f.expr, n);
313 }
314 }
315 PlanNode::Sort { input, .. } => count_plan(input, n),
316 PlanNode::Limit { input, count } => {
317 if let PlanNode::Offset {
322 input: inner,
323 count: off_count,
324 } = input.as_ref()
325 {
326 count_plan(inner, n);
327 count_expr(count, n);
328 count_expr(off_count, n);
329 } else {
330 count_plan(input, n);
331 count_expr(count, n);
332 }
333 }
334 PlanNode::Offset { input, count } => {
335 count_plan(input, n);
336 count_expr(count, n);
337 }
338 PlanNode::Aggregate { input, .. } => count_plan(input, n),
339 PlanNode::NestedLoopJoin {
340 left, right, on, ..
341 } => {
342 count_plan(left, n);
343 count_plan(right, n);
344 if let Some(pred) = on {
345 count_expr(pred, n);
346 }
347 }
348 PlanNode::Distinct { input } => count_plan(input, n),
349 PlanNode::GroupBy { input, having, .. } => {
350 count_plan(input, n);
351 if let Some(pred) = having {
352 count_expr(pred, n);
353 }
354 }
355 PlanNode::Insert { rows, .. } => {
356 for assignments in rows {
357 for a in assignments {
358 count_expr(&a.value, n);
359 }
360 }
361 }
362 PlanNode::Upsert {
363 assignments,
364 on_conflict,
365 ..
366 } => {
367 for a in assignments {
368 count_expr(&a.value, n);
369 }
370 for a in on_conflict {
371 count_expr(&a.value, n);
372 }
373 }
374 PlanNode::Update {
375 input, assignments, ..
376 } => {
377 count_plan(input, n);
378 for a in assignments {
379 count_expr(&a.value, n);
380 }
381 }
382 PlanNode::Delete { input, .. } => count_plan(input, n),
383 PlanNode::CreateTable { .. } => {}
384 PlanNode::AlterTable { .. } => {}
385 PlanNode::DropTable { .. } => {}
386 PlanNode::CreateView { .. } => {}
387 PlanNode::RefreshView { .. } => {}
388 PlanNode::DropView { .. } => {}
389 PlanNode::Window { input, windows } => {
390 count_plan(input, n);
391 for w in windows {
392 for arg in &w.args {
393 count_expr(arg, n);
394 }
395 }
396 }
397 PlanNode::Union { left, right, .. } => {
398 count_plan(left, n);
399 count_plan(right, n);
400 }
401 PlanNode::Explain { input } => {
402 count_plan(input, n);
403 }
404 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
405 }
406}
407
408fn count_expr(expr: &Expr, n: &mut usize) {
409 match expr {
410 Expr::Literal(_) => *n += 1,
411 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
412 Expr::BinaryOp(l, _, r) => {
413 count_expr(l, n);
414 count_expr(r, n);
415 }
416 Expr::UnaryOp(_, inner) => count_expr(inner, n),
417 Expr::FunctionCall(_, inner) => count_expr(inner, n),
418 Expr::Coalesce(l, r) => {
419 count_expr(l, n);
420 count_expr(r, n);
421 }
422 Expr::InList { expr, list, .. } => {
423 count_expr(expr, n);
424 for item in list {
425 count_expr(item, n);
426 }
427 }
428 Expr::ScalarFunc(_, args) => {
429 for a in args {
430 count_expr(a, n);
431 }
432 }
433 Expr::Cast(inner, _) => count_expr(inner, n),
434 Expr::Case { whens, else_expr } => {
435 for (cond, result) in whens {
436 count_expr(cond, n);
437 count_expr(result, n);
438 }
439 if let Some(e) = else_expr {
440 count_expr(e, n);
441 }
442 }
443 Expr::InSubquery { expr, .. } => {
444 count_expr(expr, n);
445 }
448 Expr::ExistsSubquery { .. } => {
449 }
452 Expr::Window { args, .. } => {
453 for a in args {
454 count_expr(a, n);
455 }
456 }
457 Expr::ValueLit(_) => {}
460 Expr::Null => {}
461 }
462}
463
464fn substitute_expr(expr: &mut Expr, literals: &[Literal], idx: &mut usize) {
465 match expr {
466 Expr::Literal(_) => {
467 *expr = Expr::Literal(literals[*idx].clone());
471 *idx += 1;
472 }
473 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
474 Expr::BinaryOp(l, _, r) => {
475 substitute_expr(l, literals, idx);
476 substitute_expr(r, literals, idx);
477 }
478 Expr::UnaryOp(_, inner) => {
479 substitute_expr(inner, literals, idx);
480 }
481 Expr::FunctionCall(_, inner) => {
482 substitute_expr(inner, literals, idx);
483 }
484 Expr::Coalesce(l, r) => {
485 substitute_expr(l, literals, idx);
486 substitute_expr(r, literals, idx);
487 }
488 Expr::InList { expr, list, .. } => {
489 substitute_expr(expr, literals, idx);
490 for item in list {
491 substitute_expr(item, literals, idx);
492 }
493 }
494 Expr::ScalarFunc(_, args) => {
495 for a in args {
496 substitute_expr(a, literals, idx);
497 }
498 }
499 Expr::Cast(inner, _) => substitute_expr(inner, literals, idx),
500 Expr::Case { whens, else_expr } => {
501 for (cond, result) in whens {
502 substitute_expr(cond, literals, idx);
503 substitute_expr(result, literals, idx);
504 }
505 if let Some(e) = else_expr {
506 substitute_expr(e, literals, idx);
507 }
508 }
509 Expr::InSubquery { expr, .. } => {
510 substitute_expr(expr, literals, idx);
511 }
512 Expr::ExistsSubquery { .. } => {
513 }
516 Expr::Window { args, .. } => {
517 for a in args {
518 substitute_expr(a, literals, idx);
519 }
520 }
521 Expr::ValueLit(_) => {}
524 Expr::Null => {}
525 }
526}
527
528#[cfg(test)]
529mod tests {
530 use super::*;
531 use crate::canonicalize::canonicalize;
532 use crate::planner;
533
534 #[test]
535 fn test_cache_hit_substitutes_literal() {
536 let mut cache = PlanCache::new(100);
537
538 let q1 = "User filter .id = 42";
540 let (h1, lits1) = canonicalize(q1).unwrap();
541 let p1 = planner::plan(q1).unwrap();
542 cache.insert(h1, p1, lits1.len());
543
544 let q2 = "User filter .id = 99";
547 let (h2, lits2) = canonicalize(q2).unwrap();
548 assert_eq!(h1, h2, "different literals must hash the same");
549
550 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
551
552 match plan {
554 PlanNode::IndexScan { key, .. } => {
555 assert_eq!(key, Expr::Literal(Literal::Int(99)));
556 }
557 other => panic!("expected IndexScan, got {other:?}"),
558 }
559
560 assert_eq!(lits1, vec![Literal::Int(42)]);
563 assert_eq!(cache.hits, 1);
564 assert_eq!(cache.misses, 0);
565 }
566
567 #[test]
568 fn test_subquery_plan_not_cached() {
569 let mut cache = PlanCache::new(100);
576 let q = "User filter .id in (Ord filter .total > 100 { .user_id })";
577 let (h, lits) = canonicalize(q).unwrap();
578 assert_eq!(lits.len(), 1, "canonicalize collects the inner literal");
579 let plan = planner::plan(q).unwrap();
580 assert_eq!(
581 count_literal_slots(&plan),
582 0,
583 "the subquery literal is not a reachable substitution slot"
584 );
585 cache.insert(h, plan, lits.len());
586 assert!(cache.is_empty(), "subquery plans must not be cached (#137)");
587 assert!(cache.get_with_substitution(h, &lits).is_none());
588 }
589
590 #[test]
591 fn test_cache_miss_returns_none_and_bumps_counter() {
592 let mut cache = PlanCache::new(100);
593 assert!(cache.get_with_substitution(99999, &[]).is_none());
594 assert_eq!(cache.misses, 1);
595 assert_eq!(cache.hits, 0);
596 }
597
598 #[test]
599 fn test_multi_literal_filter_substitution() {
600 let mut cache = PlanCache::new(100);
601 let q1 = r#"User filter .age > 30 and .status = "active" { .name }"#;
602 let (h1, lits1) = canonicalize(q1).unwrap();
603 cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
604
605 let q2 = r#"User filter .age > 50 and .status = "pending" { .name }"#;
606 let (h2, lits2) = canonicalize(q2).unwrap();
607 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
608
609 let mut found = Vec::new();
611 collect_literals_for_test(&plan, &mut found);
612 assert_eq!(
613 found,
614 vec![Literal::Int(50), Literal::String("pending".into()),]
615 );
616 }
617
618 #[test]
619 fn test_update_by_pk_substitution() {
620 let mut cache = PlanCache::new(100);
621 let q1 = "User filter .id = 1 update { age := 100 }";
622 let (h1, lits1) = canonicalize(q1).unwrap();
623 cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
624
625 let q2 = "User filter .id = 7 update { age := 200 }";
626 let (h2, lits2) = canonicalize(q2).unwrap();
627 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
628
629 let mut found = Vec::new();
630 collect_literals_for_test(&plan, &mut found);
631 assert_eq!(found, vec![Literal::Int(7), Literal::Int(200)]);
632 }
633
634 #[test]
635 fn test_insert_substitution() {
636 let mut cache = PlanCache::new(100);
637 let q1 = r#"insert User { id := 1, name := "Alice", age := 20 }"#;
638 let (h1, lits1) = canonicalize(q1).unwrap();
639 cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
640
641 let q2 = r#"insert User { id := 2, name := "Bob", age := 30 }"#;
642 let (h2, lits2) = canonicalize(q2).unwrap();
643 let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
644
645 let mut found = Vec::new();
646 collect_literals_for_test(&plan, &mut found);
647 assert_eq!(
648 found,
649 vec![
650 Literal::Int(2),
651 Literal::String("Bob".into()),
652 Literal::Int(30),
653 ]
654 );
655 }
656
657 #[test]
658 fn test_eviction_on_capacity() {
659 let mut cache = PlanCache::new(2);
660 let q1 = "User";
661 let q2 = "User filter .age > 1";
662 let _q3 = "User filter .age > 2";
663 let q3_distinct = "User filter .id = 5";
666
667 let (h1, lits1) = canonicalize(q1).unwrap();
668 let (h2, lits2) = canonicalize(q2).unwrap();
669 let (h3, lits3) = canonicalize(q3_distinct).unwrap();
670 cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
671 cache.insert(h2, planner::plan(q2).unwrap(), lits2.len());
672 cache.insert(h3, planner::plan(q3_distinct).unwrap(), lits3.len());
674 assert!(cache.cache.contains_key(&h3));
675 assert_eq!(cache.cache.len(), 1);
676 }
677
678 fn collect_literals_for_test(plan: &PlanNode, out: &mut Vec<Literal>) {
682 match plan {
683 PlanNode::SeqScan { .. } => {}
684 PlanNode::AliasScan { .. } => {}
685 PlanNode::IndexScan { key, .. } => collect_expr_literals(key, out),
686 PlanNode::RangeScan { start, end, .. } => {
687 if let Some((expr, _)) = start {
688 collect_expr_literals(expr, out);
689 }
690 if let Some((expr, _)) = end {
691 collect_expr_literals(expr, out);
692 }
693 }
694 PlanNode::Filter { input, predicate } => {
695 collect_literals_for_test(input, out);
696 collect_expr_literals(predicate, out);
697 }
698 PlanNode::Project { input, fields } => {
699 collect_literals_for_test(input, out);
700 for f in fields {
701 collect_expr_literals(&f.expr, out);
702 }
703 }
704 PlanNode::Sort { input, .. } => collect_literals_for_test(input, out),
705 PlanNode::Limit { input, count } => {
706 collect_literals_for_test(input, out);
707 collect_expr_literals(count, out);
708 }
709 PlanNode::Offset { input, count } => {
710 collect_literals_for_test(input, out);
711 collect_expr_literals(count, out);
712 }
713 PlanNode::Aggregate { input, .. } => collect_literals_for_test(input, out),
714 PlanNode::NestedLoopJoin {
715 left, right, on, ..
716 } => {
717 collect_literals_for_test(left, out);
718 collect_literals_for_test(right, out);
719 if let Some(pred) = on {
720 collect_expr_literals(pred, out);
721 }
722 }
723 PlanNode::Insert { rows, .. } => {
724 for assignments in rows {
725 for a in assignments {
726 collect_expr_literals(&a.value, out);
727 }
728 }
729 }
730 PlanNode::Upsert {
731 assignments,
732 on_conflict,
733 ..
734 } => {
735 for a in assignments {
736 collect_expr_literals(&a.value, out);
737 }
738 for a in on_conflict {
739 collect_expr_literals(&a.value, out);
740 }
741 }
742 PlanNode::Update {
743 input, assignments, ..
744 } => {
745 collect_literals_for_test(input, out);
746 for a in assignments {
747 collect_expr_literals(&a.value, out);
748 }
749 }
750 PlanNode::Distinct { input } => collect_literals_for_test(input, out),
751 PlanNode::GroupBy { input, having, .. } => {
752 collect_literals_for_test(input, out);
753 if let Some(pred) = having {
754 collect_expr_literals(pred, out);
755 }
756 }
757 PlanNode::Delete { input, .. } => collect_literals_for_test(input, out),
758 PlanNode::CreateTable { .. } => {}
759 PlanNode::AlterTable { .. } => {}
760 PlanNode::DropTable { .. } => {}
761 PlanNode::CreateView { .. } => {}
762 PlanNode::RefreshView { .. } => {}
763 PlanNode::DropView { .. } => {}
764 PlanNode::Window { input, windows } => {
765 collect_literals_for_test(input, out);
766 for w in windows {
767 for arg in &w.args {
768 collect_expr_literals(arg, out);
769 }
770 }
771 }
772 PlanNode::Union { left, right, .. } => {
773 collect_literals_for_test(left, out);
774 collect_literals_for_test(right, out);
775 }
776 PlanNode::Explain { input } => {
777 collect_literals_for_test(input, out);
778 }
779 PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
780 }
781 }
782
783 fn collect_expr_literals(expr: &Expr, out: &mut Vec<Literal>) {
784 match expr {
785 Expr::Literal(l) => out.push(l.clone()),
786 Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
787 Expr::BinaryOp(l, _, r) => {
788 collect_expr_literals(l, out);
789 collect_expr_literals(r, out);
790 }
791 Expr::UnaryOp(_, inner) => collect_expr_literals(inner, out),
792 Expr::FunctionCall(_, inner) => collect_expr_literals(inner, out),
793 Expr::Coalesce(l, r) => {
794 collect_expr_literals(l, out);
795 collect_expr_literals(r, out);
796 }
797 Expr::InList { expr, list, .. } => {
798 collect_expr_literals(expr, out);
799 for item in list {
800 collect_expr_literals(item, out);
801 }
802 }
803 Expr::ScalarFunc(_, args) => {
804 for a in args {
805 collect_expr_literals(a, out);
806 }
807 }
808 Expr::Cast(inner, _) => collect_expr_literals(inner, out),
809 Expr::Case { whens, else_expr } => {
810 for (cond, result) in whens {
811 collect_expr_literals(cond, out);
812 collect_expr_literals(result, out);
813 }
814 if let Some(e) = else_expr {
815 collect_expr_literals(e, out);
816 }
817 }
818 Expr::InSubquery { expr, .. } => {
819 collect_expr_literals(expr, out);
820 }
821 Expr::ExistsSubquery { .. } => {}
822 Expr::Window { args, .. } => {
823 for a in args {
824 collect_expr_literals(a, out);
825 }
826 }
827 Expr::ValueLit(_) => {}
828 Expr::Null => {}
829 }
830 }
831}