chryso_optimizer/
utils.rs

1use chryso_core::ast::{BinaryOperator, Expr};
2use chryso_planner::LogicalPlan;
3use std::collections::HashSet;
4
5pub fn collect_identifiers(expr: &Expr) -> HashSet<String> {
6    let mut idents = HashSet::new();
7    collect_identifiers_inner(expr, &mut idents);
8    idents
9}
10
11pub fn collect_tables(plan: &LogicalPlan) -> HashSet<String> {
12    let mut tables = HashSet::new();
13    collect_tables_inner(plan, &mut tables);
14    tables
15}
16
17pub fn split_conjuncts(expr: &Expr) -> Vec<Expr> {
18    match expr {
19        Expr::BinaryOp { left, op, right } if matches!(op, BinaryOperator::And) => {
20            let mut out = split_conjuncts(left);
21            out.extend(split_conjuncts(right));
22            out
23        }
24        _ => vec![expr.clone()],
25    }
26}
27
28pub fn combine_conjuncts(exprs: Vec<Expr>) -> Option<Expr> {
29    let mut iter = exprs.into_iter();
30    let first = iter.next()?;
31    Some(iter.fold(first, |left, right| Expr::BinaryOp {
32        left: Box::new(left),
33        op: BinaryOperator::And,
34        right: Box::new(right),
35    }))
36}
37
38pub fn table_prefix(ident: &str) -> Option<&str> {
39    ident.split_once('.').map(|(prefix, _)| prefix)
40}
41
42fn collect_identifiers_inner(expr: &Expr, idents: &mut HashSet<String>) {
43    match expr {
44        Expr::Identifier(name) => {
45            idents.insert(name.clone());
46        }
47        Expr::BinaryOp { left, right, .. } => {
48            collect_identifiers_inner(left, idents);
49            collect_identifiers_inner(right, idents);
50        }
51        Expr::IsNull { expr, .. } => {
52            collect_identifiers_inner(expr, idents);
53        }
54        Expr::UnaryOp { expr, .. } => {
55            collect_identifiers_inner(expr, idents);
56        }
57        Expr::FunctionCall { args, .. } => {
58            for arg in args {
59                collect_identifiers_inner(arg, idents);
60            }
61        }
62        Expr::WindowFunction { function, spec } => {
63            collect_identifiers_inner(function, idents);
64            for expr in &spec.partition_by {
65                collect_identifiers_inner(expr, idents);
66            }
67            for expr in &spec.order_by {
68                collect_identifiers_inner(&expr.expr, idents);
69            }
70        }
71        Expr::Subquery(select) | Expr::Exists(select) => {
72            for item in &select.projection {
73                collect_identifiers_inner(&item.expr, idents);
74            }
75        }
76        Expr::InSubquery { expr, subquery } => {
77            collect_identifiers_inner(expr, idents);
78            for item in &subquery.projection {
79                collect_identifiers_inner(&item.expr, idents);
80            }
81        }
82        Expr::Case {
83            operand,
84            when_then,
85            else_expr,
86        } => {
87            if let Some(expr) = operand {
88                collect_identifiers_inner(expr, idents);
89            }
90            for (when_expr, then_expr) in when_then {
91                collect_identifiers_inner(when_expr, idents);
92                collect_identifiers_inner(then_expr, idents);
93            }
94            if let Some(expr) = else_expr {
95                collect_identifiers_inner(expr, idents);
96            }
97        }
98        Expr::Literal(_) | Expr::Wildcard => {}
99    }
100}
101
102fn collect_tables_inner(plan: &LogicalPlan, tables: &mut HashSet<String>) {
103    match plan {
104        LogicalPlan::Scan { table } | LogicalPlan::IndexScan { table, .. } => {
105            tables.insert(table.clone());
106        }
107        LogicalPlan::Dml { .. } => {}
108        LogicalPlan::Derived { input, .. } => collect_tables_inner(input, tables),
109        LogicalPlan::Filter { input, .. }
110        | LogicalPlan::Projection { input, .. }
111        | LogicalPlan::Aggregate { input, .. }
112        | LogicalPlan::Distinct { input }
113        | LogicalPlan::TopN { input, .. }
114        | LogicalPlan::Sort { input, .. }
115        | LogicalPlan::Limit { input, .. } => collect_tables_inner(input, tables),
116        LogicalPlan::Join { left, right, .. } => {
117            collect_tables_inner(left, tables);
118            collect_tables_inner(right, tables);
119        }
120    }
121}
chryso_optimizer/utils.rs

chryso_optimizer/
utils.rs
