linkrs_query/planning/
context.rs

1use std::sync::Arc;
2
3use crate::binder::validation::ValidatedStatement;
4use crate::binder::BoundStatement;
5use crate::metadata::MetadataContext;
6use crate::parser::ast::stmt::{Ast, Stmt};
7use crate::QueryContext;
8
9/// Unified context for the plan_bound phase.
10///
11/// Bundles all inputs needed by statement planners during bound-statement
12/// planning. Callers construct a PlanContext once; planners destructure
13/// only the fields they need.
14#[derive(Debug, Clone)]
15pub struct PlanContext<'a> {
16    /// The bound (post-binder) statement to plan.
17    pub bound: &'a BoundStatement,
18    /// Query-level context (space, session, parameters).
19    pub qctx: Arc<QueryContext>,
20    /// Pre-resolved metadata for referenced tags/edges/indexes.
21    /// Only populated when schema_manager is available.
22    pub metadata: Option<&'a MetadataContext>,
23    /// AST-derived validated statement (expression analysis context).
24    /// Needed by planners that still use AST-level info (e.g., MATCH YIELD columns).
25    pub validated: &'a ValidatedStatement,
26}
27
28impl<'a> PlanContext<'a> {
29    /// Create a new PlanContext.
30    pub fn new(
31        bound: &'a BoundStatement,
32        qctx: Arc<QueryContext>,
33        metadata: Option<&'a MetadataContext>,
34        validated: &'a ValidatedStatement,
35    ) -> Self {
36        Self {
37            bound,
38            qctx,
39            metadata,
40            validated,
41        }
42    }
43
44    /// Create a PlanContext with metadata injected later.
45    pub fn without_metadata(
46        bound: &'a BoundStatement,
47        qctx: Arc<QueryContext>,
48        validated: &'a ValidatedStatement,
49    ) -> Self {
50        Self {
51            bound,
52            qctx,
53            metadata: None,
54            validated,
55        }
56    }
57
58    /// Shorthand: return a sub-context containing only bound + qctx + validated,
59    /// for recursive planners (Pipe, SetOperation) that forward to sub-planners.
60    pub fn as_inner(&self) -> PlanContext<'_> {
61        PlanContext {
62            bound: self.bound,
63            qctx: self.qctx.clone(),
64            metadata: self.metadata,
65            validated: self.validated,
66        }
67    }
68
69    /// Replace the bound statement (for recursive planning of sub-statements).
70    pub fn with_bound<'b>(&'b self, bound: &'b BoundStatement) -> PlanContext<'b> {
71        PlanContext {
72            bound,
73            qctx: self.qctx.clone(),
74            metadata: self.metadata,
75            validated: self.validated,
76        }
77    }
78
79    /// Derive the `ValidatedStatement` aligned with a child bound statement.
80    ///
81    /// Composite planners (Pipe, SetOperation, Explain) forward a shared
82    /// `validated` that describes the whole query, while each stage planner
83    /// expects the AST fragment matching its own bound sub-statement (e.g.
84    /// `MatchStatementPlanner` requires `Stmt::Match`). The binder mirrors
85    /// the AST structure exactly, so the fragment is located by walking the
86    /// parent AST and parent bound statement in parallel.
87    ///
88    /// Returns `None` when the structures do not mirror each other; callers
89    /// then keep the parent `validated` (previous behavior).
90    pub fn derive_validated(&self, child: &BoundStatement) -> Option<ValidatedStatement> {
91        let stmt = find_aligned_stmt(self.validated.stmt(), self.bound, child)?;
92        Some(ValidatedStatement::new(
93            Arc::new(Ast::new(
94                stmt.clone(),
95                self.validated.ast.expr_context().clone(),
96            )),
97            self.validated.validation_info.clone(),
98        ))
99    }
100}
101
102/// Locate the AST sub-statement aligned with `child` by walking `stmt` and
103/// `bound` in parallel. Identity is by reference: `child` must be borrowed
104/// from inside `bound`.
105fn find_aligned_stmt<'s>(
106    stmt: &'s Stmt,
107    bound: &BoundStatement,
108    child: &BoundStatement,
109) -> Option<&'s Stmt> {
110    if std::ptr::eq(bound, child) {
111        return Some(stmt);
112    }
113    match (stmt, bound) {
114        (Stmt::Pipe(pipe_stmt), BoundStatement::Pipe(bound_pipe)) => {
115            let stmts = [&pipe_stmt.left, &pipe_stmt.right];
116            if bound_pipe.statements.len() != stmts.len() {
117                return None;
118            }
119            for (sub_stmt, sub_bound) in stmts.into_iter().zip(bound_pipe.statements.iter()) {
120                if let Some(found) = find_aligned_stmt(sub_stmt, sub_bound, child) {
121                    return Some(found);
122                }
123            }
124            None
125        }
126        (Stmt::SetOperation(set_op), BoundStatement::SetOperation(bound_set_op)) => {
127            find_aligned_stmt(&set_op.left, &bound_set_op.left, child)
128                .or_else(|| find_aligned_stmt(&set_op.right, &bound_set_op.right, child))
129        }
130        (Stmt::Assignment(assign_stmt), BoundStatement::Assignment(bound_assign)) => {
131            find_aligned_stmt(&assign_stmt.statement, &bound_assign.statement, child)
132        }
133        (Stmt::Explain(explain), BoundStatement::Explain(bound_explain)) => {
134            find_aligned_stmt(&explain.statement, &bound_explain.statement, child)
135        }
136        (Stmt::Profile(profile), BoundStatement::Profile(bound_profile)) => {
137            find_aligned_stmt(&profile.statement, &bound_profile.statement, child)
138        }
139        _ => None,
140    }
141}
linkrs_query/planning/context.rs

linkrs_query/planning/
context.rs
