Skip to main content

uqa_sql/plan/scalar/
copy.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Admission-aware copies of already analyzed scalar expressions. Query children remain references into their enclosing plan, as in an ordinary scalar clone.
8
9use super::{
10    resources::{Control, Lowering, Result},
11    source::Source,
12    ScalarExpr, ScalarFrameBound, ScalarOrder, ScalarWindowFrame, ScalarWindowSpec,
13};
14use uqa_core::memory::{Produced, ProductionControl};
15
16impl ScalarExpr {
17    /// Copy a scalar tree, admitting every destination allocation before creating it and retaining the complete reservation. Subquery identities still refer to the enclosing plan's query arena.
18    pub fn clone_with_control(&self, control: &ProductionControl<'_>) -> Result<Produced<Self>> {
19        control.check()?;
20        if control.budget().is_none() {
21            return Ok(control.finish(self.clone(), None)?);
22        }
23        let mut copying = Lowering {
24            control: Some(Control::from_production(*control)),
25        };
26        let expression = copying.scalar_copy(self)?;
27        Ok(copying.finish(expression)?.into())
28    }
29}
30
31impl Lowering<'_> {
32    fn scalar_type_copy(
33        &mut self,
34        ty: Option<&crate::ast::ColumnType>,
35    ) -> Result<Option<crate::ast::ColumnType>> {
36        ty.map(|ty| {
37            let control = self.control.as_mut().expect("controlled scalar copy");
38            let (ty, memory) = ty.clone_with_control(&control.production)?.into_parts();
39            control.memory.absorb(memory.expect("controlled type copy"));
40            Ok(ty)
41        })
42        .transpose()
43    }
44
45    #[expect(
46        clippy::too_many_lines,
47        reason = "exhaustive scalar copying shares the lowerer's destination admission"
48    )]
49    fn scalar_copy(&mut self, expression: &ScalarExpr) -> Result<ScalarExpr> {
50        self.check()?;
51        Ok(match expression {
52            ScalarExpr::Star => ScalarExpr::Star,
53            ScalarExpr::QualifiedStar(name) => ScalarExpr::QualifiedStar(self.copy_text(name)?),
54            ScalarExpr::Default => ScalarExpr::Default,
55            ScalarExpr::Column(name) => ScalarExpr::Column(self.copy_text(name)?),
56            ScalarExpr::Position(position) => ScalarExpr::Position(*position),
57            ScalarExpr::InternalColumn(column) => ScalarExpr::InternalColumn(*column),
58            ScalarExpr::QualifiedColumn { qualifier, column } => ScalarExpr::QualifiedColumn {
59                qualifier: self.copy_text(qualifier)?,
60                column: self.copy_text(column)?,
61            },
62            ScalarExpr::Literal(value) => ScalarExpr::Literal(self.value(Source::Borrowed(value))?),
63            ScalarExpr::TypedLiteral {
64                value,
65                ty,
66                bound_type,
67                parameter_index,
68            } => {
69                let bound_type = self.scalar_type_copy(bound_type.as_ref())?;
70                ScalarExpr::TypedLiteral {
71                    value: self.value(Source::Borrowed(value))?,
72                    ty: self.copy_text(ty)?,
73                    bound_type,
74                    parameter_index: *parameter_index,
75                }
76            }
77            ScalarExpr::Param(index) => ScalarExpr::Param(*index),
78            ScalarExpr::Func {
79                name,
80                binding,
81                args,
82                distinct,
83                order_by,
84                order_syntax,
85                filter,
86            } => ScalarExpr::Func {
87                name: self.copy_text(name)?,
88                binding: binding
89                    .as_ref()
90                    .map(|binding| self.binding(Source::Borrowed(binding)))
91                    .transpose()?,
92                args: self.map(args.iter(), Self::scalar_copy)?,
93                distinct: *distinct,
94                order_by: self.map(order_by.iter(), Self::scalar_order_copy)?,
95                order_syntax: *order_syntax,
96                filter: self.scalar_optional_copy(filter.as_deref())?,
97            },
98            ScalarExpr::Array(items) => {
99                ScalarExpr::Array(self.map(items.iter(), Self::scalar_copy)?)
100            }
101            ScalarExpr::CompositeRow {
102                items,
103                binding,
104                bound_type,
105            } => ScalarExpr::CompositeRow {
106                bound_type: self.scalar_type_copy(bound_type.as_ref())?,
107                items: self.map(items.iter(), Self::scalar_copy)?,
108                binding: self.composite_binding(Source::Borrowed(binding))?,
109            },
110            ScalarExpr::Row(items) => ScalarExpr::Row(self.map(items.iter(), Self::scalar_copy)?),
111            ScalarExpr::Binary { op, lhs, rhs } => ScalarExpr::Binary {
112                op: *op,
113                lhs: self.scalar_box_copy(lhs)?,
114                rhs: self.scalar_box_copy(rhs)?,
115            },
116            ScalarExpr::UnaryMinus(expression) => {
117                ScalarExpr::UnaryMinus(self.scalar_box_copy(expression)?)
118            }
119            ScalarExpr::Not(expression) => ScalarExpr::Not(self.scalar_box_copy(expression)?),
120            ScalarExpr::And(items) => ScalarExpr::And(self.map(items.iter(), Self::scalar_copy)?),
121            ScalarExpr::Or(items) => ScalarExpr::Or(self.map(items.iter(), Self::scalar_copy)?),
122            ScalarExpr::IsNull { expr, negated } => ScalarExpr::IsNull {
123                expr: self.scalar_box_copy(expr)?,
124                negated: *negated,
125            },
126            ScalarExpr::Between { expr, low, high } => ScalarExpr::Between {
127                expr: self.scalar_box_copy(expr)?,
128                low: self.scalar_box_copy(low)?,
129                high: self.scalar_box_copy(high)?,
130            },
131            ScalarExpr::InList {
132                expr,
133                list,
134                negated,
135            } => ScalarExpr::InList {
136                expr: self.scalar_box_copy(expr)?,
137                list: self.map(list.iter(), Self::scalar_copy)?,
138                negated: *negated,
139            },
140            ScalarExpr::WindowCall {
141                name,
142                args,
143                spec,
144                filter,
145                modifiers,
146            } => ScalarExpr::WindowCall {
147                name: self.copy_text(name)?,
148                args: self.map(args.iter(), Self::scalar_copy)?,
149                spec: ScalarWindowSpec {
150                    definition: spec.definition,
151                    partition_by: self.map(spec.partition_by.iter(), Self::scalar_copy)?,
152                    order_by: self.map(spec.order_by.iter(), Self::scalar_order_copy)?,
153                    frame: spec
154                        .frame
155                        .as_ref()
156                        .map(|frame| {
157                            Ok::<_, crate::schema::retention::CatalogRetentionError>(
158                                ScalarWindowFrame {
159                                    mode: frame.mode,
160                                    start: self.scalar_bound_copy(&frame.start)?,
161                                    end: self.scalar_bound_copy(&frame.end)?,
162                                    between: frame.between,
163                                    exclusion: frame.exclusion,
164                                },
165                            )
166                        })
167                        .transpose()?,
168                },
169                filter: self.scalar_optional_copy(filter.as_deref())?,
170                modifiers: *modifiers,
171            },
172            ScalarExpr::Case {
173                base,
174                when,
175                else_branch,
176            } => ScalarExpr::Case {
177                base: self.scalar_optional_copy(base.as_deref())?,
178                when: self.map(when.iter(), |this, (condition, value)| {
179                    Ok((this.scalar_copy(condition)?, this.scalar_copy(value)?))
180                })?,
181                else_branch: self.scalar_optional_copy(else_branch.as_deref())?,
182            },
183            ScalarExpr::Cast { implicit, expr, ty } => ScalarExpr::Cast {
184                implicit: *implicit,
185                expr: self.scalar_box_copy(expr)?,
186                ty: self.copy_text(ty)?,
187            },
188            ScalarExpr::ScalarSubquery(query) => ScalarExpr::ScalarSubquery(*query),
189            ScalarExpr::Exists { subquery, negated } => ScalarExpr::Exists {
190                subquery: *subquery,
191                negated: *negated,
192            },
193            ScalarExpr::InSubquery {
194                expr,
195                subquery,
196                negated,
197            } => ScalarExpr::InSubquery {
198                expr: self.scalar_box_copy(expr)?,
199                subquery: *subquery,
200                negated: *negated,
201            },
202        })
203    }
204
205    fn scalar_box_copy(&mut self, expression: &ScalarExpr) -> Result<Box<ScalarExpr>> {
206        self.boxed(|this| this.scalar_copy(expression))
207    }
208
209    fn scalar_optional_copy(
210        &mut self,
211        expression: Option<&ScalarExpr>,
212    ) -> Result<Option<Box<ScalarExpr>>> {
213        expression
214            .map(|expression| self.scalar_box_copy(expression))
215            .transpose()
216    }
217
218    fn scalar_order_copy(&mut self, order: &ScalarOrder) -> Result<ScalarOrder> {
219        Ok(ScalarOrder {
220            expr: self.scalar_copy(&order.expr)?,
221            descending: order.descending,
222            nulls: order.nulls,
223        })
224    }
225
226    fn scalar_bound_copy(&mut self, bound: &ScalarFrameBound) -> Result<ScalarFrameBound> {
227        Ok(match bound {
228            ScalarFrameBound::UnboundedPreceding => ScalarFrameBound::UnboundedPreceding,
229            ScalarFrameBound::UnboundedFollowing => ScalarFrameBound::UnboundedFollowing,
230            ScalarFrameBound::CurrentRow => ScalarFrameBound::CurrentRow,
231            ScalarFrameBound::Preceding(expression) => {
232                ScalarFrameBound::Preceding(self.scalar_box_copy(expression)?)
233            }
234            ScalarFrameBound::Following(expression) => {
235                ScalarFrameBound::Following(self.scalar_box_copy(expression)?)
236            }
237        })
238    }
239}
240
241#[cfg(test)]
242mod tests;