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                    composite_source: None,
160                    value: uqa_core::Value::Null,
161                    ty: layout.declared_type.catalog_name(),
162                    bound_type: Some(layout.declared_type),
163                    parameter_index: None,
164                }));
165            }
166            if definition.volatility == crate::ast::FunctionVolatility::Immutable
167                && arguments.iter().all(is_constant)
168            {
169                let call = ScalarExpr::Func {
170                    name: binding.name.clone(),
171                    binding: Some(binding.clone()),
172                    args: arguments_for_evaluation.to_vec(),
173                    distinct: false,
174                    order_by: Vec::new(),
175                    order_syntax: crate::ast::FunctionCallSyntax::Ordinary,
176                    filter: None,
177                };
178                let value = self.expressions.evaluate_constant_routine(&call)?;
179                return Ok(Some(ScalarExpr::TypedLiteral {
180                    composite_source: None,
181                    value,
182                    ty: layout.declared_type.catalog_name(),
183                    bound_type: Some(layout.declared_type),
184                    parameter_index: None,
185                }));
186            }
187        }
188        Ok(None)
189    }
190
191    fn analyze_body(
192        &self,
193        function: &super::SQLUserFunction,
194        definition: &CreateFunction,
195    ) -> Result<Option<(ScalarExpr, Vec<crate::SQLParam>)>, SQLError> {
196        let Some(mut plan) = self.body_plan(function)? else {
197            return Ok(None);
198        };
199        let parameters = routine_parameter_values(self.types, definition);
200        let names = matches!(definition.body, FunctionBody::Source(_))
201            .then(|| sql_body_parameter_scope(definition, &parameters))
202            .transpose()?;
203        let snapshot = self.catalog.binding_snapshot()?;
204        let result = prepare_routine_body_inputs(
205            self.routines,
206            &mut plan,
207            &parameters,
208            &snapshot.context(),
209            self.catalog,
210            names.as_ref(),
211        )?;
212        if let Some(names) = &names {
213            bind_routine_parameter_references(
214                self.routines,
215                &mut plan,
216                &parameters,
217                &snapshot.context(),
218                names,
219            )?;
220        }
221        let UnifiedPlan::Query(query) = plan else {
222            return Ok(None);
223        };
224        if !query.ctes.is_empty() {
225            return Ok(None);
226        }
227        let RelationalPlan::QueryBlock(mut block) = query.root else {
228            return Ok(None);
229        };
230        if !simple_select(&block) {
231            return Ok(None);
232        }
233        let layout = sql_function_result_layout(self.types, definition, Some(&result))?;
234        if layout.kind == SQLFunctionResultKind::Tuple {
235            return Ok(None);
236        }
237        let mut expression = block.projections.remove(0).expr;
238        let actual = result
239            .column_types()
240            .and_then(|types| types.first())
241            .and_then(Option::as_ref);
242        if layout.kind == SQLFunctionResultKind::Void && actual != Some(&ColumnType::Void) {
243            return Ok(None);
244        }
245        if actual != Some(&layout.declared_type) {
246            expression = ScalarExpr::Cast {
247                implicit: true,
248                expr: Box::new(expression),
249                ty: layout.declared_type.catalog_name(),
250            };
251        }
252        Ok(Some((expression, parameters)))
253    }
254
255    /// Expand defaults, named arguments and implicit casts before the planner
256    /// simplifies children. Returning an exact invocation prevents a second
257    /// overload selection and keeps discarded or repeated arguments typed.
258    pub fn materialize_call(
259        &self,
260        binding: &FunctionBinding,
261        arguments: &[ScalarExpr],
262    ) -> Result<Option<(FunctionBinding, Vec<ScalarExpr>)>, SQLError> {
263        let Some(function) = self.selected_function(binding) else {
264            return Ok(None);
265        };
266        let Some(invocation) = binding.invocation.as_deref() else {
267            return Ok(None);
268        };
269        let specialized = super::invocation::specialized_definition(&function.def, invocation)?;
270        let definition = specialized.as_ref().unwrap_or(&function.def);
271        let Some(mut arguments) = arguments::prepare(self, definition, invocation, arguments)?
272        else {
273            return Ok(None);
274        };
275        let mut selected = binding.clone();
276        let invocation = selected.invocation.as_mut().expect("selected invocation");
277        invocation.argument_positions = definition
278            .params
279            .iter()
280            .enumerate()
281            .filter_map(|(index, parameter)| {
282                super::body_parameters::is_sql_body_parameter(parameter).then_some(index)
283            })
284            .collect();
285        invocation.argument_targets = invocation
286            .argument_positions
287            .iter()
288            .map(|index| invocation.parameter_types[*index].clone())
289            .collect();
290        invocation.argument_sources = invocation
291            .argument_targets
292            .iter()
293            .cloned()
294            .map(Some)
295            .collect();
296        if let crate::ast::RoutineVariadicMode::Expanded { parameter_index } =
297            invocation.variadic_mode
298        {
299            invocation.variadic_mode =
300                crate::ast::RoutineVariadicMode::Explicit { parameter_index };
301        }
302        if let crate::ast::RoutineVariadicMode::Explicit { parameter_index } =
303            invocation.variadic_mode
304        {
305            let position = invocation
306                .argument_positions
307                .iter()
308                .position(|index| *index == parameter_index)
309                .ok_or_else(|| SQLError::Internal("variadic call lost its argument".into()))?;
310            let value = std::mem::replace(
311                &mut arguments[position],
312                ScalarExpr::Literal(uqa_core::Value::Null),
313            );
314            let marker =
315                FunctionBinding::dispatched(crate::ast::FunctionDispatch::VariadicArgument);
316            arguments[position] = ScalarExpr::Func {
317                name: marker.name.clone(),
318                binding: Some(marker),
319                args: vec![value],
320                distinct: false,
321                order_by: Vec::new(),
322                filter: None,
323                order_syntax: crate::ast::FunctionCallSyntax::Ordinary,
324            };
325        }
326        Ok(Some((selected, arguments)))
327    }
328
329    fn selected_function(
330        &self,
331        binding: &FunctionBinding,
332    ) -> Option<std::sync::Arc<super::SQLUserFunction>> {
333        if binding.builtin || binding.resolution_error.is_some() {
334            return None;
335        }
336        let identity = binding.object_id?;
337        self.routines
338            .lookup_bound_sql_functions_by_binding(binding)?
339            .into_iter()
340            .find(|function| {
341                !function.def.is_procedure
342                    && function.def.object_id == Some(identity)
343                    && super::routine_signature_types(&function.def) == binding.argument_types
344            })
345    }
346
347    /// Ordinary messages repeat optimizer-time routine work. Prepared plans
348    /// intentionally keep their own custom/generic planning lifetime instead.
349    pub fn requires_replanning(&self, plan: &UnifiedPlan) -> bool {
350        let mut required = false;
351        plan.visit_scalar_expressions(&mut |root| {
352            root.visit(&mut |expression| {
353                let ScalarExpr::Func {
354                    binding: Some(binding),
355                    ..
356                } = expression
357                else {
358                    return;
359                };
360                if binding.builtin {
361                    return;
362                }
363                required |= self
364                    .routines
365                    .lookup_bound_sql_functions_by_binding(binding)
366                    .is_some_and(|functions| {
367                        functions.iter().any(|function| {
368                            function.def.object_id == binding.object_id
369                                && super::routine_signature_types(&function.def)
370                                    == binding.argument_types
371                                && (eligible_definition(&function.def)
372                                    || function.def.volatility
373                                        == crate::ast::FunctionVolatility::Immutable)
374                        })
375                    });
376            });
377        });
378        required
379    }
380
381    fn body_plan(
382        &self,
383        function: &super::SQLUserFunction,
384    ) -> Result<Option<UnifiedPlan>, SQLError> {
385        match &function.body {
386            RoutineBody::Bound(body) => match body.as_ref() {
387                CompiledFunctionBody::SQL(plans) if plans.len() == 1 => Ok(Some(plans[0].clone())),
388                _ => Ok(None),
389            },
390            RoutineBody::Source => {
391                let FunctionBody::Source(source) = &function.def.body else {
392                    return Err(SQLError::Internal(
393                        "source routine has no SQL source".into(),
394                    ));
395                };
396                // Parse the complete message once, then decline a multi-statement
397                // body before semantic compilation. Lexical notices still belong
398                // to every attempted expansion, as in inline_function.
399                super::compilation::with_parser_context(self.parsers, || {
400                    let statements = crate::parse_statements(source)?;
401                    let [statement] = statements.as_slice() else {
402                        return Ok(None);
403                    };
404                    Ok(Some(UnifiedPlan::lower_with(
405                        statement.compile()?,
406                        &|name: &str| self.catalog.has_registered_aggregate_function(name),
407                    )))
408                })
409            }
410        }
411    }
412}
413
414fn eligible_definition(definition: &CreateFunction) -> bool {
415    definition.language == "sql"
416        && !definition.is_procedure
417        && !definition.security.security_definer
418        && !definition.returns_set()
419        && definition.config.is_empty()
420        && definition.output_params().len() <= 1
421        && !matches!(&definition.returns, FunctionReturns::Scalar { type_name }
422            if crate::type_resolution::canonical_routine_type_name(type_name) == "record")
423}
424
425fn simple_select(block: &QueryBlockPlan) -> bool {
426    matches!(block.compute, ComputePlan::Project)
427        && block.from.is_none()
428        && block.r#where.is_none()
429        && block.group_by.is_empty()
430        && block.grouping_sets.is_empty()
431        && block.having.is_none()
432        && block.windows.is_empty()
433        && !block.distinct
434        && block.distinct_on.is_empty()
435        && block.order_by.is_empty()
436        && block.limit.is_none()
437        && block.offset.is_none()
438        && block.locking.is_empty()
439        && block.subqueries.is_empty()
440        && block.projections.len() == 1
441}
442
443fn is_constant(expression: &ScalarExpr) -> bool {
444    matches!(
445        expression,
446        ScalarExpr::Literal(_) | ScalarExpr::TypedLiteral { .. }
447    )
448}
449
450fn is_null_constant(expression: &ScalarExpr) -> bool {
451    matches!(
452        expression,
453        ScalarExpr::Literal(uqa_core::Value::Null)
454            | ScalarExpr::TypedLiteral {
455                value: uqa_core::Value::Null,
456                ..
457            }
458    )
459}
uqa_sql/routines/inlining.rs

uqa_sql/routines/
inlining.rs
