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