Skip to main content

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