1use super::{
10 CommandPlan, CtePlan, CtePlanBody, QueryPlan, RelationalPlan, SourcePlan, UnifiedPlan,
11};
12use crate::ir::ScalarExpr;
13
14impl UnifiedPlan {
15 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 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 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 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}