uqa_sql/routines/
inlining.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Caller-plan expansion of selected SQL routines. Source text is analyzed for
8//! each planning attempt; declining expansion never changes the session body cache.
9
10mod arguments;
11mod operators;
12mod properties;
13#[cfg(test)]
14mod tests;
15
16use super::{
17    body_parameters::sql_body_parameter_scope,
18    body_validation::routine_parameter_values,
19    compilation::{RoutineCompilationCatalog, RoutineParserCatalog},
20    declaration::RoutineTypeCatalog,
21    result_check::{sql_function_result_layout, SQLFunctionResultKind},
22    security::{ensure_routine_execute_privilege, RoutineExecutionAuthority},
23    CompiledFunctionBody, RoutineBody, RoutineResolution,
24};
25use crate::{
26    ast::{CreateFunction, FunctionBinding, FunctionBody, FunctionReturns},
27    binding::{bind_routine_parameter_references, prepare_routine_body_inputs},
28    plan::{ComputePlan, QueryBlockPlan, RelationalPlan, UnifiedPlan},
29    semantics::volatility::VolatilityCatalog,
30    ColumnType, SQLError, ScalarExpr,
31};
32
33/// Metadata and analysis capabilities borrowed for one optimizer invocation.
34pub trait RoutinePlanExpressions {
35    /// Execute only an immutable call whose inputs the planner proved constant.
36    /// The existing runtime boundary retains authorization and session body caches.
37    fn evaluate_constant_routine(
38        &self,
39        expression: &ScalarExpr,
40    ) -> Result<uqa_core::Value, SQLError>;
41}
42
43#[derive(Clone, Copy)]
44pub struct RoutineInliningContext<'a> {
45    pub routines: &'a dyn RoutineResolution,
46    pub types: &'a dyn RoutineTypeCatalog,
47    pub parsers: &'a dyn RoutineParserCatalog,
48    pub catalog: &'a dyn RoutineCompilationCatalog,
49    pub authority: &'a dyn RoutineExecutionAuthority,
50    pub volatility: &'a dyn VolatilityCatalog,
51    pub expressions: &'a dyn RoutinePlanExpressions,
52}
53
54impl std::fmt::Debug for RoutineInliningContext<'_> {
55    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        formatter
57            .debug_struct("RoutineInliningContext")
58            .finish_non_exhaustive()
59    }
60}
61
62/// A typed replacement, together with the routine identity whose recursive
63/// expansion the planner must suppress while simplifying the replacement.
64pub struct InlineRoutineExpression {
65    pub identity: [u8; 16],
66    pub expression: ScalarExpr,
67}
68
69impl RoutineInliningContext<'_> {
70    /// Classify an analyzed/planned expression in its own row scope, including casts and selected operators.
71    pub fn expression_volatility(
72        &self,
73        expression: &ScalarExpr,
74        schema: &crate::RowSchema,
75    ) -> Result<crate::ast::FunctionVolatility, SQLError> {
76        properties::volatility(self, expression, schema)
77    }
78
79    /// Try the selected routine only; a same-named overload cannot replace it.
80    /// Argument simplification and lazy branches belong to the caller's planner.
81    pub fn prepare(
82        &self,
83        binding: &FunctionBinding,
84        arguments: &[ScalarExpr],
85        active: &[[u8; 16]],
86    ) -> Result<Option<InlineRoutineExpression>, SQLError> {
87        if binding.builtin || binding.resolution_error.is_some() {
88            return Ok(None);
89        }
90        let Some(identity) = binding.object_id else {
91            return Ok(None);
92        };
93        let Some(function) = self.selected_function(binding) else {
94            return Ok(None);
95        };
96        let Some(invocation) = binding.invocation.as_deref() else {
97            return Ok(None);
98        };
99        let specialized = super::invocation::specialized_definition(&function.def, invocation)?;
100        let definition = specialized.as_ref().unwrap_or(&function.def);
101        let arguments_for_evaluation = arguments;
102        let Some(arguments) = arguments::prepare(self, definition, invocation, arguments)? else {
103            return Ok(None);
104        };
105        if let Some(expression) =
106            self.constant_call(definition, binding, &arguments, arguments_for_evaluation)?
107        {
108            return Ok(Some(InlineRoutineExpression {
109                identity,
110                expression,
111            }));
112        }
113        if !eligible_definition(definition) || active.contains(&identity) {
114            return Ok(None);
115        }
116        match ensure_routine_execute_privilege(self.authority, definition) {
117            Ok(()) => {}
118            Err(error) if error.sqlstate() == Some("42501") => return Ok(None),
119            Err(error) => return Err(error),
120        }
121        let Some((mut expression, parameters)) = self.analyze_body(&function, definition)? else {
122            return Ok(None);
123        };
124        let properties = properties::inspect(self, &expression, &parameters)?;
125        if !properties.permits(definition) {
126            return Ok(None);
127        }
128        let uses = arguments::use_counts(&expression, arguments.len())?;
129        for (argument, count) in arguments.iter().zip(uses) {
130            if count == 0 && definition.strict {
131                return Ok(None);
132            }
133            if count > 1 && !properties::can_duplicate(self, argument)? {
134                return Ok(None);
135            }
136        }
137        arguments::substitute(&mut expression, &arguments);
138        Ok(Some(InlineRoutineExpression {
139            identity,
140            expression,
141        }))
142    }
143
144    fn constant_call(
145        &self,
146        definition: &CreateFunction,
147        binding: &FunctionBinding,
148        arguments: &[ScalarExpr],
149        arguments_for_evaluation: &[ScalarExpr],
150    ) -> Result<Option<ScalarExpr>, SQLError> {
151        if !definition.is_procedure
152            && !definition.returns_set()
153            && definition.output_params().len() <= 1
154            && !super::routine_returns_anonymous_record(definition)
155        {
156            let layout = super::result_check::declared_sql_function_result(self.types, definition)?;
157            if definition.strict && arguments.iter().any(is_null_constant) {
158                return Ok(Some(ScalarExpr::TypedLiteral {
159                    collation_oid: None,
160                    composite_source: None,
161                    value: uqa_core::Value::Null,
162                    ty: layout.declared_type.catalog_name(),
163                    bound_type: Some(layout.declared_type),
164                    parameter_index: None,
165                }));
166            }
167            if definition.volatility == crate::ast::FunctionVolatility::Immutable
168                && arguments.iter().all(is_constant)
169            {
170                let call = ScalarExpr::Func {
171                    name: binding.name.clone(),
172                    binding: Some(binding.clone()),
173                    args: arguments_for_evaluation.to_vec(),
174                    distinct: false,
175                    order_by: Vec::new(),
176                    order_syntax: crate::ast::FunctionCallSyntax::Ordinary,
177                    filter: None,
178                };
179                let value = self.expressions.evaluate_constant_routine(&call)?;
180                return Ok(Some(ScalarExpr::TypedLiteral {
181                    collation_oid: None,
182                    composite_source: None,
183                    value,
184                    ty: layout.declared_type.catalog_name(),
185                    bound_type: Some(layout.declared_type),
186                    parameter_index: None,
187                }));
188            }
189        }
190        Ok(None)
191    }
192
193    fn analyze_body(
194        &self,
195        function: &super::SQLUserFunction,
196        definition: &CreateFunction,
197    ) -> Result<Option<(ScalarExpr, Vec<crate::SQLParam>)>, SQLError> {
198        let Some(mut plan) = self.body_plan(function)? else {
199            return Ok(None);
200        };
201        let parameters = routine_parameter_values(self.types, definition);
202        let names = matches!(definition.body, FunctionBody::Source(_))
203            .then(|| sql_body_parameter_scope(definition, &parameters))
204            .transpose()?;
205        let snapshot = self.catalog.binding_snapshot()?;
206        let result = prepare_routine_body_inputs(
207            self.routines,
208            &mut plan,
209            &parameters,
210            &snapshot.context(),
211            self.catalog,
212            names.as_ref(),
213        )?;
214        if let Some(names) = &names {
215            bind_routine_parameter_references(
216                self.routines,
217                &mut plan,
218                &parameters,
219                &snapshot.context(),
220                names,
221            )?;
222        }
223        let UnifiedPlan::Query(query) = plan else {
224            return Ok(None);
225        };
226        if !query.ctes.is_empty() {
227            return Ok(None);
228        }
229        let RelationalPlan::QueryBlock(mut block) = query.root else {
230            return Ok(None);
231        };
232        if !simple_select(&block) {
233            return Ok(None);
234        }
235        let layout = sql_function_result_layout(self.types, definition, Some(&result))?;
236        if layout.kind == SQLFunctionResultKind::Tuple {
237            return Ok(None);
238        }
239        let mut expression = block.projections.remove(0).expr;
240        let actual = result
241            .column_types()
242            .and_then(|types| types.first())
243            .and_then(Option::as_ref);
244        if layout.kind == SQLFunctionResultKind::Void && actual != Some(&ColumnType::Void) {
245            return Ok(None);
246        }
247        if actual != Some(&layout.declared_type) {
248            expression = ScalarExpr::Cast {
249                implicit: true,
250                expr: Box::new(expression),
251                ty: layout.declared_type.catalog_name(),
252            };
253        }
254        Ok(Some((expression, parameters)))
255    }
256
257    /// Expand defaults, named arguments and implicit casts before the planner
258    /// simplifies children. Returning an exact invocation prevents a second
259    /// overload selection and keeps discarded or repeated arguments typed.
260    pub fn materialize_call(
261        &self,
262        binding: &FunctionBinding,
263        arguments: &[ScalarExpr],
264    ) -> Result<Option<(FunctionBinding, Vec<ScalarExpr>)>, SQLError> {
265        let Some(function) = self.selected_function(binding) else {
266            return Ok(None);
267        };
268        let Some(invocation) = binding.invocation.as_deref() else {
269            return Ok(None);
270        };
271        let specialized = super::invocation::specialized_definition(&function.def, invocation)?;
272        let definition = specialized.as_ref().unwrap_or(&function.def);
273        let Some(mut arguments) = arguments::prepare(self, definition, invocation, arguments)?
274        else {
275            return Ok(None);
276        };
277        let mut selected = binding.clone();
278        let invocation = selected.invocation.as_mut().expect("selected invocation");
279        invocation.argument_positions = definition
280            .params
281            .iter()
282            .enumerate()
283            .filter_map(|(index, parameter)| {
284                super::body_parameters::is_sql_body_parameter(parameter).then_some(index)
285            })
286            .collect();
287        invocation.argument_targets = invocation
288            .argument_positions
289            .iter()
290            .map(|index| invocation.parameter_types[*index].clone())
291            .collect();
292        invocation.argument_sources = invocation
293            .argument_targets
294            .iter()
295            .cloned()
296            .map(Some)
297            .collect();
298        if let crate::ast::RoutineVariadicMode::Expanded { parameter_index } =
299            invocation.variadic_mode
300        {
301            invocation.variadic_mode =
302                crate::ast::RoutineVariadicMode::Explicit { parameter_index };
303        }
304        if let crate::ast::RoutineVariadicMode::Explicit { parameter_index } =
305            invocation.variadic_mode
306        {
307            let position = invocation
308                .argument_positions
309                .iter()
310                .position(|index| *index == parameter_index)
311                .ok_or_else(|| SQLError::Internal("variadic call lost its argument".into()))?;
312            let value = std::mem::replace(
313                &mut arguments[position],
314                ScalarExpr::Literal(uqa_core::Value::Null),
315            );
316            let marker =
317                FunctionBinding::dispatched(crate::ast::FunctionDispatch::VariadicArgument);
318            arguments[position] = ScalarExpr::Func {
319                name: marker.name.clone(),
320                binding: Some(marker),
321                args: vec![value],
322                distinct: false,
323                order_by: Vec::new(),
324                filter: None,
325                order_syntax: crate::ast::FunctionCallSyntax::Ordinary,
326            };
327        }
328        Ok(Some((selected, arguments)))
329    }
330
331    fn selected_function(
332        &self,
333        binding: &FunctionBinding,
334    ) -> Option<std::sync::Arc<super::SQLUserFunction>> {
335        if binding.builtin || binding.resolution_error.is_some() {
336            return None;
337        }
338        let identity = binding.object_id?;
339        self.routines
340            .lookup_bound_sql_functions_by_binding(binding)?
341            .into_iter()
342            .find(|function| {
343                !function.def.is_procedure
344                    && function.def.object_id == Some(identity)
345                    && super::routine_signature_types(&function.def) == binding.argument_types
346            })
347    }
348
349    /// Ordinary messages repeat optimizer-time routine work. Prepared plans
350    /// intentionally keep their own custom/generic planning lifetime instead.
351    pub fn requires_replanning(&self, plan: &UnifiedPlan) -> bool {
352        let mut required = false;
353        plan.visit_scalar_expressions(&mut |root| {
354            root.visit(&mut |expression| {
355                let ScalarExpr::Func {
356                    binding: Some(binding),
357                    ..
358                } = expression
359                else {
360                    return;
361                };
362                if binding.builtin {
363                    return;
364                }
365                required |= self
366                    .routines
367                    .lookup_bound_sql_functions_by_binding(binding)
368                    .is_some_and(|functions| {
369                        functions.iter().any(|function| {
370                            function.def.object_id == binding.object_id
371                                && super::routine_signature_types(&function.def)
372                                    == binding.argument_types
373                                && (eligible_definition(&function.def)
374                                    || function.def.volatility
375                                        == crate::ast::FunctionVolatility::Immutable)
376                        })
377                    });
378            });
379        });
380        required
381    }
382
383    fn body_plan(
384        &self,
385        function: &super::SQLUserFunction,
386    ) -> Result<Option<UnifiedPlan>, SQLError> {
387        match &function.body {
388            RoutineBody::Bound(body) => match body.as_ref() {
389                CompiledFunctionBody::SQL(plans) if plans.len() == 1 => Ok(Some(plans[0].clone())),
390                _ => Ok(None),
391            },
392            RoutineBody::Source => {
393                let FunctionBody::Source(source) = &function.def.body else {
394                    return Err(SQLError::Internal(
395                        "source routine has no SQL source".into(),
396                    ));
397                };
398                // Parse the complete message once, then decline a multi-statement
399                // body before semantic compilation. Lexical notices still belong
400                // to every attempted expansion, as in inline_function.
401                super::compilation::with_parser_context(self.parsers, || {
402                    let statements = crate::parse_statements(source)?;
403                    let [statement] = statements.as_slice() else {
404                        return Ok(None);
405                    };
406                    Ok(Some(UnifiedPlan::lower_with(
407                        statement.compile()?,
408                        &|name: &str| self.catalog.has_registered_aggregate_function(name),
409                    )))
410                })
411            }
412        }
413    }
414}
415
416fn eligible_definition(definition: &CreateFunction) -> bool {
417    definition.language == "sql"
418        && !definition.is_procedure
419        && !definition.security.security_definer
420        && !definition.returns_set()
421        && definition.config.is_empty()
422        && definition.output_params().len() <= 1
423        && !matches!(&definition.returns, FunctionReturns::Scalar { type_name }
424            if crate::type_resolution::canonical_routine_type_name(type_name) == "record")
425}
426
427fn simple_select(block: &QueryBlockPlan) -> bool {
428    matches!(block.compute, ComputePlan::Project)
429        && block.from.is_none()
430        && block.r#where.is_none()
431        && block.group_by.is_empty()
432        && block.grouping_sets.is_empty()
433        && block.having.is_none()
434        && block.windows.is_empty()
435        && !block.distinct
436        && block.distinct_on.is_empty()
437        && block.order_by.is_empty()
438        && block.limit.is_none()
439        && block.offset.is_none()
440        && block.locking.is_empty()
441        && block.subqueries.is_empty()
442        && block.projections.len() == 1
443}
444
445fn is_constant(expression: &ScalarExpr) -> bool {
446    matches!(
447        expression,
448        ScalarExpr::Literal(_) | ScalarExpr::TypedLiteral { .. }
449    )
450}
451
452fn is_null_constant(expression: &ScalarExpr) -> bool {
453    matches!(
454        expression,
455        ScalarExpr::Literal(uqa_core::Value::Null)
456            | ScalarExpr::TypedLiteral {
457                value: uqa_core::Value::Null,
458                ..
459            }
460    )
461}
uqa_sql/routines/inlining.rs

uqa_sql/routines/
inlining.rs
