Skip to main content

powdb_query/
plan_cache.rs

1//! Plan cache — Mission D9.
2//!
3//! Two queries that differ only in literal values share the same parsed +
4//! planned tree. Re-running the lexer + parser + planner on every call is
5//! pure overhead — easily 1-3μs per query, which is the entire budget on
6//! sub-microsecond workloads like `update_by_pk`. SQLite gets around this
7//! with `prepare_cached`; we get around it with this cache.
8//!
9//! ## How it works
10//!
11//! 1. [`crate::canonicalize::canonicalize`] lexes the input and produces
12//!    `(canonical_hash, literals)`. The hash collapses literal *values*
13//!    into placeholders, so `User filter .id = 1` and `User filter .id = 2`
14//!    have the same hash.
15//! 2. On the first call, [`PlanCache::insert`] stores the planned tree
16//!    keyed by the canonical hash. The plan still has the *first call's*
17//!    literal values baked into its `Expr::Literal` nodes — that's fine,
18//!    we'll overwrite them on subsequent hits.
19//! 3. On a subsequent call, [`PlanCache::get_with_substitution`] clones
20//!    the cached plan and walks it depth-first, replacing each
21//!    `Expr::Literal` it finds (in source order) with the corresponding
22//!    literal from the new query.
23//!
24//! The walk order is **deterministic and matches the source order of the
25//! literal list collected by `canonicalize`** — see the per-PlanNode
26//! comments below for the exact traversal contract.
27
28use crate::ast::{Assignment, Expr, Literal};
29use crate::plan::PlanNode;
30use rustc_hash::FxHashMap;
31
32/// LRU-ish plan cache keyed by canonical query hash.
33///
34/// Mission F: uses FxHashMap. The keys are u64 hashes (already pre-hashed
35/// by `canonicalize`), so SipHash is pure overhead — Fx is much cheaper for
36/// the integer-keyed lookup.
37pub 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    /// Store a planned query under its canonical hash. The plan can have
55    /// any literal values inside it — those will be overwritten on hit.
56    ///
57    /// `source_literal_count` is the number of literals
58    /// [`crate::canonicalize::canonicalize`] collected from the query text.
59    /// The cache only works because, on a hit, every collected literal maps
60    /// 1:1 to a substitutable slot the `substitute_plan` walk can reach.
61    /// If those counts disagree, the plan has literals the walk cannot
62    /// re-bind — the one case today is a subquery (`in (<subquery>)` /
63    /// `exists (...)`), whose inner literals `canonicalize` collects at the
64    /// token level but which live in an un-walked `QueryExpr` AST inside the
65    /// predicate. Caching such a plan would serve the *first* call's inner
66    /// literal to every later same-shape call: silent wrong rows in release,
67    /// a substitution-count assert in debug (issue #137). We refuse to cache
68    /// it; the engine then plans from source on every call, which is always
69    /// correct. This check runs only on the populating miss, so the hot
70    /// hit-path pays nothing.
71    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            // Crude eviction: when full, drop everything. Plan cache is
77            // small (capacity ~256) and bench loops only ever fill a
78            // handful of slots, so this is acceptable for now. A real LRU
79            // would matter once we have hundreds of distinct query shapes.
80            self.cache.clear();
81        }
82        self.cache.insert(hash, plan);
83    }
84
85    /// Look up a plan by canonical hash and return a clone with the new
86    /// literals substituted into every `Expr::Literal` slot in source
87    /// order.
88    ///
89    /// Returns `Some(plan)` on a hit and bumps `self.hits`. Returns `None`
90    /// on a miss and bumps `self.misses`. Returning `None` instead of
91    /// reaching for the planner here keeps this module dependency-free
92    /// from `planner` — the engine handles the miss path.
93    ///
94    /// The substitution is done on a **clone** of the cached plan, not the
95    /// stored copy. The cached plan stays pristine for the next call.
96    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
131/// Walk a plan tree depth-first, replacing every `Expr::Literal` with the
132/// next literal from `literals` (consumed by index). The traversal order
133/// is deterministic and matches the source order produced by
134/// [`crate::canonicalize::canonicalize`].
135///
136/// **Walk contract** — the cache only works because both ends agree on
137/// this order:
138///   - Children of recursive nodes (`Filter`, `Project`, `Sort`, `Limit`,
139///     `Offset`, `Aggregate`, `Update`, `Delete`) are visited *before*
140///     the local expressions, because the source `User filter ... { ... }`
141///     reads the table → predicate → projection in that order, and the
142///     planner wraps `SeqScan → Filter → Project` accordingly.
143///   - For `Update`, the input plan is visited first (which holds the
144///     filter literal), then assignments in declaration order — same
145///     order as the source `User filter .id = 42 update { age := 31 }`.
146///
147/// `pub(crate)` so the executor's prepared-statement API can reuse the
148/// exact same walk — same order as canonicalise, same as the cache.
149pub(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            // Source order for `filter ... limit N offset M` is
179            // [filter literals, N, M]. The planner now builds
180            // Limit(Offset(...)) so that execution skips M rows *before*
181            // taking N. Naively walking "input then count" would yield
182            // [filter, M, N] — wrong. Special-case `Limit(Offset(...))`
183            // to descend into Offset's own input (which holds the filter
184            // literals), then visit Limit.count, then Offset.count, so
185            // the literal stream stays in source order.
186            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            // Bare Offset (no wrapping Limit) — source order is
201            // [..., offset literal] so descend first then visit count.
202            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            // Walk order: left subtree → right subtree → on predicate.
212            // Matches canonicalise's source-order literal collection for
213            // joined queries: left source tokens come first, then right
214            // source tokens, then the `on` expression's literals (if any).
215            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::ListTypes | PlanNode::Describe { .. } => {}
272        PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
273    }
274}
275
276fn substitute_assignments(assignments: &mut [Assignment], literals: &[Literal], idx: &mut usize) {
277    for a in assignments {
278        substitute_expr(&mut a.value, literals, idx);
279    }
280}
281
282/// Count every `Expr::Literal` slot reachable from `plan` using the same
283/// walk order as [`substitute_plan`]. Used by `Engine::prepare` to validate
284/// that calls to `execute_prepared` pass the right number of literals, and
285/// to fail early if a caller prepares a query with zero literals (which
286/// would be a no-op for the prepared API — better to catch that up front).
287pub(crate) fn count_literal_slots(plan: &PlanNode) -> usize {
288    let mut n = 0usize;
289    count_plan(plan, &mut n);
290    n
291}
292
293fn count_plan(plan: &PlanNode, n: &mut usize) {
294    match plan {
295        PlanNode::SeqScan { .. } => {}
296        PlanNode::AliasScan { .. } => {}
297        PlanNode::IndexScan { key, .. } => count_expr(key, n),
298        PlanNode::RangeScan { start, end, .. } => {
299            if let Some((expr, _)) = start {
300                count_expr(expr, n);
301            }
302            if let Some((expr, _)) = end {
303                count_expr(expr, n);
304            }
305        }
306        PlanNode::Filter { input, predicate } => {
307            count_plan(input, n);
308            count_expr(predicate, n);
309        }
310        PlanNode::Project { input, fields } => {
311            count_plan(input, n);
312            for f in fields {
313                count_expr(&f.expr, n);
314            }
315        }
316        PlanNode::Sort { input, .. } => count_plan(input, n),
317        PlanNode::Limit { input, count } => {
318            // Mirror the substitute walk: `Limit(Offset(...))` descends
319            // into the offset's child first, then counts Limit.count,
320            // then Offset.count. Source order is
321            // [..., limit literal, offset literal].
322            if let PlanNode::Offset {
323                input: inner,
324                count: off_count,
325            } = input.as_ref()
326            {
327                count_plan(inner, n);
328                count_expr(count, n);
329                count_expr(off_count, n);
330            } else {
331                count_plan(input, n);
332                count_expr(count, n);
333            }
334        }
335        PlanNode::Offset { input, count } => {
336            count_plan(input, n);
337            count_expr(count, n);
338        }
339        PlanNode::Aggregate { input, .. } => count_plan(input, n),
340        PlanNode::NestedLoopJoin {
341            left, right, on, ..
342        } => {
343            count_plan(left, n);
344            count_plan(right, n);
345            if let Some(pred) = on {
346                count_expr(pred, n);
347            }
348        }
349        PlanNode::Distinct { input } => count_plan(input, n),
350        PlanNode::GroupBy { input, having, .. } => {
351            count_plan(input, n);
352            if let Some(pred) = having {
353                count_expr(pred, n);
354            }
355        }
356        PlanNode::Insert { rows, .. } => {
357            for assignments in rows {
358                for a in assignments {
359                    count_expr(&a.value, n);
360                }
361            }
362        }
363        PlanNode::Upsert {
364            assignments,
365            on_conflict,
366            ..
367        } => {
368            for a in assignments {
369                count_expr(&a.value, n);
370            }
371            for a in on_conflict {
372                count_expr(&a.value, n);
373            }
374        }
375        PlanNode::Update {
376            input, assignments, ..
377        } => {
378            count_plan(input, n);
379            for a in assignments {
380                count_expr(&a.value, n);
381            }
382        }
383        PlanNode::Delete { input, .. } => count_plan(input, n),
384        PlanNode::CreateTable { .. } => {}
385        PlanNode::AlterTable { .. } => {}
386        PlanNode::DropTable { .. } => {}
387        PlanNode::CreateView { .. } => {}
388        PlanNode::RefreshView { .. } => {}
389        PlanNode::DropView { .. } => {}
390        PlanNode::Window { input, windows } => {
391            count_plan(input, n);
392            for w in windows {
393                for arg in &w.args {
394                    count_expr(arg, n);
395                }
396            }
397        }
398        PlanNode::Union { left, right, .. } => {
399            count_plan(left, n);
400            count_plan(right, n);
401        }
402        PlanNode::Explain { input } => {
403            count_plan(input, n);
404        }
405        PlanNode::ListTypes | PlanNode::Describe { .. } => {}
406        PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
407    }
408}
409
410fn count_expr(expr: &Expr, n: &mut usize) {
411    match expr {
412        Expr::Literal(_) => *n += 1,
413        Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
414        Expr::BinaryOp(l, _, r) => {
415            count_expr(l, n);
416            count_expr(r, n);
417        }
418        Expr::UnaryOp(_, inner) => count_expr(inner, n),
419        Expr::FunctionCall(_, inner) => count_expr(inner, n),
420        Expr::Coalesce(l, r) => {
421            count_expr(l, n);
422            count_expr(r, n);
423        }
424        Expr::InList { expr, list, .. } => {
425            count_expr(expr, n);
426            for item in list {
427                count_expr(item, n);
428            }
429        }
430        Expr::ScalarFunc(_, args) => {
431            for a in args {
432                count_expr(a, n);
433            }
434        }
435        Expr::Cast(inner, _) => count_expr(inner, n),
436        Expr::Case { whens, else_expr } => {
437            for (cond, result) in whens {
438                count_expr(cond, n);
439                count_expr(result, n);
440            }
441            if let Some(e) = else_expr {
442                count_expr(e, n);
443            }
444        }
445        Expr::InSubquery { expr, .. } => {
446            count_expr(expr, n);
447            // Subquery literals are not counted — the subquery is
448            // re-planned/executed separately.
449        }
450        Expr::ExistsSubquery { .. } => {
451            // Subquery literals are not counted — the subquery is
452            // re-planned/executed separately.
453        }
454        Expr::Window { args, .. } => {
455            for a in args {
456                count_expr(a, n);
457            }
458        }
459        // JSON path segments are STRUCTURAL, never literal slots (#137): only
460        // the base can carry literals (it can't today — it is a Field — but
461        // recursing keeps this correct if the base grammar ever widens).
462        Expr::JsonPath { base, .. } => count_expr(base, n),
463        // Runtime-only literal (correlated/subquery substitution); never
464        // reaches the plan cache, and it occupies no source literal slot.
465        Expr::ValueLit(_) => {}
466        Expr::Null => {}
467    }
468}
469
470fn substitute_expr(expr: &mut Expr, literals: &[Literal], idx: &mut usize) {
471    match expr {
472        Expr::Literal(_) => {
473            // The cached plan held the *first* call's literal at this
474            // slot; replace with the new call's value at the matching
475            // source position.
476            *expr = Expr::Literal(literals[*idx].clone());
477            *idx += 1;
478        }
479        Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
480        Expr::BinaryOp(l, _, r) => {
481            substitute_expr(l, literals, idx);
482            substitute_expr(r, literals, idx);
483        }
484        Expr::UnaryOp(_, inner) => {
485            substitute_expr(inner, literals, idx);
486        }
487        Expr::FunctionCall(_, inner) => {
488            substitute_expr(inner, literals, idx);
489        }
490        Expr::Coalesce(l, r) => {
491            substitute_expr(l, literals, idx);
492            substitute_expr(r, literals, idx);
493        }
494        Expr::InList { expr, list, .. } => {
495            substitute_expr(expr, literals, idx);
496            for item in list {
497                substitute_expr(item, literals, idx);
498            }
499        }
500        Expr::ScalarFunc(_, args) => {
501            for a in args {
502                substitute_expr(a, literals, idx);
503            }
504        }
505        Expr::Cast(inner, _) => substitute_expr(inner, literals, idx),
506        Expr::Case { whens, else_expr } => {
507            for (cond, result) in whens {
508                substitute_expr(cond, literals, idx);
509                substitute_expr(result, literals, idx);
510            }
511            if let Some(e) = else_expr {
512                substitute_expr(e, literals, idx);
513            }
514        }
515        Expr::InSubquery { expr, .. } => {
516            substitute_expr(expr, literals, idx);
517        }
518        Expr::ExistsSubquery { .. } => {
519            // Subquery has its own literal list; nothing to substitute
520            // at this level.
521        }
522        Expr::Window { args, .. } => {
523            for a in args {
524                substitute_expr(a, literals, idx);
525            }
526        }
527        // JSON path segments are STRUCTURAL (#137): substitution only recurses
528        // into the base, mirroring `count_expr`, so the slot walk stays aligned.
529        Expr::JsonPath { base, .. } => substitute_expr(base, literals, idx),
530        // Runtime-only literal (correlated/subquery substitution); never
531        // reaches the plan cache, so there is nothing to substitute.
532        Expr::ValueLit(_) => {}
533        Expr::Null => {}
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540    use crate::canonicalize::canonicalize;
541    use crate::planner;
542
543    #[test]
544    fn test_cache_hit_substitutes_literal() {
545        let mut cache = PlanCache::new(100);
546
547        // First call: "User filter .id = 42" — miss, plan + insert.
548        let q1 = "User filter .id = 42";
549        let (h1, lits1) = canonicalize(q1).unwrap();
550        let p1 = planner::plan(q1).unwrap();
551        cache.insert(h1, p1, lits1.len());
552
553        // Second call with a different literal — should hit and produce
554        // a plan with the new literal substituted in.
555        let q2 = "User filter .id = 99";
556        let (h2, lits2) = canonicalize(q2).unwrap();
557        assert_eq!(h1, h2, "different literals must hash the same");
558
559        let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
560
561        // The plan should be `IndexScan { key: Literal::Int(99) }`.
562        match plan {
563            PlanNode::IndexScan { key, .. } => {
564                assert_eq!(key, Expr::Literal(Literal::Int(99)));
565            }
566            other => panic!("expected IndexScan, got {other:?}"),
567        }
568
569        // First call's literal vector still holds 42, untouched — proves
570        // we substituted on a clone, not the cached template.
571        assert_eq!(lits1, vec![Literal::Int(42)]);
572        assert_eq!(cache.hits, 1);
573        assert_eq!(cache.misses, 0);
574    }
575
576    #[test]
577    fn test_subquery_plan_not_cached() {
578        // #137: `canonicalize` collects the inner `100` literal at the token
579        // level, but it lives in an un-walked subquery AST that
580        // `substitute_plan` can't reach (`count_literal_slots` returns 0).
581        // The counts disagree, so the cache must refuse to store the plan —
582        // otherwise a later same-shape call with a different inner literal
583        // would be served this plan's stale `100`.
584        let mut cache = PlanCache::new(100);
585        let q = "User filter .id in (Ord filter .total > 100 { .user_id })";
586        let (h, lits) = canonicalize(q).unwrap();
587        assert_eq!(lits.len(), 1, "canonicalize collects the inner literal");
588        let plan = planner::plan(q).unwrap();
589        assert_eq!(
590            count_literal_slots(&plan),
591            0,
592            "the subquery literal is not a reachable substitution slot"
593        );
594        cache.insert(h, plan, lits.len());
595        assert!(cache.is_empty(), "subquery plans must not be cached (#137)");
596        assert!(cache.get_with_substitution(h, &lits).is_none());
597    }
598
599    #[test]
600    fn test_cache_miss_returns_none_and_bumps_counter() {
601        let mut cache = PlanCache::new(100);
602        assert!(cache.get_with_substitution(99999, &[]).is_none());
603        assert_eq!(cache.misses, 1);
604        assert_eq!(cache.hits, 0);
605    }
606
607    #[test]
608    fn test_multi_literal_filter_substitution() {
609        let mut cache = PlanCache::new(100);
610        let q1 = r#"User filter .age > 30 and .status = "active" { .name }"#;
611        let (h1, lits1) = canonicalize(q1).unwrap();
612        cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
613
614        let q2 = r#"User filter .age > 50 and .status = "pending" { .name }"#;
615        let (h2, lits2) = canonicalize(q2).unwrap();
616        let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
617
618        // Walk the plan and pull out every literal — should be [50, "pending"].
619        let mut found = Vec::new();
620        collect_literals_for_test(&plan, &mut found);
621        assert_eq!(
622            found,
623            vec![Literal::Int(50), Literal::String("pending".into()),]
624        );
625    }
626
627    #[test]
628    fn test_grouped_join_query_shares_plan_across_literals() {
629        // P-5: a qualified-key grouped-join query with a HAVING literal must
630        // canonicalize identically regardless of that literal, cache cleanly
631        // (#137 slot-count invariant holds: group keys and aggregate args are
632        // structural identifiers, never literal slots), and substitute the new
633        // HAVING literal on a hit.
634        let mut cache = PlanCache::new(100);
635        let q1 = "User as u join Order as o on u.id = o.user_id \
636                  group u.status having count(o.total) > 1 { u.status, n: count(o.total) }";
637        let (h1, lits1) = canonicalize(q1).unwrap();
638        let p1 = planner::plan(q1).unwrap();
639
640        // Only the HAVING `1` is a substitutable literal; the qualified keys
641        // and aggregate arguments are identifiers, so the slot count matches
642        // the canonical literal count and the plan is cacheable.
643        assert_eq!(lits1.len(), 1, "only the HAVING literal is collected");
644        assert_eq!(
645            count_literal_slots(&p1),
646            lits1.len(),
647            "group keys/args are structural, so slots == literals (#137)"
648        );
649        cache.insert(h1, p1, lits1.len());
650        assert_eq!(cache.len(), 1, "grouped-join plan must cache");
651
652        let q2 = "User as u join Order as o on u.id = o.user_id \
653                  group u.status having count(o.total) > 5 { u.status, n: count(o.total) }";
654        let (h2, lits2) = canonicalize(q2).unwrap();
655        assert_eq!(h1, h2, "different HAVING literal must hash the same");
656
657        let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
658        let mut found = Vec::new();
659        collect_literals_for_test(&plan, &mut found);
660        assert_eq!(
661            found,
662            vec![Literal::Int(5)],
663            "new HAVING literal substituted"
664        );
665        assert_eq!(cache.hits, 1);
666    }
667
668    #[test]
669    fn json_path_slot_count_invariant() {
670        // #137 property: over path-bearing queries, count_literal_slots(plan)
671        // must equal the source literal count from canonicalize. Path segments
672        // (keys and array indexes) are structural and must not inflate either
673        // side, so the plan is cacheable.
674        for q in [
675            r#"Post filter .data->author->name = "x""#,
676            r#"Post filter .data->tags->0 = "rust""#,
677            r#"Post filter .data->age > 21 and .data->year = 2026"#,
678            r#"Post filter .data->"weird key" = 1 { .id }"#,
679            r#"Post { author: .data->author, first_tag: .data->tags->0 }"#,
680        ] {
681            let (_, lits) = canonicalize(q).unwrap();
682            let plan = planner::plan(q).unwrap();
683            assert_eq!(
684                count_literal_slots(&plan),
685                lits.len(),
686                "slot count must equal source literal count for `{q}`"
687            );
688        }
689    }
690
691    #[test]
692    fn json_path_plan_round_trips_cache() {
693        // A path-bearing query caches on first call and, on a second call with
694        // a DIFFERENT comparison literal but the SAME path, hits the cache and
695        // substitutes the new literal (no panic, no stale value).
696        let mut cache = PlanCache::new(100);
697        let q1 = r#"Post filter .data->age > 21"#;
698        let (h1, lits1) = canonicalize(q1).unwrap();
699        let p1 = planner::plan(q1).unwrap();
700        assert_eq!(count_literal_slots(&p1), lits1.len());
701        cache.insert(h1, p1, lits1.len());
702        assert_eq!(cache.len(), 1, "path plan must cache");
703
704        let q2 = r#"Post filter .data->age > 65"#;
705        let (h2, lits2) = canonicalize(q2).unwrap();
706        assert_eq!(h1, h2, "same path, different literal → same hash");
707        let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
708
709        let mut found = Vec::new();
710        collect_literals_for_test(&plan, &mut found);
711        assert_eq!(found, vec![Literal::Int(65)], "new literal substituted");
712        assert_eq!(cache.hits, 1);
713    }
714
715    #[test]
716    fn json_path_different_path_is_a_cache_miss() {
717        let mut cache = PlanCache::new(100);
718        let q1 = r#"Post filter .data->age > 21"#;
719        let (h1, lits1) = canonicalize(q1).unwrap();
720        cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
721
722        // Different key → different shape → miss (would otherwise serve a plan
723        // that walks the wrong path).
724        let q2 = r#"Post filter .data->year > 21"#;
725        let (h2, lits2) = canonicalize(q2).unwrap();
726        assert!(
727            cache.get_with_substitution(h2, &lits2).is_none(),
728            "a different path must not hit the cached plan"
729        );
730    }
731
732    #[test]
733    fn test_update_by_pk_substitution() {
734        let mut cache = PlanCache::new(100);
735        let q1 = "User filter .id = 1 update { age := 100 }";
736        let (h1, lits1) = canonicalize(q1).unwrap();
737        cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
738
739        let q2 = "User filter .id = 7 update { age := 200 }";
740        let (h2, lits2) = canonicalize(q2).unwrap();
741        let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
742
743        let mut found = Vec::new();
744        collect_literals_for_test(&plan, &mut found);
745        assert_eq!(found, vec![Literal::Int(7), Literal::Int(200)]);
746    }
747
748    #[test]
749    fn test_insert_substitution() {
750        let mut cache = PlanCache::new(100);
751        let q1 = r#"insert User { id := 1, name := "Alice", age := 20 }"#;
752        let (h1, lits1) = canonicalize(q1).unwrap();
753        cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
754
755        let q2 = r#"insert User { id := 2, name := "Bob", age := 30 }"#;
756        let (h2, lits2) = canonicalize(q2).unwrap();
757        let plan = cache.get_with_substitution(h2, &lits2).expect("hit");
758
759        let mut found = Vec::new();
760        collect_literals_for_test(&plan, &mut found);
761        assert_eq!(
762            found,
763            vec![
764                Literal::Int(2),
765                Literal::String("Bob".into()),
766                Literal::Int(30),
767            ]
768        );
769    }
770
771    /// #151/#140-class: `uuid("…")` const-fold sugar plans as
772    /// `Cast(Literal::String, Uuid)` — one substitutable slot. It must cache
773    /// AND rebind the inner string on a same-shape second call (a bulk load).
774    #[test]
775    fn test_insert_uuid_sugar_cacheable_and_substitutes() {
776        let mut cache = PlanCache::new(100);
777        let q1 = r#"insert User { id := uuid("00000000-0000-0000-0000-000000000001") }"#;
778        let (h1, lits1) = canonicalize(q1).unwrap();
779        assert_eq!(lits1.len(), 1, "the inner string is the only literal");
780        let plan = planner::plan(q1).unwrap();
781        assert_eq!(
782            count_literal_slots(&plan),
783            1,
784            "Cast wrapping a Literal is a reachable substitution slot"
785        );
786        cache.insert(h1, plan, lits1.len());
787        assert!(!cache.is_empty(), "uuid() insert must be cacheable");
788
789        let q2 = r#"insert User { id := uuid("00000000-0000-0000-0000-000000000002") }"#;
790        let (h2, lits2) = canonicalize(q2).unwrap();
791        assert_eq!(h1, h2, "same shape hashes identically");
792        let subst = cache.get_with_substitution(h2, &lits2).expect("hit");
793
794        let mut found = Vec::new();
795        collect_literals_for_test(&subst, &mut found);
796        assert_eq!(
797            found,
798            vec![Literal::String(
799                "00000000-0000-0000-0000-000000000002".into()
800            )],
801            "the second call's uuid must be substituted in, not the cached one"
802        );
803    }
804
805    /// Two-arg `cast(.x, "uuid")` bakes the target into the AST but leaves the
806    /// `"uuid"` string as a collected literal with no matching plan slot, so
807    /// the count-mismatch guard refuses to cache it (pre-existing behavior,
808    /// now confirmed for the uuid target).
809    #[test]
810    fn test_two_arg_cast_uuid_not_cached() {
811        let mut cache = PlanCache::new(100);
812        let q = r#"User filter .id = cast(.other, "uuid")"#;
813        let (h, lits) = canonicalize(q).unwrap();
814        assert_eq!(
815            lits.len(),
816            1,
817            "canonicalize collects the cast-target string"
818        );
819        let plan = planner::plan(q).unwrap();
820        assert_eq!(
821            count_literal_slots(&plan),
822            0,
823            "the cast target is baked into the AST, not a slot"
824        );
825        cache.insert(h, plan, lits.len());
826        assert!(cache.is_empty(), "cast(x, \"uuid\") must not be cached");
827    }
828
829    #[test]
830    fn test_eviction_on_capacity() {
831        let mut cache = PlanCache::new(2);
832        let q1 = "User";
833        let q2 = "User filter .age > 1";
834        let _q3 = "User filter .age > 2";
835        // q3 has same canonical as q2 — won't trigger eviction.
836        // Use a different shape to force eviction.
837        let q3_distinct = "User filter .id = 5";
838
839        let (h1, lits1) = canonicalize(q1).unwrap();
840        let (h2, lits2) = canonicalize(q2).unwrap();
841        let (h3, lits3) = canonicalize(q3_distinct).unwrap();
842        cache.insert(h1, planner::plan(q1).unwrap(), lits1.len());
843        cache.insert(h2, planner::plan(q2).unwrap(), lits2.len());
844        // Cache full → inserting a third *new* shape should clear.
845        cache.insert(h3, planner::plan(q3_distinct).unwrap(), lits3.len());
846        assert!(cache.cache.contains_key(&h3));
847        assert_eq!(cache.cache.len(), 1);
848    }
849
850    /// Test helper — depth-first walk that pulls out every Literal in the
851    /// same order `substitute_plan` would visit them. Used to verify
852    /// substitution actually wrote to the right slots.
853    fn collect_literals_for_test(plan: &PlanNode, out: &mut Vec<Literal>) {
854        match plan {
855            PlanNode::SeqScan { .. } => {}
856            PlanNode::AliasScan { .. } => {}
857            PlanNode::IndexScan { key, .. } => collect_expr_literals(key, out),
858            PlanNode::RangeScan { start, end, .. } => {
859                if let Some((expr, _)) = start {
860                    collect_expr_literals(expr, out);
861                }
862                if let Some((expr, _)) = end {
863                    collect_expr_literals(expr, out);
864                }
865            }
866            PlanNode::Filter { input, predicate } => {
867                collect_literals_for_test(input, out);
868                collect_expr_literals(predicate, out);
869            }
870            PlanNode::Project { input, fields } => {
871                collect_literals_for_test(input, out);
872                for f in fields {
873                    collect_expr_literals(&f.expr, out);
874                }
875            }
876            PlanNode::Sort { input, .. } => collect_literals_for_test(input, out),
877            PlanNode::Limit { input, count } => {
878                collect_literals_for_test(input, out);
879                collect_expr_literals(count, out);
880            }
881            PlanNode::Offset { input, count } => {
882                collect_literals_for_test(input, out);
883                collect_expr_literals(count, out);
884            }
885            PlanNode::Aggregate { input, .. } => collect_literals_for_test(input, out),
886            PlanNode::NestedLoopJoin {
887                left, right, on, ..
888            } => {
889                collect_literals_for_test(left, out);
890                collect_literals_for_test(right, out);
891                if let Some(pred) = on {
892                    collect_expr_literals(pred, out);
893                }
894            }
895            PlanNode::Insert { rows, .. } => {
896                for assignments in rows {
897                    for a in assignments {
898                        collect_expr_literals(&a.value, out);
899                    }
900                }
901            }
902            PlanNode::Upsert {
903                assignments,
904                on_conflict,
905                ..
906            } => {
907                for a in assignments {
908                    collect_expr_literals(&a.value, out);
909                }
910                for a in on_conflict {
911                    collect_expr_literals(&a.value, out);
912                }
913            }
914            PlanNode::Update {
915                input, assignments, ..
916            } => {
917                collect_literals_for_test(input, out);
918                for a in assignments {
919                    collect_expr_literals(&a.value, out);
920                }
921            }
922            PlanNode::Distinct { input } => collect_literals_for_test(input, out),
923            PlanNode::GroupBy { input, having, .. } => {
924                collect_literals_for_test(input, out);
925                if let Some(pred) = having {
926                    collect_expr_literals(pred, out);
927                }
928            }
929            PlanNode::Delete { input, .. } => collect_literals_for_test(input, out),
930            PlanNode::CreateTable { .. } => {}
931            PlanNode::AlterTable { .. } => {}
932            PlanNode::DropTable { .. } => {}
933            PlanNode::CreateView { .. } => {}
934            PlanNode::RefreshView { .. } => {}
935            PlanNode::DropView { .. } => {}
936            PlanNode::Window { input, windows } => {
937                collect_literals_for_test(input, out);
938                for w in windows {
939                    for arg in &w.args {
940                        collect_expr_literals(arg, out);
941                    }
942                }
943            }
944            PlanNode::Union { left, right, .. } => {
945                collect_literals_for_test(left, out);
946                collect_literals_for_test(right, out);
947            }
948            PlanNode::Explain { input } => {
949                collect_literals_for_test(input, out);
950            }
951            PlanNode::ListTypes | PlanNode::Describe { .. } => {}
952            PlanNode::Begin | PlanNode::Commit | PlanNode::Rollback => {}
953        }
954    }
955
956    fn collect_expr_literals(expr: &Expr, out: &mut Vec<Literal>) {
957        match expr {
958            Expr::Literal(l) => out.push(l.clone()),
959            Expr::Field(_) | Expr::QualifiedField { .. } | Expr::Param(_) => {}
960            Expr::BinaryOp(l, _, r) => {
961                collect_expr_literals(l, out);
962                collect_expr_literals(r, out);
963            }
964            Expr::UnaryOp(_, inner) => collect_expr_literals(inner, out),
965            Expr::FunctionCall(_, inner) => collect_expr_literals(inner, out),
966            Expr::Coalesce(l, r) => {
967                collect_expr_literals(l, out);
968                collect_expr_literals(r, out);
969            }
970            Expr::InList { expr, list, .. } => {
971                collect_expr_literals(expr, out);
972                for item in list {
973                    collect_expr_literals(item, out);
974                }
975            }
976            Expr::ScalarFunc(_, args) => {
977                for a in args {
978                    collect_expr_literals(a, out);
979                }
980            }
981            Expr::Cast(inner, _) => collect_expr_literals(inner, out),
982            Expr::Case { whens, else_expr } => {
983                for (cond, result) in whens {
984                    collect_expr_literals(cond, out);
985                    collect_expr_literals(result, out);
986                }
987                if let Some(e) = else_expr {
988                    collect_expr_literals(e, out);
989                }
990            }
991            Expr::InSubquery { expr, .. } => {
992                collect_expr_literals(expr, out);
993            }
994            Expr::ExistsSubquery { .. } => {}
995            Expr::Window { args, .. } => {
996                for a in args {
997                    collect_expr_literals(a, out);
998                }
999            }
1000            // JSON path segments are structural, never literals — mirror
1001            // count_expr/substitute_expr and recurse into the base only.
1002            Expr::JsonPath { base, .. } => collect_expr_literals(base, out),
1003            Expr::ValueLit(_) => {}
1004            Expr::Null => {}
1005        }
1006    }
1007}