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    pub(super) 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                composite_source,
69            } => {
70                let bound_type = self.scalar_type_copy(bound_type.as_ref())?;
71                ScalarExpr::TypedLiteral {
72                    composite_source: self
73                        .composite_source(composite_source.as_ref().map(Source::Borrowed))?,
74                    value: self.value(Source::Borrowed(value))?,
75                    ty: self.copy_text(ty)?,
76                    bound_type,
77                    parameter_index: *parameter_index,
78                }
79            }
80            ScalarExpr::Param(index) => ScalarExpr::Param(*index),
81            ScalarExpr::Func {
82                name,
83                binding,
84                args,
85                distinct,
86                order_by,
87                order_syntax,
88                filter,
89            } => ScalarExpr::Func {
90                name: self.copy_text(name)?,
91                binding: binding
92                    .as_ref()
93                    .map(|binding| self.binding(Source::Borrowed(binding)))
94                    .transpose()?,
95                args: self.map(args.iter(), Self::scalar_copy)?,
96                distinct: *distinct,
97                order_by: self.map(order_by.iter(), Self::scalar_order_copy)?,
98                order_syntax: *order_syntax,
99                filter: self.scalar_optional_copy(filter.as_deref())?,
100            },
101            ScalarExpr::Array(items) => {
102                ScalarExpr::Array(self.map(items.iter(), Self::scalar_copy)?)
103            }
104            ScalarExpr::CompositeRow {
105                items,
106                binding,
107                bound_type,
108            } => ScalarExpr::CompositeRow {
109                bound_type: self.scalar_type_copy(bound_type.as_ref())?,
110                items: self.map(items.iter(), Self::scalar_copy)?,
111                binding: self.composite_binding(Source::Borrowed(binding))?,
112            },
113            ScalarExpr::Row(items) => ScalarExpr::Row(self.map(items.iter(), Self::scalar_copy)?),
114            ScalarExpr::Binary { op, lhs, rhs } => ScalarExpr::Binary {
115                op: *op,
116                lhs: self.scalar_box_copy(lhs)?,
117                rhs: self.scalar_box_copy(rhs)?,
118            },
119            ScalarExpr::UnaryMinus(expression) => {
120                ScalarExpr::UnaryMinus(self.scalar_box_copy(expression)?)
121            }
122            ScalarExpr::Not(expression) => ScalarExpr::Not(self.scalar_box_copy(expression)?),
123            ScalarExpr::And(items) => ScalarExpr::And(self.map(items.iter(), Self::scalar_copy)?),
124            ScalarExpr::Or(items) => ScalarExpr::Or(self.map(items.iter(), Self::scalar_copy)?),
125            ScalarExpr::IsNull { expr, negated } => ScalarExpr::IsNull {
126                expr: self.scalar_box_copy(expr)?,
127                negated: *negated,
128            },
129            ScalarExpr::Between { expr, low, high } => ScalarExpr::Between {
130                expr: self.scalar_box_copy(expr)?,
131                low: self.scalar_box_copy(low)?,
132                high: self.scalar_box_copy(high)?,
133            },
134            ScalarExpr::InList {
135                expr,
136                list,
137                negated,
138            } => ScalarExpr::InList {
139                expr: self.scalar_box_copy(expr)?,
140                list: self.map(list.iter(), Self::scalar_copy)?,
141                negated: *negated,
142            },
143            ScalarExpr::WindowCall {
144                name,
145                args,
146                spec,
147                filter,
148                modifiers,
149            } => ScalarExpr::WindowCall {
150                name: self.copy_text(name)?,
151                args: self.map(args.iter(), Self::scalar_copy)?,
152                spec: ScalarWindowSpec {
153                    definition: spec.definition,
154                    partition_by: self.map(spec.partition_by.iter(), Self::scalar_copy)?,
155                    order_by: self.map(spec.order_by.iter(), Self::scalar_order_copy)?,
156                    frame: spec
157                        .frame
158                        .as_ref()
159                        .map(|frame| {
160                            Ok::<_, crate::schema::retention::CatalogRetentionError>(
161                                ScalarWindowFrame {
162                                    mode: frame.mode,
163                                    start: self.scalar_bound_copy(&frame.start)?,
164                                    end: self.scalar_bound_copy(&frame.end)?,
165                                    between: frame.between,
166                                    exclusion: frame.exclusion,
167                                },
168                            )
169                        })
170                        .transpose()?,
171                },
172                filter: self.scalar_optional_copy(filter.as_deref())?,
173                modifiers: *modifiers,
174            },
175            ScalarExpr::Case {
176                base,
177                when,
178                else_branch,
179            } => ScalarExpr::Case {
180                base: self.scalar_optional_copy(base.as_deref())?,
181                when: self.map(when.iter(), |this, (condition, value)| {
182                    Ok((this.scalar_copy(condition)?, this.scalar_copy(value)?))
183                })?,
184                else_branch: self.scalar_optional_copy(else_branch.as_deref())?,
185            },
186            ScalarExpr::Cast { implicit, expr, ty } => ScalarExpr::Cast {
187                implicit: *implicit,
188                expr: self.scalar_box_copy(expr)?,
189                ty: self.copy_text(ty)?,
190            },
191            ScalarExpr::ScalarSubquery(query) => ScalarExpr::ScalarSubquery(*query),
192            ScalarExpr::Exists { subquery, negated } => ScalarExpr::Exists {
193                subquery: *subquery,
194                negated: *negated,
195            },
196            ScalarExpr::InSubquery {
197                expr,
198                subquery,
199                negated,
200            } => ScalarExpr::InSubquery {
201                expr: self.scalar_box_copy(expr)?,
202                subquery: *subquery,
203                negated: *negated,
204            },
205        })
206    }
207
208    fn scalar_box_copy(&mut self, expression: &ScalarExpr) -> Result<Box<ScalarExpr>> {
209        self.boxed(|this| this.scalar_copy(expression))
210    }
211
212    fn scalar_optional_copy(
213        &mut self,
214        expression: Option<&ScalarExpr>,
215    ) -> Result<Option<Box<ScalarExpr>>> {
216        expression
217            .map(|expression| self.scalar_box_copy(expression))
218            .transpose()
219    }
220
221    fn scalar_order_copy(&mut self, order: &ScalarOrder) -> Result<ScalarOrder> {
222        Ok(ScalarOrder {
223            expr: self.scalar_copy(&order.expr)?,
224            descending: order.descending,
225            nulls: order.nulls,
226        })
227    }
228
229    fn scalar_bound_copy(&mut self, bound: &ScalarFrameBound) -> Result<ScalarFrameBound> {
230        Ok(match bound {
231            ScalarFrameBound::UnboundedPreceding => ScalarFrameBound::UnboundedPreceding,
232            ScalarFrameBound::UnboundedFollowing => ScalarFrameBound::UnboundedFollowing,
233            ScalarFrameBound::CurrentRow => ScalarFrameBound::CurrentRow,
234            ScalarFrameBound::Preceding(expression) => {
235                ScalarFrameBound::Preceding(self.scalar_box_copy(expression)?)
236            }
237            ScalarFrameBound::Following(expression) => {
238                ScalarFrameBound::Following(self.scalar_box_copy(expression)?)
239            }
240        })
241    }
242}
243
244#[cfg(test)]
245mod tests;