1use super::{ScalarExpr, ScalarFrameBound};
10
11impl ScalarExpr {
12 pub fn visit(&self, visitor: &mut impl FnMut(&Self)) {
14 visitor(self);
15 match self {
16 Self::And(parts) | Self::Or(parts) | Self::Array(parts) | Self::Row(parts) => {
17 for part in parts {
18 part.visit(visitor);
19 }
20 }
21 Self::Not(inner)
22 | Self::UnaryMinus(inner)
23 | Self::Cast { expr: inner, .. }
24 | Self::IsNull { expr: inner, .. }
25 | Self::InSubquery { expr: inner, .. } => inner.visit(visitor),
26 Self::Binary { lhs, rhs, .. } => {
27 lhs.visit(visitor);
28 rhs.visit(visitor);
29 }
30 Self::Between { expr, low, high } => {
31 expr.visit(visitor);
32 low.visit(visitor);
33 high.visit(visitor);
34 }
35 Self::InList { expr, list, .. } => {
36 expr.visit(visitor);
37 for part in list {
38 part.visit(visitor);
39 }
40 }
41 Self::Func {
42 args,
43 order_by,
44 filter,
45 ..
46 } => {
47 for argument in args {
48 argument.visit(visitor);
49 }
50 for order in order_by {
51 order.expr.visit(visitor);
52 }
53 if let Some(filter) = filter {
54 filter.visit(visitor);
55 }
56 }
57 Self::WindowCall { args, spec, .. } => {
58 for argument in args {
59 argument.visit(visitor);
60 }
61 for partition in &spec.partition_by {
62 partition.visit(visitor);
63 }
64 for order in &spec.order_by {
65 order.expr.visit(visitor);
66 }
67 if let Some(frame) = &spec.frame {
68 for bound in [&frame.start, &frame.end] {
69 match bound {
70 ScalarFrameBound::Preceding(expression)
71 | ScalarFrameBound::Following(expression) => expression.visit(visitor),
72 ScalarFrameBound::UnboundedPreceding
73 | ScalarFrameBound::UnboundedFollowing
74 | ScalarFrameBound::CurrentRow => {}
75 }
76 }
77 }
78 }
79 Self::Case {
80 base,
81 when,
82 else_branch,
83 } => {
84 if let Some(base) = base {
85 base.visit(visitor);
86 }
87 for (condition, result) in when {
88 condition.visit(visitor);
89 result.visit(visitor);
90 }
91 if let Some(else_branch) = else_branch {
92 else_branch.visit(visitor);
93 }
94 }
95 Self::Default
96 | Self::Star
97 | Self::QualifiedStar(_)
98 | Self::Column(_)
99 | Self::Position(_)
100 | Self::InternalColumn(_)
101 | Self::QualifiedColumn { .. }
102 | Self::Literal(_)
103 | Self::TypedLiteral { .. }
104 | Self::Param(_)
105 | Self::ScalarSubquery(_)
106 | Self::Exists { .. } => {}
107 }
108 }
109
110 pub fn collect_columns(&self, output: &mut std::collections::BTreeSet<String>) -> bool {
112 match self {
113 Self::Column(name) | Self::QualifiedColumn { column: name, .. } => {
114 output.insert(name.clone());
115 true
116 }
117 Self::Literal(_)
118 | Self::TypedLiteral { .. }
119 | Self::Param(_)
120 | Self::InternalColumn(_) => true,
121 Self::Func {
122 args,
123 order_by,
124 filter,
125 ..
126 } => {
127 args.iter().all(|arg| arg.collect_columns(output))
128 && order_by
129 .iter()
130 .all(|order| order.expr.collect_columns(output))
131 && filter
132 .as_deref()
133 .is_none_or(|filter| filter.collect_columns(output))
134 }
135 Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
136 items.iter().all(|item| item.collect_columns(output))
137 }
138 Self::Binary { lhs, rhs, .. } => {
139 lhs.collect_columns(output) && rhs.collect_columns(output)
140 }
141 Self::UnaryMinus(expr)
142 | Self::Not(expr)
143 | Self::IsNull { expr, .. }
144 | Self::Cast { expr, .. } => expr.collect_columns(output),
145 Self::Between { expr, low, high } => {
146 expr.collect_columns(output)
147 && low.collect_columns(output)
148 && high.collect_columns(output)
149 }
150 Self::InList { expr, list, .. } => {
151 expr.collect_columns(output) && list.iter().all(|item| item.collect_columns(output))
152 }
153 Self::Case {
154 base,
155 when,
156 else_branch,
157 } => {
158 base.as_deref()
159 .is_none_or(|base| base.collect_columns(output))
160 && when.iter().all(|(condition, result)| {
161 condition.collect_columns(output) && result.collect_columns(output)
162 })
163 && else_branch
164 .as_deref()
165 .is_none_or(|branch| branch.collect_columns(output))
166 }
167 Self::Default
168 | Self::Star
169 | Self::QualifiedStar(_)
170 | Self::Position(_)
171 | Self::WindowCall { .. }
172 | Self::ScalarSubquery(_)
173 | Self::Exists { .. }
174 | Self::InSubquery { .. } => false,
175 }
176 }
177
178 #[must_use]
179 pub fn contains_window(&self) -> bool {
180 match self {
181 Self::WindowCall { .. } => true,
182 Self::Func {
183 args,
184 order_by,
185 filter,
186 ..
187 } => {
188 args.iter().any(Self::contains_window)
189 || order_by.iter().any(|order| order.expr.contains_window())
190 || filter.as_deref().is_some_and(Self::contains_window)
191 }
192 Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
193 items.iter().any(Self::contains_window)
194 }
195 Self::Binary { lhs, rhs, .. } => lhs.contains_window() || rhs.contains_window(),
196 Self::UnaryMinus(expr)
197 | Self::Not(expr)
198 | Self::IsNull { expr, .. }
199 | Self::Cast { expr, .. }
200 | Self::InSubquery { expr, .. } => expr.contains_window(),
201 Self::Between { expr, low, high } => {
202 expr.contains_window() || low.contains_window() || high.contains_window()
203 }
204 Self::InList { expr, list, .. } => {
205 expr.contains_window() || list.iter().any(Self::contains_window)
206 }
207 Self::Case {
208 base,
209 when,
210 else_branch,
211 } => {
212 base.as_deref().is_some_and(Self::contains_window)
213 || when.iter().any(|(condition, result)| {
214 condition.contains_window() || result.contains_window()
215 })
216 || else_branch.as_deref().is_some_and(Self::contains_window)
217 }
218 Self::Default
219 | Self::Star
220 | Self::QualifiedStar(_)
221 | Self::Column(_)
222 | Self::QualifiedColumn { .. }
223 | Self::Position(_)
224 | Self::InternalColumn(_)
225 | Self::Literal(_)
226 | Self::TypedLiteral { .. }
227 | Self::Param(_)
228 | Self::ScalarSubquery(_)
229 | Self::Exists { .. } => false,
230 }
231 }
232
233 #[must_use]
234 pub fn contains_subquery(&self) -> bool {
235 match self {
236 Self::ScalarSubquery(_) | Self::Exists { .. } | Self::InSubquery { .. } => true,
237 Self::Func {
238 args,
239 order_by,
240 filter,
241 ..
242 } => {
243 args.iter().any(Self::contains_subquery)
244 || order_by.iter().any(|order| order.expr.contains_subquery())
245 || filter.as_deref().is_some_and(Self::contains_subquery)
246 }
247 Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
248 items.iter().any(Self::contains_subquery)
249 }
250 Self::Binary { lhs, rhs, .. } => lhs.contains_subquery() || rhs.contains_subquery(),
251 Self::UnaryMinus(expr)
252 | Self::Not(expr)
253 | Self::IsNull { expr, .. }
254 | Self::Cast { expr, .. } => expr.contains_subquery(),
255 Self::Between { expr, low, high } => {
256 expr.contains_subquery() || low.contains_subquery() || high.contains_subquery()
257 }
258 Self::InList { expr, list, .. } => {
259 expr.contains_subquery() || list.iter().any(Self::contains_subquery)
260 }
261 Self::WindowCall { args, spec, .. } => {
262 args.iter().any(Self::contains_subquery)
263 || spec.partition_by.iter().any(Self::contains_subquery)
264 || spec
265 .order_by
266 .iter()
267 .any(|order| order.expr.contains_subquery())
268 || spec.frame.as_ref().is_some_and(|frame| {
269 frame_has(&frame.start, Self::contains_subquery)
270 || frame_has(&frame.end, Self::contains_subquery)
271 })
272 }
273 Self::Case {
274 base,
275 when,
276 else_branch,
277 } => {
278 base.as_deref().is_some_and(Self::contains_subquery)
279 || when.iter().any(|(condition, result)| {
280 condition.contains_subquery() || result.contains_subquery()
281 })
282 || else_branch.as_deref().is_some_and(Self::contains_subquery)
283 }
284 Self::Default
285 | Self::Star
286 | Self::QualifiedStar(_)
287 | Self::Column(_)
288 | Self::QualifiedColumn { .. }
289 | Self::Position(_)
290 | Self::InternalColumn(_)
291 | Self::Literal(_)
292 | Self::TypedLiteral { .. }
293 | Self::Param(_) => false,
294 }
295 }
296
297 #[must_use]
298 pub fn contains_parameter(&self) -> bool {
299 match self {
300 Self::Param(_) => true,
301 Self::Func {
302 args,
303 order_by,
304 filter,
305 ..
306 } => {
307 args.iter().any(Self::contains_parameter)
308 || order_by.iter().any(|order| order.expr.contains_parameter())
309 || filter.as_deref().is_some_and(Self::contains_parameter)
310 }
311 Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
312 items.iter().any(Self::contains_parameter)
313 }
314 Self::Binary { lhs, rhs, .. } => lhs.contains_parameter() || rhs.contains_parameter(),
315 Self::UnaryMinus(expr)
316 | Self::Not(expr)
317 | Self::IsNull { expr, .. }
318 | Self::Cast { expr, .. }
319 | Self::InSubquery { expr, .. } => expr.contains_parameter(),
320 Self::Between { expr, low, high } => {
321 expr.contains_parameter() || low.contains_parameter() || high.contains_parameter()
322 }
323 Self::InList { expr, list, .. } => {
324 expr.contains_parameter() || list.iter().any(Self::contains_parameter)
325 }
326 Self::WindowCall { args, spec, .. } => {
327 args.iter().any(Self::contains_parameter)
328 || spec.partition_by.iter().any(Self::contains_parameter)
329 || spec
330 .order_by
331 .iter()
332 .any(|order| order.expr.contains_parameter())
333 || spec.frame.as_ref().is_some_and(|frame| {
334 frame_has(&frame.start, Self::contains_parameter)
335 || frame_has(&frame.end, Self::contains_parameter)
336 })
337 }
338 Self::Case {
339 base,
340 when,
341 else_branch,
342 } => {
343 base.as_deref().is_some_and(Self::contains_parameter)
344 || when.iter().any(|(condition, result)| {
345 condition.contains_parameter() || result.contains_parameter()
346 })
347 || else_branch.as_deref().is_some_and(Self::contains_parameter)
348 }
349 Self::Default
350 | Self::Star
351 | Self::QualifiedStar(_)
352 | Self::Column(_)
353 | Self::QualifiedColumn { .. }
354 | Self::Position(_)
355 | Self::InternalColumn(_)
356 | Self::Literal(_)
357 | Self::TypedLiteral { .. }
358 | Self::ScalarSubquery(_)
359 | Self::Exists { .. } => false,
360 }
361 }
362
363 #[must_use]
364 pub fn contains_aggregate(&self, is_aggregate: &dyn Fn(&str) -> bool) -> bool {
365 match self {
366 Self::Func {
367 name,
368 args,
369 order_by,
370 filter,
371 ..
372 } => {
373 is_aggregate(name)
374 || args
375 .iter()
376 .any(|expression| expression.contains_aggregate(is_aggregate))
377 || order_by
378 .iter()
379 .any(|order| order.expr.contains_aggregate(is_aggregate))
380 || filter
381 .as_deref()
382 .is_some_and(|expression| expression.contains_aggregate(is_aggregate))
383 }
384 Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => items
385 .iter()
386 .any(|expression| expression.contains_aggregate(is_aggregate)),
387 Self::Binary { lhs, rhs, .. } => {
388 lhs.contains_aggregate(is_aggregate) || rhs.contains_aggregate(is_aggregate)
389 }
390 Self::UnaryMinus(expr)
391 | Self::Not(expr)
392 | Self::IsNull { expr, .. }
393 | Self::Cast { expr, .. }
394 | Self::InSubquery { expr, .. } => expr.contains_aggregate(is_aggregate),
395 Self::Between { expr, low, high } => {
396 expr.contains_aggregate(is_aggregate)
397 || low.contains_aggregate(is_aggregate)
398 || high.contains_aggregate(is_aggregate)
399 }
400 Self::InList { expr, list, .. } => {
401 expr.contains_aggregate(is_aggregate)
402 || list
403 .iter()
404 .any(|item| item.contains_aggregate(is_aggregate))
405 }
406 Self::Case {
407 base,
408 when,
409 else_branch,
410 } => {
411 base.as_deref()
412 .is_some_and(|expression| expression.contains_aggregate(is_aggregate))
413 || when.iter().any(|(condition, result)| {
414 condition.contains_aggregate(is_aggregate)
415 || result.contains_aggregate(is_aggregate)
416 })
417 || else_branch
418 .as_deref()
419 .is_some_and(|expression| expression.contains_aggregate(is_aggregate))
420 }
421 Self::Default
422 | Self::Star
423 | Self::QualifiedStar(_)
424 | Self::Column(_)
425 | Self::QualifiedColumn { .. }
426 | Self::Position(_)
427 | Self::InternalColumn(_)
428 | Self::Literal(_)
429 | Self::TypedLiteral { .. }
430 | Self::Param(_)
431 | Self::ScalarSubquery(_)
432 | Self::Exists { .. }
433 | Self::WindowCall { .. } => false,
434 }
435 }
436}
437
438fn frame_has(bound: &ScalarFrameBound, predicate: fn(&ScalarExpr) -> bool) -> bool {
439 match bound {
440 ScalarFrameBound::Preceding(expression) | ScalarFrameBound::Following(expression) => {
441 predicate(expression)
442 }
443 ScalarFrameBound::UnboundedPreceding
444 | ScalarFrameBound::UnboundedFollowing
445 | ScalarFrameBound::CurrentRow => false,
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use super::{ScalarExpr, ScalarFrameBound};
452 use crate::ast::FrameMode;
453 use uqa_core::Value;
454
455 #[test]
456 fn visit_includes_root_and_nested_expressions() {
457 let expression = ScalarExpr::Binary {
458 op: crate::ast::BinaryOp::Add,
459 lhs: Box::new(ScalarExpr::Column("amount".into())),
460 rhs: Box::new(ScalarExpr::Literal(Value::Int(1))),
461 };
462 let mut visited = Vec::new();
463 expression.visit(&mut |part| visited.push(part.clone()));
464 assert_eq!(visited.len(), 3);
465 assert_eq!(visited[0], expression);
466 }
467
468 #[test]
469 fn traversal_includes_window_frame_expressions() {
470 let expression = ScalarExpr::WindowCall {
471 name: "sum".into(),
472 args: vec![ScalarExpr::Column("amount".into())],
473 spec: super::super::ScalarWindowSpec {
474 partition_by: vec![ScalarExpr::QualifiedColumn {
475 qualifier: "orders".into(),
476 column: "account_id".into(),
477 }],
478 order_by: Vec::new(),
479 frame: Some(super::super::ScalarWindowFrame {
480 mode: FrameMode::Rows,
481 start: ScalarFrameBound::Preceding(Box::new(ScalarExpr::Param(0))),
482 end: ScalarFrameBound::CurrentRow,
483 }),
484 },
485 };
486 let mut visited_parameter = false;
487 expression.visit(&mut |part| {
488 visited_parameter |= matches!(part, ScalarExpr::Param(0));
489 });
490 assert!(visited_parameter);
491 assert!(expression.contains_window());
492 assert!(expression.contains_parameter());
493 }
494
495 #[test]
496 fn owned_walkers_preserve_column_and_aggregate_policy() {
497 let expression = ScalarExpr::Func {
498 name: "sum".into(),
499 binding: None,
500 args: vec![ScalarExpr::QualifiedColumn {
501 qualifier: "orders".into(),
502 column: "amount".into(),
503 }],
504 distinct: false,
505 order_by: Vec::new(),
506 filter: None,
507 };
508 let mut columns = std::collections::BTreeSet::new();
509 assert!(expression.collect_columns(&mut columns));
510 assert_eq!(columns, std::collections::BTreeSet::from(["amount".into()]));
511 assert!(expression.contains_aggregate(&|name| name == "sum"));
512 assert!(!expression.contains_subquery());
513 }
514}