uqa_sql/ir/
traversal.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Complete scalar IR traversal.
8
9use super::{ScalarExpr, ScalarFrameBound};
10
11impl ScalarExpr {
12    /// Visit this expression and every nested scalar expression in pre-order.
13    pub fn visit(&self, visitor: &mut impl FnMut(&Self)) {
14        self.try_visit(&mut |expression| {
15            visitor(expression);
16            Ok::<_, std::convert::Infallible>(true)
17        })
18        .unwrap_or_else(|never| match never {});
19    }
20
21    /// Visit in pre-order, skipping a subtree when the visitor returns false and stopping immediately on its first error.
22    pub fn try_visit<E>(
23        &self,
24        visitor: &mut impl FnMut(&Self) -> Result<bool, E>,
25    ) -> Result<(), E> {
26        if !visitor(self)? {
27            return Ok(());
28        }
29        match self {
30            Self::And(parts) | Self::Or(parts) | Self::Array(parts) | Self::Row(parts) => {
31                for part in parts {
32                    part.try_visit(visitor)?;
33                }
34            }
35            Self::Not(inner)
36            | Self::UnaryMinus(inner)
37            | Self::Cast { expr: inner, .. }
38            | Self::IsNull { expr: inner, .. }
39            | Self::InSubquery { expr: inner, .. } => inner.try_visit(visitor)?,
40            Self::Binary { lhs, rhs, .. } => {
41                lhs.try_visit(visitor)?;
42                rhs.try_visit(visitor)?;
43            }
44            Self::Between { expr, low, high } => {
45                expr.try_visit(visitor)?;
46                low.try_visit(visitor)?;
47                high.try_visit(visitor)?;
48            }
49            Self::InList { expr, list, .. } => {
50                expr.try_visit(visitor)?;
51                for part in list {
52                    part.try_visit(visitor)?;
53                }
54            }
55            Self::Func {
56                args,
57                order_by,
58                filter,
59                ..
60            } => {
61                for argument in args {
62                    argument.try_visit(visitor)?;
63                }
64                for order in order_by {
65                    order.expr.try_visit(visitor)?;
66                }
67                if let Some(filter) = filter {
68                    filter.try_visit(visitor)?;
69                }
70            }
71            Self::WindowCall { args, spec, .. } => {
72                for argument in args {
73                    argument.try_visit(visitor)?;
74                }
75                for partition in &spec.partition_by {
76                    partition.try_visit(visitor)?;
77                }
78                for order in &spec.order_by {
79                    order.expr.try_visit(visitor)?;
80                }
81                if let Some(frame) = &spec.frame {
82                    for bound in [&frame.start, &frame.end] {
83                        match bound {
84                            ScalarFrameBound::Preceding(expression)
85                            | ScalarFrameBound::Following(expression) => {
86                                expression.try_visit(visitor)?;
87                            }
88                            ScalarFrameBound::UnboundedPreceding
89                            | ScalarFrameBound::UnboundedFollowing
90                            | ScalarFrameBound::CurrentRow => {}
91                        }
92                    }
93                }
94            }
95            Self::Case {
96                base,
97                when,
98                else_branch,
99            } => {
100                if let Some(base) = base {
101                    base.try_visit(visitor)?;
102                }
103                for (condition, result) in when {
104                    condition.try_visit(visitor)?;
105                    result.try_visit(visitor)?;
106                }
107                if let Some(else_branch) = else_branch {
108                    else_branch.try_visit(visitor)?;
109                }
110            }
111            Self::Default
112            | Self::Star
113            | Self::QualifiedStar(_)
114            | Self::Column(_)
115            | Self::Position(_)
116            | Self::InternalColumn(_)
117            | Self::QualifiedColumn { .. }
118            | Self::Literal(_)
119            | Self::TypedLiteral { .. }
120            | Self::Param(_)
121            | Self::ScalarSubquery(_)
122            | Self::Exists { .. } => {}
123        }
124        Ok(())
125    }
126
127    /// Collect every column needed to evaluate this expression. Returns `false` when evaluation needs row shape or a relational child that a projected field scan cannot provide.
128    pub fn collect_columns(&self, output: &mut std::collections::BTreeSet<String>) -> bool {
129        match self.try_visit_columns(&mut |name| {
130            output.insert(name.to_owned());
131            Ok::<_, std::convert::Infallible>(())
132        }) {
133            Ok(projectable) => projectable,
134            Err(never) => match never {},
135        }
136    }
137
138    /// Borrow referenced column names in evaluation-tree order, stopping at the first visitor failure or unprojectable expression. Qualified references yield their column component, matching `collect_columns`; repeated references remain visible to the visitor. No name or result container is allocated by this traversal.
139    pub fn try_visit_columns<'a, E>(
140        &'a self,
141        visitor: &mut impl FnMut(&'a str) -> Result<(), E>,
142    ) -> Result<bool, E> {
143        match self {
144            Self::Column(name) | Self::QualifiedColumn { column: name, .. } => {
145                visitor(name)?;
146                Ok(true)
147            }
148            Self::Literal(_)
149            | Self::TypedLiteral { .. }
150            | Self::Param(_)
151            | Self::InternalColumn(_) => Ok(true),
152            Self::Func {
153                args,
154                order_by,
155                filter,
156                ..
157            } => {
158                for expression in args
159                    .iter()
160                    .chain(order_by.iter().map(|order| &order.expr))
161                    .chain(filter.as_deref())
162                {
163                    if !expression.try_visit_columns(visitor)? {
164                        return Ok(false);
165                    }
166                }
167                Ok(true)
168            }
169            Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
170                for item in items {
171                    if !item.try_visit_columns(visitor)? {
172                        return Ok(false);
173                    }
174                }
175                Ok(true)
176            }
177            Self::Binary { lhs, rhs, .. } => {
178                Ok(lhs.try_visit_columns(visitor)? && rhs.try_visit_columns(visitor)?)
179            }
180            Self::UnaryMinus(expr)
181            | Self::Not(expr)
182            | Self::IsNull { expr, .. }
183            | Self::Cast { expr, .. } => expr.try_visit_columns(visitor),
184            Self::Between { expr, low, high } => Ok(expr.try_visit_columns(visitor)?
185                && low.try_visit_columns(visitor)?
186                && high.try_visit_columns(visitor)?),
187            Self::InList { expr, list, .. } => {
188                for item in std::iter::once(expr.as_ref()).chain(list) {
189                    if !item.try_visit_columns(visitor)? {
190                        return Ok(false);
191                    }
192                }
193                Ok(true)
194            }
195            Self::Case {
196                base,
197                when,
198                else_branch,
199            } => {
200                for expression in base
201                    .as_deref()
202                    .into_iter()
203                    .chain(
204                        when.iter()
205                            .flat_map(|(condition, result)| [condition, result]),
206                    )
207                    .chain(else_branch.as_deref())
208                {
209                    if !expression.try_visit_columns(visitor)? {
210                        return Ok(false);
211                    }
212                }
213                Ok(true)
214            }
215            Self::Default
216            | Self::Star
217            | Self::QualifiedStar(_)
218            | Self::Position(_)
219            | Self::WindowCall { .. }
220            | Self::ScalarSubquery(_)
221            | Self::Exists { .. }
222            | Self::InSubquery { .. } => Ok(false),
223        }
224    }
225
226    #[must_use]
227    pub fn contains_window(&self) -> bool {
228        match self {
229            Self::WindowCall { .. } => true,
230            Self::Func {
231                args,
232                order_by,
233                filter,
234                ..
235            } => {
236                args.iter().any(Self::contains_window)
237                    || order_by.iter().any(|order| order.expr.contains_window())
238                    || filter.as_deref().is_some_and(Self::contains_window)
239            }
240            Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
241                items.iter().any(Self::contains_window)
242            }
243            Self::Binary { lhs, rhs, .. } => lhs.contains_window() || rhs.contains_window(),
244            Self::UnaryMinus(expr)
245            | Self::Not(expr)
246            | Self::IsNull { expr, .. }
247            | Self::Cast { expr, .. }
248            | Self::InSubquery { expr, .. } => expr.contains_window(),
249            Self::Between { expr, low, high } => {
250                expr.contains_window() || low.contains_window() || high.contains_window()
251            }
252            Self::InList { expr, list, .. } => {
253                expr.contains_window() || list.iter().any(Self::contains_window)
254            }
255            Self::Case {
256                base,
257                when,
258                else_branch,
259            } => {
260                base.as_deref().is_some_and(Self::contains_window)
261                    || when.iter().any(|(condition, result)| {
262                        condition.contains_window() || result.contains_window()
263                    })
264                    || else_branch.as_deref().is_some_and(Self::contains_window)
265            }
266            Self::Default
267            | Self::Star
268            | Self::QualifiedStar(_)
269            | Self::Column(_)
270            | Self::QualifiedColumn { .. }
271            | Self::Position(_)
272            | Self::InternalColumn(_)
273            | Self::Literal(_)
274            | Self::TypedLiteral { .. }
275            | Self::Param(_)
276            | Self::ScalarSubquery(_)
277            | Self::Exists { .. } => false,
278        }
279    }
280
281    #[must_use]
282    pub fn contains_subquery(&self) -> bool {
283        match self {
284            Self::ScalarSubquery(_) | Self::Exists { .. } | Self::InSubquery { .. } => true,
285            Self::Func {
286                args,
287                order_by,
288                filter,
289                ..
290            } => {
291                args.iter().any(Self::contains_subquery)
292                    || order_by.iter().any(|order| order.expr.contains_subquery())
293                    || filter.as_deref().is_some_and(Self::contains_subquery)
294            }
295            Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
296                items.iter().any(Self::contains_subquery)
297            }
298            Self::Binary { lhs, rhs, .. } => lhs.contains_subquery() || rhs.contains_subquery(),
299            Self::UnaryMinus(expr)
300            | Self::Not(expr)
301            | Self::IsNull { expr, .. }
302            | Self::Cast { expr, .. } => expr.contains_subquery(),
303            Self::Between { expr, low, high } => {
304                expr.contains_subquery() || low.contains_subquery() || high.contains_subquery()
305            }
306            Self::InList { expr, list, .. } => {
307                expr.contains_subquery() || list.iter().any(Self::contains_subquery)
308            }
309            Self::WindowCall { args, spec, .. } => {
310                args.iter().any(Self::contains_subquery)
311                    || spec.partition_by.iter().any(Self::contains_subquery)
312                    || spec
313                        .order_by
314                        .iter()
315                        .any(|order| order.expr.contains_subquery())
316                    || spec.frame.as_ref().is_some_and(|frame| {
317                        frame_has(&frame.start, Self::contains_subquery)
318                            || frame_has(&frame.end, Self::contains_subquery)
319                    })
320            }
321            Self::Case {
322                base,
323                when,
324                else_branch,
325            } => {
326                base.as_deref().is_some_and(Self::contains_subquery)
327                    || when.iter().any(|(condition, result)| {
328                        condition.contains_subquery() || result.contains_subquery()
329                    })
330                    || else_branch.as_deref().is_some_and(Self::contains_subquery)
331            }
332            Self::Default
333            | Self::Star
334            | Self::QualifiedStar(_)
335            | Self::Column(_)
336            | Self::QualifiedColumn { .. }
337            | Self::Position(_)
338            | Self::InternalColumn(_)
339            | Self::Literal(_)
340            | Self::TypedLiteral { .. }
341            | Self::Param(_) => false,
342        }
343    }
344
345    #[must_use]
346    pub fn contains_parameter(&self) -> bool {
347        match self {
348            Self::Param(_) => true,
349            Self::Func {
350                args,
351                order_by,
352                filter,
353                ..
354            } => {
355                args.iter().any(Self::contains_parameter)
356                    || order_by.iter().any(|order| order.expr.contains_parameter())
357                    || filter.as_deref().is_some_and(Self::contains_parameter)
358            }
359            Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => {
360                items.iter().any(Self::contains_parameter)
361            }
362            Self::Binary { lhs, rhs, .. } => lhs.contains_parameter() || rhs.contains_parameter(),
363            Self::UnaryMinus(expr)
364            | Self::Not(expr)
365            | Self::IsNull { expr, .. }
366            | Self::Cast { expr, .. }
367            | Self::InSubquery { expr, .. } => expr.contains_parameter(),
368            Self::Between { expr, low, high } => {
369                expr.contains_parameter() || low.contains_parameter() || high.contains_parameter()
370            }
371            Self::InList { expr, list, .. } => {
372                expr.contains_parameter() || list.iter().any(Self::contains_parameter)
373            }
374            Self::WindowCall { args, spec, .. } => {
375                args.iter().any(Self::contains_parameter)
376                    || spec.partition_by.iter().any(Self::contains_parameter)
377                    || spec
378                        .order_by
379                        .iter()
380                        .any(|order| order.expr.contains_parameter())
381                    || spec.frame.as_ref().is_some_and(|frame| {
382                        frame_has(&frame.start, Self::contains_parameter)
383                            || frame_has(&frame.end, Self::contains_parameter)
384                    })
385            }
386            Self::Case {
387                base,
388                when,
389                else_branch,
390            } => {
391                base.as_deref().is_some_and(Self::contains_parameter)
392                    || when.iter().any(|(condition, result)| {
393                        condition.contains_parameter() || result.contains_parameter()
394                    })
395                    || else_branch.as_deref().is_some_and(Self::contains_parameter)
396            }
397            Self::Default
398            | Self::Star
399            | Self::QualifiedStar(_)
400            | Self::Column(_)
401            | Self::QualifiedColumn { .. }
402            | Self::Position(_)
403            | Self::InternalColumn(_)
404            | Self::Literal(_)
405            | Self::TypedLiteral { .. }
406            | Self::ScalarSubquery(_)
407            | Self::Exists { .. } => false,
408        }
409    }
410
411    #[must_use]
412    pub fn contains_aggregate(&self, is_aggregate: &dyn Fn(&str) -> bool) -> bool {
413        match self {
414            Self::Func {
415                name,
416                args,
417                order_by,
418                filter,
419                ..
420            } => {
421                is_aggregate(name)
422                    || args
423                        .iter()
424                        .any(|expression| expression.contains_aggregate(is_aggregate))
425                    || order_by
426                        .iter()
427                        .any(|order| order.expr.contains_aggregate(is_aggregate))
428                    || filter
429                        .as_deref()
430                        .is_some_and(|expression| expression.contains_aggregate(is_aggregate))
431            }
432            Self::Array(items) | Self::Row(items) | Self::And(items) | Self::Or(items) => items
433                .iter()
434                .any(|expression| expression.contains_aggregate(is_aggregate)),
435            Self::Binary { lhs, rhs, .. } => {
436                lhs.contains_aggregate(is_aggregate) || rhs.contains_aggregate(is_aggregate)
437            }
438            Self::UnaryMinus(expr)
439            | Self::Not(expr)
440            | Self::IsNull { expr, .. }
441            | Self::Cast { expr, .. }
442            | Self::InSubquery { expr, .. } => expr.contains_aggregate(is_aggregate),
443            Self::Between { expr, low, high } => {
444                expr.contains_aggregate(is_aggregate)
445                    || low.contains_aggregate(is_aggregate)
446                    || high.contains_aggregate(is_aggregate)
447            }
448            Self::InList { expr, list, .. } => {
449                expr.contains_aggregate(is_aggregate)
450                    || list
451                        .iter()
452                        .any(|item| item.contains_aggregate(is_aggregate))
453            }
454            Self::Case {
455                base,
456                when,
457                else_branch,
458            } => {
459                base.as_deref()
460                    .is_some_and(|expression| expression.contains_aggregate(is_aggregate))
461                    || when.iter().any(|(condition, result)| {
462                        condition.contains_aggregate(is_aggregate)
463                            || result.contains_aggregate(is_aggregate)
464                    })
465                    || else_branch
466                        .as_deref()
467                        .is_some_and(|expression| expression.contains_aggregate(is_aggregate))
468            }
469            Self::Default
470            | Self::Star
471            | Self::QualifiedStar(_)
472            | Self::Column(_)
473            | Self::QualifiedColumn { .. }
474            | Self::Position(_)
475            | Self::InternalColumn(_)
476            | Self::Literal(_)
477            | Self::TypedLiteral { .. }
478            | Self::Param(_)
479            | Self::ScalarSubquery(_)
480            | Self::Exists { .. }
481            | Self::WindowCall { .. } => false,
482        }
483    }
484}
485
486fn frame_has(bound: &ScalarFrameBound, predicate: fn(&ScalarExpr) -> bool) -> bool {
487    match bound {
488        ScalarFrameBound::Preceding(expression) | ScalarFrameBound::Following(expression) => {
489            predicate(expression)
490        }
491        ScalarFrameBound::UnboundedPreceding
492        | ScalarFrameBound::UnboundedFollowing
493        | ScalarFrameBound::CurrentRow => false,
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::{ScalarExpr, ScalarFrameBound};
500    use crate::ast::FrameMode;
501    use uqa_core::Value;
502
503    #[test]
504    fn visit_includes_root_and_nested_expressions() {
505        let expression = ScalarExpr::Binary {
506            op: crate::ast::BinaryOp::Add,
507            lhs: Box::new(ScalarExpr::Column("amount".into())),
508            rhs: Box::new(ScalarExpr::Literal(Value::Int(1))),
509        };
510        let mut visited = Vec::new();
511        expression.visit(&mut |part| visited.push(part.clone()));
512        assert_eq!(visited.len(), 3);
513        assert_eq!(visited[0], expression);
514    }
515
516    #[test]
517    fn fallible_visits_skip_selected_subtrees_and_stop_before_later_siblings() {
518        let expression = ScalarExpr::Row(vec![
519            ScalarExpr::Array(vec![ScalarExpr::Column("hidden".into())]),
520            ScalarExpr::Column("reject".into()),
521            ScalarExpr::Column("unvisited".into()),
522        ]);
523        let mut visited = Vec::new();
524        let result = expression.try_visit(&mut |part| {
525            visited.push(part.clone());
526            match part {
527                ScalarExpr::Array(_) => Ok(false),
528                ScalarExpr::Column(name) if name == "reject" => Err("grouping"),
529                _ => Ok(true),
530            }
531        });
532        assert_eq!(result, Err("grouping"));
533        assert_eq!(visited.len(), 3);
534        assert!(matches!(&visited[2], ScalarExpr::Column(name) if name == "reject"));
535    }
536
537    #[test]
538    fn traversal_includes_window_frame_expressions() {
539        let expression = ScalarExpr::WindowCall {
540            name: "sum".into(),
541            args: vec![ScalarExpr::Column("amount".into())],
542            spec: super::super::ScalarWindowSpec {
543                partition_by: vec![ScalarExpr::QualifiedColumn {
544                    qualifier: "orders".into(),
545                    column: "account_id".into(),
546                }],
547                order_by: Vec::new(),
548                frame: Some(super::super::ScalarWindowFrame {
549                    mode: FrameMode::Rows,
550                    start: ScalarFrameBound::Preceding(Box::new(ScalarExpr::Param(0))),
551                    end: ScalarFrameBound::CurrentRow,
552                }),
553            },
554        };
555        let mut visited_parameter = false;
556        expression.visit(&mut |part| {
557            visited_parameter |= matches!(part, ScalarExpr::Param(0));
558        });
559        assert!(visited_parameter);
560        assert!(expression.contains_window());
561        assert!(expression.contains_parameter());
562    }
563
564    #[test]
565    fn owned_walkers_preserve_column_and_aggregate_policy() {
566        let expression = ScalarExpr::Func {
567            name: "sum".into(),
568            binding: None,
569            args: vec![ScalarExpr::QualifiedColumn {
570                qualifier: "orders".into(),
571                column: "amount".into(),
572            }],
573            distinct: false,
574            order_by: Vec::new(),
575            filter: None,
576        };
577        let mut columns = std::collections::BTreeSet::new();
578        assert!(expression.collect_columns(&mut columns));
579        assert_eq!(columns, std::collections::BTreeSet::from(["amount".into()]));
580        assert!(expression.contains_aggregate(&|name| name == "sum"));
581        assert!(!expression.contains_subquery());
582    }
583
584    #[test]
585    fn borrowed_column_visits_keep_names_and_stop_at_the_first_failure() {
586        let expression = ScalarExpr::Row(vec![
587            ScalarExpr::Column("first".into()),
588            ScalarExpr::QualifiedColumn {
589                qualifier: "table".into(),
590                column: "second".into(),
591            },
592            ScalarExpr::Column("first".into()),
593        ]);
594        let mut borrowed = Vec::new();
595        assert!(expression
596            .try_visit_columns(&mut |name| {
597                borrowed.push(name);
598                Ok::<_, &str>(())
599            })
600            .unwrap());
601        assert_eq!(borrowed, ["first", "second", "first"]);
602        let ScalarExpr::Row(items) = &expression else {
603            unreachable!()
604        };
605        let ScalarExpr::Column(first) = &items[0] else {
606            unreachable!()
607        };
608        assert_eq!(borrowed[0].as_ptr(), first.as_ptr());
609        let mut visits = 0;
610        let result = expression.try_visit_columns(&mut |_| {
611            visits += 1;
612            if visits == 2 {
613                Err("quota")
614            } else {
615                Ok(())
616            }
617        });
618        assert_eq!(result, Err("quota"));
619        assert_eq!(visits, 2);
620    }
621
622    #[test]
623    fn borrowed_column_visits_preserve_unprojectable_prefix_semantics() {
624        let expression = ScalarExpr::Array(vec![
625            ScalarExpr::Column("before".into()),
626            ScalarExpr::Position(0),
627            ScalarExpr::Column("after".into()),
628        ]);
629        let mut borrowed = Vec::new();
630        assert!(!expression
631            .try_visit_columns(&mut |name| {
632                borrowed.push(name);
633                Ok::<_, &str>(())
634            })
635            .unwrap());
636        let mut owned = std::collections::BTreeSet::new();
637        assert!(!expression.collect_columns(&mut owned));
638        assert_eq!(borrowed, ["before"]);
639        assert_eq!(owned, std::collections::BTreeSet::from(["before".into()]));
640    }
641}
uqa_sql/ir/traversal.rs

uqa_sql/ir/
traversal.rs
