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