1use super::{
10 AssignmentPlan, CommandPlan, ConflictActionPlan, MergeWhenPlan, OrderPlan, ProjectionPlan,
11 QueryPlan, RelationalPlan, ScalarExpr, ScalarFrameBound, SourcePlan,
12};
13
14pub(super) fn rewrite_query_scalars(
15 query: &mut QueryPlan,
16 rewrite: &mut dyn FnMut(&mut ScalarExpr),
17) {
18 for cte in &mut query.ctes {
19 match &mut cte.body {
20 super::CtePlanBody::Query(query) => rewrite_query_scalars(query, rewrite),
21 super::CtePlanBody::Command(command) => rewrite_command_scalars(command, rewrite),
22 }
23 }
24 match &mut query.root {
25 RelationalPlan::QueryBlock(block) => {
26 if let Some(source) = &mut block.from {
27 rewrite_source_scalars(source, rewrite);
28 }
29 rewrite_optional_scalar(&mut block.r#where, rewrite);
30 for projection in &mut block.projections {
31 rewrite_scalar(&mut projection.expr, rewrite);
32 }
33 for expression in &mut block.group_by {
34 rewrite_scalar(expression, rewrite);
35 }
36 for set in &mut block.grouping_sets {
37 for expression in set {
38 rewrite_scalar(expression, rewrite);
39 }
40 }
41 rewrite_optional_scalar(&mut block.having, rewrite);
42 rewrite_orders(&mut block.order_by, rewrite);
43 rewrite_optional_scalar(&mut block.limit, rewrite);
44 rewrite_optional_scalar(&mut block.offset, rewrite);
45 for expression in &mut block.distinct_on {
46 rewrite_scalar(expression, rewrite);
47 }
48 for subquery in &mut block.subqueries {
49 rewrite_query_scalars(subquery, rewrite);
50 }
51 }
52 RelationalPlan::SetOp {
53 left,
54 right,
55 order_by,
56 limit,
57 offset,
58 subqueries,
59 ..
60 } => {
61 rewrite_query_scalars(left, rewrite);
62 rewrite_query_scalars(right, rewrite);
63 rewrite_orders(order_by, rewrite);
64 if let Some(limit) = limit {
65 rewrite_scalar(limit, rewrite);
66 }
67 if let Some(offset) = offset {
68 rewrite_scalar(offset, rewrite);
69 }
70 for subquery in subqueries {
71 rewrite_query_scalars(subquery, rewrite);
72 }
73 }
74 RelationalPlan::Values { rows, subqueries } => {
75 for row in rows {
76 for expression in row {
77 rewrite_scalar(expression, rewrite);
78 }
79 }
80 for subquery in subqueries {
81 rewrite_query_scalars(subquery, rewrite);
82 }
83 }
84 }
85}
86
87pub(super) fn rewrite_source_scalars(
88 source: &mut SourcePlan,
89 rewrite: &mut dyn FnMut(&mut ScalarExpr),
90) {
91 match source {
92 SourcePlan::Table { .. } => {}
93 SourcePlan::Join {
94 left, right, on, ..
95 } => {
96 rewrite_source_scalars(left, rewrite);
97 rewrite_source_scalars(right, rewrite);
98 rewrite_optional_scalar(on, rewrite);
99 }
100 SourcePlan::Values { rows, .. } => {
101 for row in rows {
102 for expression in row {
103 rewrite_scalar(expression, rewrite);
104 }
105 }
106 }
107 SourcePlan::Function { args, .. } => {
108 for expression in args {
109 rewrite_scalar(expression, rewrite);
110 }
111 }
112 SourcePlan::FunctionGroup { functions, .. } => {
113 for function in functions {
114 for expression in &mut function.args {
115 rewrite_scalar(expression, rewrite);
116 }
117 }
118 }
119 SourcePlan::Subquery { body, .. } => rewrite_query_scalars(body, rewrite),
120 }
121}
122
123#[expect(
124 clippy::too_many_lines,
125 reason = "optimizer rewrite preserves exhaustive variants and fixed-point order"
126)]
127pub(super) fn rewrite_command_scalars(
128 command: &mut CommandPlan,
129 rewrite: &mut dyn FnMut(&mut ScalarExpr),
130) {
131 match command {
132 CommandPlan::Insert(plan) => {
133 for cte in &mut plan.ctes {
134 match &mut cte.body {
135 super::CtePlanBody::Query(query) => rewrite_query_scalars(query, rewrite),
136 super::CtePlanBody::Command(command) => {
137 rewrite_command_scalars(command, rewrite);
138 }
139 }
140 }
141 for row in &mut plan.rows {
142 for expression in row {
143 rewrite_scalar(expression, rewrite);
144 }
145 }
146 if let Some(source) = &mut plan.source {
147 rewrite_query_scalars(source, rewrite);
148 }
149 if let Some(conflict) = &mut plan.on_conflict {
150 for expression in &mut conflict.expressions {
151 rewrite_scalar(expression, rewrite);
152 }
153 if let Some(predicate) = &mut conflict.predicate {
154 rewrite_scalar(predicate, rewrite);
155 }
156 if let ConflictActionPlan::Update {
157 assignments,
158 predicate,
159 } = &mut conflict.action
160 {
161 rewrite_assignments(assignments, rewrite);
162 if let Some(predicate) = predicate {
163 rewrite_scalar(predicate, rewrite);
164 }
165 }
166 }
167 rewrite_projections(&mut plan.returning, rewrite);
168 for check in &mut plan.view_checks {
169 rewrite_scalar(&mut check.predicate, rewrite);
170 }
171 rewrite_subqueries(&mut plan.subqueries, rewrite);
172 }
173 CommandPlan::Update(plan) => {
174 for cte in &mut plan.ctes {
175 match &mut cte.body {
176 super::CtePlanBody::Query(query) => rewrite_query_scalars(query, rewrite),
177 super::CtePlanBody::Command(command) => {
178 rewrite_command_scalars(command, rewrite);
179 }
180 }
181 }
182 if let Some(source) = &mut plan.source {
183 rewrite_source_scalars(source, rewrite);
184 }
185 rewrite_assignments(&mut plan.assignments, rewrite);
186 rewrite_optional_scalar(&mut plan.predicate, rewrite);
187 rewrite_projections(&mut plan.returning, rewrite);
188 for check in &mut plan.view_checks {
189 rewrite_scalar(&mut check.predicate, rewrite);
190 }
191 rewrite_subqueries(&mut plan.subqueries, rewrite);
192 }
193 CommandPlan::Delete(plan) => {
194 for cte in &mut plan.ctes {
195 match &mut cte.body {
196 super::CtePlanBody::Query(query) => rewrite_query_scalars(query, rewrite),
197 super::CtePlanBody::Command(command) => {
198 rewrite_command_scalars(command, rewrite);
199 }
200 }
201 }
202 if let Some(source) = &mut plan.source {
203 rewrite_source_scalars(source, rewrite);
204 }
205 rewrite_optional_scalar(&mut plan.predicate, rewrite);
206 rewrite_projections(&mut plan.returning, rewrite);
207 rewrite_subqueries(&mut plan.subqueries, rewrite);
208 }
209 CommandPlan::Merge(plan) => {
210 for cte in &mut plan.ctes {
211 match &mut cte.body {
212 super::CtePlanBody::Query(query) => rewrite_query_scalars(query, rewrite),
213 super::CtePlanBody::Command(command) => {
214 rewrite_command_scalars(command, rewrite);
215 }
216 }
217 }
218 rewrite_source_scalars(&mut plan.source, rewrite);
219 rewrite_optional_scalar(&mut plan.target_predicate, rewrite);
220 rewrite_scalar(&mut plan.join_condition, rewrite);
221 for clause in &mut plan.when_clauses {
222 match clause {
223 MergeWhenPlan::UpdateMatched {
224 condition,
225 assignments,
226 }
227 | MergeWhenPlan::UpdateNotMatchedBySource {
228 condition,
229 assignments,
230 } => {
231 rewrite_optional_scalar(condition, rewrite);
232 rewrite_assignments(assignments, rewrite);
233 }
234 MergeWhenPlan::DeleteMatched { condition }
235 | MergeWhenPlan::DeleteNotMatchedBySource { condition }
236 | MergeWhenPlan::NothingMatched { condition }
237 | MergeWhenPlan::NothingNotMatched { condition }
238 | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
239 rewrite_optional_scalar(condition, rewrite);
240 }
241 MergeWhenPlan::InsertNotMatched {
242 condition, values, ..
243 } => {
244 rewrite_optional_scalar(condition, rewrite);
245 for value in values {
246 rewrite_scalar(value, rewrite);
247 }
248 }
249 }
250 }
251 rewrite_projections(&mut plan.returning, rewrite);
252 for check in &mut plan.view_checks {
253 rewrite_scalar(&mut check.predicate, rewrite);
254 }
255 rewrite_subqueries(&mut plan.subqueries, rewrite);
256 }
257 CommandPlan::CreateView { query, .. }
258 | CommandPlan::CreateMaterializedView { query, .. }
259 | CommandPlan::CreateTableAs { query, .. }
260 | CommandPlan::DeclareCursor { query, .. } => {
261 rewrite_query_scalars(query, rewrite);
262 }
263 CommandPlan::Explain { body, .. } | CommandPlan::Prepare { body, .. } => {
264 body.rewrite_scalar_expressions(rewrite);
265 }
266 CommandPlan::Execute { params, .. } | CommandPlan::Call { args: params, .. } => {
267 for expression in params {
268 rewrite_scalar(&mut expression.scalar, rewrite);
269 rewrite_subqueries(&mut expression.subqueries, rewrite);
270 }
271 }
272 CommandPlan::CreateTable(_)
273 | CommandPlan::CreateTableIfNotExists(_)
274 | CommandPlan::CreateIndex(_)
275 | CommandPlan::Drop(_)
276 | CommandPlan::AlterTable(_)
277 | CommandPlan::AlterForeignTable(_)
278 | CommandPlan::AlterView(_)
279 | CommandPlan::RefreshMaterializedView { .. }
280 | CommandPlan::CreateSchema { .. }
281 | CommandPlan::AlterSchemaOwner { .. }
282 | CommandPlan::Notify { .. }
283 | CommandPlan::Listen { .. }
284 | CommandPlan::Unlisten { .. }
285 | CommandPlan::SetVariable { .. }
286 | CommandPlan::ResetVariable { .. }
287 | CommandPlan::ResetAllVariables
288 | CommandPlan::SetConstraints { .. }
289 | CommandPlan::ShowVariable { .. }
290 | CommandPlan::Discard { .. }
291 | CommandPlan::Load { .. }
292 | CommandPlan::Analyze { .. }
293 | CommandPlan::Vacuum(_)
294 | CommandPlan::Truncate { .. }
295 | CommandPlan::Transaction(_)
296 | CommandPlan::FetchCursor(_)
297 | CommandPlan::CloseCursor { .. }
298 | CommandPlan::CreateSequence(_)
299 | CommandPlan::CreateDomain(_)
300 | CommandPlan::AlterSequence(_)
301 | CommandPlan::Deallocate { .. }
302 | CommandPlan::CreateForeignServer(_)
303 | CommandPlan::CreateForeignTable(_)
304 | CommandPlan::CreateForeignTableIfNotExists(_)
305 | CommandPlan::CreateFunction(_)
306 | CommandPlan::DropFunction(_)
307 | CommandPlan::AlterRoutine(_)
308 | CommandPlan::AlterRoutineOwner(_)
309 | CommandPlan::RenameRoutine(_)
310 | CommandPlan::GrantRoutine(_)
311 | CommandPlan::GrantTable(_)
312 | CommandPlan::GrantSequence(_)
313 | CommandPlan::GrantDatabase(_)
314 | CommandPlan::GrantSchema(_)
315 | CommandPlan::GrantRole(_)
316 | CommandPlan::CreateRole(_)
317 | CommandPlan::AlterRole(_)
318 | CommandPlan::DropRole(_)
319 | CommandPlan::CreateTrigger(_)
320 | CommandPlan::DropTrigger(_)
321 | CommandPlan::CreateRule(_)
322 | CommandPlan::DropRule(_)
323 | CommandPlan::DoBlock { .. } => {}
324 }
325}
326
327pub(super) fn rewrite_assignments(
328 assignments: &mut [AssignmentPlan],
329 rewrite: &mut dyn FnMut(&mut ScalarExpr),
330) {
331 for assignment in assignments {
332 rewrite_scalar(&mut assignment.value, rewrite);
333 }
334}
335
336pub(super) fn rewrite_projections(
337 projections: &mut [ProjectionPlan],
338 rewrite: &mut dyn FnMut(&mut ScalarExpr),
339) {
340 for projection in projections {
341 rewrite_scalar(&mut projection.expr, rewrite);
342 }
343}
344
345pub(super) fn rewrite_orders(orders: &mut [OrderPlan], rewrite: &mut dyn FnMut(&mut ScalarExpr)) {
346 for order in orders {
347 rewrite_scalar(&mut order.expr, rewrite);
348 }
349}
350
351pub(super) fn rewrite_subqueries(
352 subqueries: &mut [QueryPlan],
353 rewrite: &mut dyn FnMut(&mut ScalarExpr),
354) {
355 for subquery in subqueries {
356 rewrite_query_scalars(subquery, rewrite);
357 }
358}
359
360pub(super) fn rewrite_optional_scalar(
361 expression: &mut Option<ScalarExpr>,
362 rewrite: &mut dyn FnMut(&mut ScalarExpr),
363) {
364 if let Some(expression) = expression {
365 rewrite_scalar(expression, rewrite);
366 }
367}
368
369pub(super) fn rewrite_scalar(
370 expression: &mut ScalarExpr,
371 rewrite: &mut dyn FnMut(&mut ScalarExpr),
372) {
373 match expression {
374 ScalarExpr::Func {
375 args,
376 order_by,
377 filter,
378 ..
379 } => {
380 for argument in args {
381 rewrite_scalar(argument, rewrite);
382 }
383 for order in order_by {
384 rewrite_scalar(&mut order.expr, rewrite);
385 }
386 if let Some(filter) = filter {
387 rewrite_scalar(filter, rewrite);
388 }
389 }
390 ScalarExpr::Array(items)
391 | ScalarExpr::Row(items)
392 | ScalarExpr::And(items)
393 | ScalarExpr::Or(items) => {
394 for item in items {
395 rewrite_scalar(item, rewrite);
396 }
397 }
398 ScalarExpr::Binary { lhs, rhs, .. } => {
399 rewrite_scalar(lhs, rewrite);
400 rewrite_scalar(rhs, rewrite);
401 }
402 ScalarExpr::UnaryMinus(inner)
403 | ScalarExpr::Not(inner)
404 | ScalarExpr::IsNull { expr: inner, .. }
405 | ScalarExpr::Cast { expr: inner, .. } => rewrite_scalar(inner, rewrite),
406 ScalarExpr::Between { expr, low, high } => {
407 rewrite_scalar(expr, rewrite);
408 rewrite_scalar(low, rewrite);
409 rewrite_scalar(high, rewrite);
410 }
411 ScalarExpr::InList { expr, list, .. } => {
412 rewrite_scalar(expr, rewrite);
413 for item in list {
414 rewrite_scalar(item, rewrite);
415 }
416 }
417 ScalarExpr::WindowCall { args, spec, .. } => {
418 for argument in args {
419 rewrite_scalar(argument, rewrite);
420 }
421 for expression in &mut spec.partition_by {
422 rewrite_scalar(expression, rewrite);
423 }
424 for order in &mut spec.order_by {
425 rewrite_scalar(&mut order.expr, rewrite);
426 }
427 if let Some(frame) = &mut spec.frame {
428 rewrite_frame_bound(&mut frame.start, rewrite);
429 rewrite_frame_bound(&mut frame.end, rewrite);
430 }
431 }
432 ScalarExpr::Case {
433 base,
434 when,
435 else_branch,
436 } => {
437 if let Some(base) = base {
438 rewrite_scalar(base, rewrite);
439 }
440 for (condition, result) in when {
441 rewrite_scalar(condition, rewrite);
442 rewrite_scalar(result, rewrite);
443 }
444 if let Some(branch) = else_branch {
445 rewrite_scalar(branch, rewrite);
446 }
447 }
448 ScalarExpr::InSubquery { expr, .. } => rewrite_scalar(expr, rewrite),
449 ScalarExpr::Default
450 | ScalarExpr::Star
451 | ScalarExpr::QualifiedStar(_)
452 | ScalarExpr::Column(_)
453 | ScalarExpr::Position(_)
454 | ScalarExpr::InternalColumn(_)
455 | ScalarExpr::QualifiedColumn { .. }
456 | ScalarExpr::Literal(_)
457 | ScalarExpr::TypedLiteral { .. }
458 | ScalarExpr::Param(_)
459 | ScalarExpr::ScalarSubquery(_)
460 | ScalarExpr::Exists { .. } => {}
461 }
462 rewrite(expression);
463}
464
465pub fn rewrite_scalar_expression(
470 expression: &mut ScalarExpr,
471 rewrite: &mut dyn FnMut(&mut ScalarExpr),
472) {
473 rewrite_scalar(expression, rewrite);
474}
475
476pub(super) fn rewrite_frame_bound(
477 bound: &mut ScalarFrameBound,
478 rewrite: &mut dyn FnMut(&mut ScalarExpr),
479) {
480 match bound {
481 ScalarFrameBound::Preceding(expression) | ScalarFrameBound::Following(expression) => {
482 rewrite_scalar(expression, rewrite);
483 }
484 ScalarFrameBound::UnboundedPreceding
485 | ScalarFrameBound::UnboundedFollowing
486 | ScalarFrameBound::CurrentRow => {}
487 }
488}