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