Skip to main content

uqa_sql/plan/
visit.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Read-only traversal of every scalar expression a plan owns: the counterpart of the plan rewrite for analyses that check a plan without changing it.
8
9use super::{
10    CommandPlan, CtePlan, CtePlanBody, QueryPlan, RelationalPlan, SourcePlan, UnifiedPlan,
11};
12use crate::ir::ScalarExpr;
13
14impl UnifiedPlan {
15    /// Visit every scalar expression this plan owns, each root once, including those of its CTEs, relational sources, scalar subqueries and command inputs. The nested nodes of a root are reached through [`ScalarExpr::visit`].
16    pub fn visit_scalar_expressions(&self, visit: &mut dyn FnMut(&ScalarExpr)) {
17        match self {
18            Self::Query(query) => query.visit_scalar_expressions(visit),
19            Self::Command(command) => command.visit_scalar_expressions(visit),
20        }
21    }
22}
23
24fn visit_ctes(ctes: &[CtePlan], visit: &mut dyn FnMut(&ScalarExpr)) {
25    for cte in ctes {
26        match &cte.body {
27            CtePlanBody::Query(query) => query.visit_scalar_expressions(visit),
28            CtePlanBody::Command(command) => command.visit_scalar_expressions(visit),
29        }
30        if let Some(cycle) = &cte.cycle {
31            visit(&cycle.mark_value);
32            visit(&cycle.mark_default);
33        }
34    }
35}
36
37impl QueryPlan {
38    /// Visit every scalar expression this query owns, each root once, including those of its CTEs, relational sources and scalar subqueries.
39    pub fn visit_scalar_expressions(&self, visit: &mut dyn FnMut(&ScalarExpr)) {
40        visit_ctes(&self.ctes, visit);
41        match &self.root {
42            RelationalPlan::QueryBlock(block) => {
43                if let Some(source) = &block.from {
44                    source.visit_scalar_expressions(visit);
45                }
46                for expression in block
47                    .r#where
48                    .iter()
49                    .chain(block.projections.iter().map(|projection| &projection.expr))
50                    .chain(&block.group_by)
51                    .chain(block.grouping_sets.iter().flatten())
52                    .chain(block.having.iter())
53                    .chain(block.order_by.iter().map(|order| &order.expr))
54                    .chain(block.limit.iter())
55                    .chain(block.offset.iter())
56                    .chain(&block.distinct_on)
57                {
58                    visit(expression);
59                }
60                for window in &block.windows {
61                    for expression in window.spec.expressions() {
62                        visit(expression);
63                    }
64                }
65                for subquery in &block.subqueries {
66                    subquery.visit_scalar_expressions(visit);
67                }
68            }
69            RelationalPlan::SetOp {
70                left,
71                right,
72                order_by,
73                limit,
74                offset,
75                subqueries,
76                ..
77            } => {
78                left.visit_scalar_expressions(visit);
79                right.visit_scalar_expressions(visit);
80                for expression in order_by
81                    .iter()
82                    .map(|order| &order.expr)
83                    .chain(limit.as_deref())
84                    .chain(offset.as_deref())
85                {
86                    visit(expression);
87                }
88                for subquery in subqueries {
89                    subquery.visit_scalar_expressions(visit);
90                }
91            }
92            RelationalPlan::Values { rows, subqueries } => {
93                for expression in rows.iter().flatten() {
94                    visit(expression);
95                }
96                for subquery in subqueries {
97                    subquery.visit_scalar_expressions(visit);
98                }
99            }
100        }
101    }
102}
103
104impl SourcePlan {
105    /// Visit every scalar expression this source owns, each root once: join conditions, `VALUES` rows, function arguments and the expressions of the query bodies nested in derived tables.
106    pub fn visit_scalar_expressions(&self, visit: &mut dyn FnMut(&ScalarExpr)) {
107        match self {
108            Self::Table { .. } => {}
109            Self::Subquery { body, .. } => body.visit_scalar_expressions(visit),
110            Self::Join {
111                left, right, on, ..
112            } => {
113                left.visit_scalar_expressions(visit);
114                right.visit_scalar_expressions(visit);
115                if let Some(on) = on {
116                    visit(on);
117                }
118            }
119            Self::Values { rows, .. } => {
120                for expression in rows.iter().flatten() {
121                    visit(expression);
122                }
123            }
124            Self::Function { args, .. } => {
125                for argument in args {
126                    visit(argument);
127                }
128            }
129            Self::FunctionGroup { functions, .. } => {
130                for argument in functions.iter().flat_map(|function| &function.args) {
131                    visit(argument);
132                }
133            }
134        }
135    }
136}
137
138impl CommandPlan {
139    /// Visit every scalar expression this command owns, each root once, including those of its CTEs, query inputs and source.
140    pub fn visit_scalar_expressions(&self, visit: &mut dyn FnMut(&ScalarExpr)) {
141        visit_ctes(self.ctes(), visit);
142        for expression in self.expressions() {
143            visit(expression);
144        }
145        for query in self.query_inputs() {
146            query.visit_scalar_expressions(visit);
147        }
148        if let Some(source) = self.source_input() {
149            source.visit_scalar_expressions(visit);
150        }
151    }
152}