use std::{collections::HashMap, mem, sync::Arc};
use arrow_array::{Array, ArrayRef, RecordBatch, UInt64Array};
use arrow_schema::FieldRef;
use reifydb_catalog::catalog::{Catalog, procedure::ResolvedProcedure};
use reifydb_core::{
interface::catalog::{
policy::{CallableOp, PolicyTargetType},
procedure::{Procedure, ProcedureParam},
},
internal_error,
value::{
batch::{batch, is_scalar, scalar_value},
column::{
arith::length_mismatch,
builder::ColumnBuilder,
cast::{cast_column_data, convert::Convert},
factory,
},
},
};
use reifydb_evaluate::{
expression::branch::BranchLayout,
stack::{Callable, ClosureValue, Variable, strip_dollar_prefix},
};
use reifydb_policy::inject_from_policies;
use reifydb_routine_abi::{
Function as RoutineFunction, Procedure as RoutineProcedure,
context::{FunctionContext as RoutineFunctionContext, ProcedureContext as RoutineProcedureContext},
};
use reifydb_rql::{
compiler::{CompilationResult, Compiled},
instruction::{CompiledClosure, CompiledFunction, Instruction},
nodes::FunctionParameter,
};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
error::{Error as ReifyError, FunctionErrorKind, ProcedureErrorKind, TypeError},
fragment::Fragment,
params::Params,
reifydb_assertions,
value::{
Value,
column_view::ColumnView,
constraint::TypeConstraint,
frame::frame::Frame,
system_columns::{SystemColumn, system_column, user_columns, with_system_column},
value_type::ValueType,
},
};
use crate::{
Result,
error::EngineError,
policy::PolicyEvaluator,
vm::{
callable::{CallSite, ProcedureCall, invoke_procedure_routine},
services::Services,
stack::ControlFlow,
vm::{EMPTY_PARAMS, UdfCall, Vm},
volcano::udf::is_vectorizable,
},
};
pub(crate) struct CallContext<'a, 'b> {
pub services: &'a Arc<Services>,
pub tx: &'a mut Transaction<'b>,
}
impl<'a> Vm<'a> {
pub(crate) fn coerce_return_value(
&self,
result: Variable,
return_type: Option<&TypeConstraint>,
name: &Fragment,
) -> Result<Variable> {
let Some(declared) = return_type else {
return Ok(result);
};
match result {
Variable::Columns {
batch: columns,
} => {
let ctx = self.eval_ctx();
let coerced: Vec<(FieldRef, ArrayRef)> = user_columns(&columns)
.map(|(field, array)| {
let casted = cast_to_declared_return_type(
&ctx,
&ColumnView::try_from((array, field.as_ref()))?,
declared,
name.text(),
name,
)?;
Ok(factory::rename(casted, field.name()))
})
.collect::<Result<Vec<_>>>()?;
Ok(Variable::columns(batch(coerced)?))
}
other => Ok(other),
}
}
fn coerce_value(&self, value: Value, target: &ValueType, fragment: Fragment) -> Result<Value> {
match (value, target) {
(Value::Any(inner), ValueType::List(_)) => self.coerce_value(*inner, target, fragment),
(Value::List(items), ValueType::List(item)) => Ok(Value::List(
items.into_iter()
.map(|value| self.coerce_value(value, item, fragment.clone()))
.collect::<Result<Vec<_>>>()?,
)),
(value, target) => {
let mut data = ColumnBuilder::with_capacity(value.get_type(), 1);
data.push_value(value);
let data = data.finish("value");
let ctx = self.eval_ctx();
let cast = cast_column_data(
&ctx,
&ColumnView::try_from(&data)?,
target.clone(),
fragment,
)?;
Ok(ColumnView::try_from(&cast)?.get_value(0))
}
}
}
}
pub(crate) fn cast_to_declared_return_type(
ctx: impl Convert + Copy,
view: &ColumnView<'_>,
declared: &TypeConstraint,
name: &str,
fragment: &Fragment,
) -> Result<(FieldRef, ArrayRef)> {
let target = declared.get_type();
let casted = cast_column_data(ctx, view, target.clone(), fragment.clone()).map_err(|err| {
EngineError::ReturnTypeMismatch {
name: name.to_string(),
declared: target,
fragment: fragment.clone(),
cause: Box::new(err.diagnostic()),
}
})?;
Ok(factory::rename(casted, name))
}
pub(crate) fn cast_to_parameter_type(
ctx: impl Convert + Copy,
parameter: &FunctionParameter,
argument: (FieldRef, ArrayRef),
) -> Result<(FieldRef, ArrayRef)> {
let Some(declared) = ¶meter.type_constraint else {
return Ok(argument);
};
let fragment = Fragment::internal(strip_dollar_prefix(parameter.name.text()));
cast_column_data(ctx, &ColumnView::try_from(&argument)?, declared.get_type(), fragment)
}
pub(crate) fn declared_return_column(
ctx: impl Convert + Copy,
values: Vec<Value>,
declared: &TypeConstraint,
name: &str,
fragment: &Fragment,
) -> Result<(FieldRef, ArrayRef)> {
let results = results_column(&values, name)?;
match cast_to_declared_return_type(ctx, &ColumnView::try_from(&results)?, declared, name, fragment) {
Ok(casted) => {
let mut data = ColumnBuilder::with_capacity(declared.get_type(), values.len());
data.extend(&ColumnView::try_from(&casted)?)?;
Ok(data.finish(name))
}
Err(_) => declared_return_column_per_value(ctx, values, declared, name, fragment),
}
}
fn results_column(values: &[Value], name: &str) -> Result<(FieldRef, ArrayRef)> {
let wrapped = |value: &Value| {
matches!(value, Value::Any(_) | Value::List(_) | Value::Record(_) | Value::Tuple(_) | Value::Type(_))
};
let mut types = values.iter().filter(|value| !matches!(value, Value::None { .. })).map(Value::get_type);
let uniform = match types.next() {
Some(first) if !values.iter().any(wrapped) && types.all(|ty| ty == first) => Some(first),
_ => None,
};
let any = uniform.is_none();
let mut data = ColumnBuilder::with_capacity(uniform.unwrap_or(ValueType::Any), values.len());
for value in values {
match value {
Value::None {
..
} => data.push_none(),
value if any && !wrapped(value) => data.push_value(Value::Any(Box::new(value.clone()))),
value => data.push_value(value.clone()),
}
}
Ok(data.finish(name))
}
fn declared_return_column_per_value(
ctx: impl Convert + Copy,
values: Vec<Value>,
declared: &TypeConstraint,
name: &str,
fragment: &Fragment,
) -> Result<(FieldRef, ArrayRef)> {
let mut data = ColumnBuilder::with_capacity(declared.get_type(), values.len());
for value in values {
let single = factory::from_one(name, value);
let casted =
cast_to_declared_return_type(ctx, &ColumnView::try_from(&single)?, declared, name, fragment)?;
data.extend(&ColumnView::try_from(&casted)?)?;
}
Ok(data.finish(name))
}
pub(crate) fn untyped_return_column(
values: Vec<Value>,
name: &str,
fragment: &Fragment,
) -> Result<(FieldRef, ArrayRef)> {
let mut layout: Option<BranchLayout> = None;
for value in &values {
let fits_any = matches!(
value,
Value::None {
inner: ValueType::Any
}
);
let named_type = [(name, value.get_type(), fits_any)];
let Some(expected) = layout.as_mut() else {
layout = Some(BranchLayout::new(named_type));
continue;
};
expected.admit(named_type, fragment)?;
}
let column_type = values
.iter()
.find(|value| !matches!(value, Value::None { .. }))
.or(values.first())
.map(Value::get_type)
.unwrap_or(ValueType::Any);
let mut data = ColumnBuilder::with_capacity(column_type, values.len());
for value in values {
data.push_value(value);
}
Ok(data.finish(name))
}
pub(crate) fn check_arity(
parameters: &[FunctionParameter],
given: usize,
name: &str,
fragment: &Fragment,
) -> Result<()> {
if given == parameters.len() {
return Ok(());
}
Err(TypeError::Function {
kind: FunctionErrorKind::ArityMismatch {
expected: parameters.len(),
actual: given,
},
message: format!("function {} expects {} arguments, got {}", name, parameters.len(), given),
fragment: fragment.clone(),
}
.into())
}
fn unknown_procedure_error(func_name: &str, name: &Fragment) -> ReifyError {
TypeError::Procedure {
kind: ProcedureErrorKind::UndefinedProcedure {
name: func_name.to_string(),
},
message: format!("Unknown procedure: {}", func_name),
fragment: name.clone(),
}
.into()
}
fn assign_row_numbers_if_absent(columns: RecordBatch) -> Result<RecordBatch> {
if system_column(&columns, SystemColumn::RowNumbers).is_none() && columns.num_rows() > 0 {
let n = columns.num_rows() as u64;
let rns: ArrayRef = Arc::new(UInt64Array::from_iter_values(1..=n));
with_system_column(columns, SystemColumn::RowNumbers, rns)
} else {
Ok(columns)
}
}
pub(crate) fn collect_call_result(vm: &mut Vm, func_result: &mut Vec<Frame>) -> Variable {
match mem::replace(&mut vm.control_flow, ControlFlow::Normal) {
ControlFlow::Return(c) => match c {
Some(columns) => Variable::columns(columns),
None => Variable::scalar(Value::none()),
},
_ => {
if let Some(frame) = func_result.pop() {
if user_columns(&frame.batch).next().is_some() && frame.batch.num_rows() > 0 {
Variable::columns(frame.batch)
} else {
Variable::scalar(Value::none())
}
} else {
vm.stack.pop().ok().unwrap_or(Variable::scalar(Value::none()))
}
}
}
}
impl<'a> Vm<'a> {
pub(crate) fn exec_call(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
name: &Fragment,
arity: u8,
is_procedure_call: bool,
type_arguments: &[(usize, Fragment)],
) -> Result<()> {
let arity = arity as usize;
let func_name = name.text();
if !type_arguments.is_empty() && self.resolves_to_user_routine(services, tx, func_name)? {
self.rebind_type_arguments(arity, type_arguments)?;
}
if self.try_call_columnar(services, tx, func_name, arity, name)? {
return Ok(());
}
let args = self.pop_scalar_args(arity)?;
if let Some(callable) = self.symbols.resolve_callable(func_name) {
return self.call_callable(services, tx, &callable, args, name);
}
self.dispatch_resolved_procedure(services, tx, args, name, func_name, is_procedure_call)
}
fn resolves_to_user_routine(
&self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
func_name: &str,
) -> Result<bool> {
if self.symbols.resolve_callable(func_name).is_some() {
return Ok(true);
}
let mut tx_tmp = tx.reborrow();
Ok(services.catalog.find_procedure_by_qualified_name(&mut tx_tmp, func_name)?.is_some())
}
fn rebind_type_arguments(&mut self, arity: usize, type_arguments: &[(usize, Fragment)]) -> Result<()> {
let mut args = Vec::with_capacity(arity);
for _ in 0..arity {
args.push(self.stack.pop()?);
}
args.reverse();
for (index, arg) in args.into_iter().enumerate() {
match type_arguments.iter().find(|(position, _)| *position == index) {
Some((_, name)) => self.exec_load_var(name)?,
None => self.stack.push(arg),
}
}
Ok(())
}
#[inline]
fn try_call_columnar(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
func_name: &str,
arity: usize,
name: &Fragment,
) -> Result<bool> {
if self.batch_size == 1 {
return Ok(false);
}
if let Some(callable) = self.symbols.resolve_callable(func_name) {
self.call_callable_columnar(services, tx, &callable, arity, name)?;
return Ok(true);
}
if !self.resolves_to_plain_function(services, tx, func_name)? {
return Ok(false);
}
self.call_plain_function_columnar(services, tx, arity, name, func_name)?;
Ok(true)
}
fn resolves_to_plain_function(
&self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
func_name: &str,
) -> Result<bool> {
if services.routines.get_function(func_name).is_none() {
return Ok(false);
}
if services.routines.get_procedure(func_name).is_some()
|| services.routines.get_generator_function(func_name).is_some()
{
return Ok(false);
}
let mut tx_tmp = tx.reborrow();
Ok(services.catalog.find_procedure_by_qualified_name(&mut tx_tmp, func_name)?.is_none())
}
fn call_plain_function_columnar(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
arity: usize,
name: &Fragment,
func_name: &str,
) -> Result<()> {
let function = services.routines.get_function(func_name).ok_or_else(|| {
ReifyError::from(EngineError::UnknownCallable {
name: func_name.to_string(),
fragment: name.clone(),
})
})?;
function.arity().check(name, arity)?;
let arg_columns = self.pop_args_as_columns(arity, name)?;
let identity = tx.identity();
let mut fn_ctx = RoutineFunctionContext {
fragment: name.clone(),
identity,
row_count: arg_columns.first().map_or(0, |c| c.1.len()),
runtime_context: &services.runtime_context,
};
let result_column =
function.call(&mut fn_ctx, &arg_columns).map_err(|e| e.with_context(name.clone(), false))?;
self.stack.push(Variable::columns(batch(vec![result_column])?));
Ok(())
}
#[inline]
fn pop_scalar_args(&mut self, arity: usize) -> Result<Vec<Value>> {
let mut args = Vec::with_capacity(arity);
for _ in 0..arity {
args.push(self.pop_value()?);
}
args.reverse();
reifydb_assertions! {
let popped = args.len();
assert!(
popped == arity,
"argument pop produced {popped} values for an arity-{arity} call; downstream \
callees positionally zip params against args, so a miscount silently mis-binds \
or drops arguments instead of failing the call"
);
}
Ok(args)
}
#[inline]
fn dispatch_resolved_procedure(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
is_procedure_call: bool,
) -> Result<()> {
let proc_def = {
let mut tx_tmp = tx.reborrow();
services.catalog.find_procedure_by_qualified_name(&mut tx_tmp, func_name)?
};
match proc_def {
Some(ResolvedProcedure::Local(proc_def)) => {
let ctx = CallContext {
services,
tx,
};
self.call_local_procedure(ctx, &proc_def, args, name, func_name)
}
Some(ResolvedProcedure::Test(proc_def)) => {
let ctx = CallContext {
services,
tx,
};
self.call_test_procedure(ctx, &proc_def, args, name, func_name)
}
#[cfg(not(reifydb_single_threaded))]
Some(ResolvedProcedure::Remote {
address,
token,
}) => self.call_remote_procedure(services, args, name, func_name, &address, token.as_deref()),
#[cfg(reifydb_single_threaded)]
Some(ResolvedProcedure::Remote {
..
}) => Err(TypeError::Procedure {
kind: ProcedureErrorKind::UndefinedProcedure {
name: func_name.to_string(),
},
message: format!("Unknown procedure: {}", func_name),
fragment: name.clone(),
}
.into()),
None => {
let ctx = CallContext {
services,
tx,
};
self.call_builtin_or_error(ctx, args, name, func_name, is_procedure_call)
}
}
}
fn call_callable_columnar(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
callable: &Callable,
arity: usize,
name: &Fragment,
) -> Result<()> {
let arg_columns = self.pop_args_as_columns(arity, name)?;
check_arity(&callable.parameters, arity, name.text(), name)?;
if is_vectorizable(&callable.body, &services.routines, &self.symbols) {
let row_count = arg_columns.first().map(|c| c.1.len()).unwrap_or(self.batch_size);
self.run_function_body_batch(
services,
tx,
&callable.body,
&callable.parameters,
arg_columns,
row_count,
&callable.captured,
callable.return_type.as_ref(),
name,
)
} else {
self.run_function_body_per_row(
services,
tx,
&callable.body,
&callable.parameters,
arg_columns,
&callable.captured,
name,
callable.return_type.as_ref(),
)
}
}
fn cast_arguments(
&self,
parameters: &[FunctionParameter],
arguments: Vec<(FieldRef, ArrayRef)>,
) -> Result<Vec<(FieldRef, ArrayRef)>> {
let ctx = self.eval_ctx();
parameters
.iter()
.zip(arguments)
.map(|(parameter, argument)| cast_to_parameter_type(&ctx, parameter, argument))
.collect()
}
fn pop_args_as_columns(&mut self, arity: usize, name: &Fragment) -> Result<Vec<(FieldRef, ArrayRef)>> {
let mut arg_columns = Vec::with_capacity(arity);
for _ in 0..arity {
arg_columns.push(self.pop_as_column()?);
}
arg_columns.reverse();
if let Some(first) = arg_columns.first() {
let row_count = first.1.len();
if let Some(other) = arg_columns.iter().find(|column| column.1.len() != row_count) {
return Err(length_mismatch(row_count, other.1.len(), name).into());
}
}
Ok(arg_columns)
}
#[allow(clippy::too_many_arguments)]
fn run_function_body_batch(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
body: &[Instruction],
parameters: &[FunctionParameter],
arg_columns: Vec<(FieldRef, ArrayRef)>,
_row_count: usize,
captured: &HashMap<String, Variable>,
return_type: Option<&TypeConstraint>,
name: &Fragment,
) -> Result<()> {
let arguments = self.cast_arguments(parameters, arg_columns)?;
let saved_ip = self.ip;
let saved_pending_return = self.pending_return.take();
let saved_returned_mask = self.returned_mask.take();
let saved_udf_call = mem::replace(
&mut self.udf_call,
UdfCall {
fragment: name.clone(),
return_type: return_type.cloned(),
},
);
self.symbols.enter_scope();
for (cap_name, cap_var) in captured {
self.symbols.set(cap_name.clone(), cap_var.clone(), true)?;
}
for (param, arg_col) in parameters.iter().zip(arguments) {
let param_name = strip_dollar_prefix(param.name.text()).to_string();
let col_var = Variable::columns(batch(vec![arg_col])?);
self.symbols.set(param_name, col_var, true)?;
}
self.ip = 0;
let mut func_result = Vec::new();
self.run_isolated_body(services, tx, body, &mut func_result)?;
let stack_value = collect_call_result(self, &mut func_result);
self.ip = saved_ip;
self.pending_return = saved_pending_return;
self.returned_mask = saved_returned_mask;
self.udf_call = saved_udf_call;
let _ = self.symbols.exit_scope();
let coerced = self.coerce_return_value(stack_value, return_type, name)?;
self.stack.push(coerced);
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn run_function_body_per_row(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
body: &[Instruction],
parameters: &[FunctionParameter],
arg_columns: Vec<(FieldRef, ArrayRef)>,
captured: &HashMap<String, Variable>,
name: &Fragment,
return_type: Option<&TypeConstraint>,
) -> Result<()> {
let arg_columns = self.cast_arguments(parameters, arg_columns)?;
let arg_views = arg_columns.iter().map(ColumnView::try_from).collect::<Result<Vec<_>>>()?;
let row_count = arg_columns.first().map(|c| c.1.len()).unwrap_or(0);
let mut results: Vec<Value> = Vec::with_capacity(row_count);
let mut func_symbols = self.symbols.clone();
for row_idx in 0..row_count {
func_symbols.enter_scope();
for (cap_name, cap_var) in captured {
func_symbols.set(cap_name.clone(), cap_var.clone(), true)?;
}
for (param, arg_view) in parameters.iter().zip(arg_views.iter()) {
let param_name = strip_dollar_prefix(param.name.text()).to_string();
let value = arg_view.get_value(row_idx);
func_symbols.set(
param_name.clone(),
Variable::scalar_named(¶m_name, value),
true,
)?;
}
let mut vm = Vm::from_services(func_symbols, services, &EMPTY_PARAMS, tx.identity());
let mut func_result: Vec<Frame> = Vec::new();
vm.run(services, tx, body, &mut func_result)?;
let result_var = collect_call_result(&mut vm, &mut func_result);
let value = match result_var {
Variable::Columns {
batch: c,
} if is_scalar(&c) => scalar_value(&c)?,
_ => Value::none(),
};
func_symbols = vm.symbols;
let _ = func_symbols.exit_scope();
results.push(value);
}
let result_col = match return_type {
Some(declared) => {
declared_return_column(&self.eval_ctx(), results, declared, name.text(), name)?
}
None => untyped_return_column(results, name.text(), name)?,
};
self.stack.push(Variable::columns(batch(vec![result_col])?));
Ok(())
}
fn call_callable(
&mut self,
services: &Arc<Services>,
tx: &mut Transaction<'_>,
callable: &Callable,
args: Vec<Value>,
name: &Fragment,
) -> Result<()> {
check_arity(&callable.parameters, args.len(), name.text(), name)?;
let ctx = self.eval_ctx();
let mut arguments = Vec::with_capacity(args.len());
for (param, arg) in callable.parameters.iter().zip(args) {
let argument = factory::from_one(strip_dollar_prefix(param.name.text()), arg);
let cast = cast_to_parameter_type(&ctx, param, argument)?;
arguments.push(ColumnView::try_from(&cast)?.get_value(0));
}
let saved_ip = self.ip;
self.symbols.enter_scope();
for (name, var) in &callable.captured {
self.symbols.set(name.clone(), var.clone(), true)?;
}
for (param, arg) in callable.parameters.iter().zip(arguments) {
let param_name = strip_dollar_prefix(param.name.text()).to_string();
self.symbols.set(param_name.clone(), Variable::scalar_named(¶m_name, arg), true)?;
}
self.ip = 0;
let mut result = Vec::new();
self.run_isolated_body(services, tx, &callable.body, &mut result)?;
let stack_value = collect_call_result(self, &mut result);
self.ip = saved_ip;
let _ = self.symbols.exit_scope();
let coerced = self.coerce_return_value(stack_value, callable.return_type.as_ref(), name)?;
self.stack.push(coerced);
Ok(())
}
fn call_local_procedure(
&mut self,
ctx: CallContext<'_, '_>,
proc_def: &Procedure,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
) -> Result<()> {
let (pol_ns, pol_name) = if let Some((ns, n)) = Catalog::split_qualified_name(func_name) {
(ns, n.to_string())
} else {
("default".to_string(), func_name.to_string())
};
PolicyEvaluator::new(ctx.services, &self.symbols).enforce_identity_policy(
ctx.tx,
&pol_ns,
&pol_name,
CallableOp::Call,
PolicyTargetType::Procedure,
)?;
match proc_def {
Procedure::InProcess {
handler_name,
..
}
| Procedure::ExternC {
handler_name,
..
}
| Procedure::ExternWasm {
handler_name,
..
} => {
let handler_name = handler_name.clone();
if let Some(routine) = ctx.services.routines.get_procedure(&handler_name) {
let call_params = Params::Positional(Arc::new(args));
let identity = ctx.tx.identity();
let mut proc_ctx = RoutineProcedureContext {
fragment: name.clone(),
identity,
row_count: 1,
runtime_context: &ctx.services.runtime_context,
tx: ctx.tx,
params: &call_params,
catalog: &ctx.services.catalog,
routines: &ctx.services.routines,
ioc: &ctx.services.ioc,
};
let attach_metadata = routine.attaches_row_metadata();
let columns = routine
.call(&mut proc_ctx, &[])
.map_err(|e| e.with_context(name.clone(), true))?;
let columns = if attach_metadata {
assign_row_numbers_if_absent(columns)?
} else {
columns
};
self.stack.push(Variable::columns(columns));
Ok(())
} else {
Err(TypeError::Procedure {
kind: ProcedureErrorKind::NoRegisteredImplementation {
name: handler_name.clone(),
},
message: format!(
"procedure handler '{}' has no registered implementation",
handler_name
),
fragment: name.clone(),
}
.into())
}
}
Procedure::Rql {
body,
params,
..
}
| Procedure::Test {
body,
params,
..
} => {
let source = body.clone();
let params = params.clone();
let compiled = ctx.services.compiler.compile_with_policy(
ctx.tx,
&source,
inject_from_policies,
)?;
match compiled {
CompilationResult::Ready(compiled_list) => {
self.execute_procedure_body(ctx, &compiled_list, ¶ms, args, name)?;
Ok(())
}
CompilationResult::Incremental(_) => Err(internal_error!(
"Procedure body should not require incremental compilation"
)),
}
}
}
}
fn call_test_procedure(
&mut self,
ctx: CallContext<'_, '_>,
proc_def: &Procedure,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
) -> Result<()> {
if !matches!(ctx.tx, Transaction::Test(..)) {
return Err(TypeError::Procedure {
kind: ProcedureErrorKind::UndefinedProcedure {
name: func_name.to_string(),
},
message: format!("test procedure {} can only be called from test context", func_name),
fragment: name.clone(),
}
.into());
}
let source = proc_def.body().unwrap_or_default().to_string();
let params = proc_def.params().to_vec();
let compiled = ctx.services.compiler.compile_with_policy(ctx.tx, &source, inject_from_policies)?;
match compiled {
CompilationResult::Ready(compiled_list) => {
self.execute_procedure_body(ctx, &compiled_list, ¶ms, args, name)?;
Ok(())
}
CompilationResult::Incremental(_) => {
Err(internal_error!("Procedure body should not require incremental compilation"))
}
}
}
fn execute_procedure_body(
&mut self,
ctx: CallContext<'_, '_>,
compiled_list: &[Compiled],
proc_params: &[ProcedureParam],
args: Vec<Value>,
name: &Fragment,
) -> Result<()> {
let args = if args.is_empty() && !proc_params.is_empty() && matches!(self.params, Params::Named(_)) {
let named: Vec<Value> = proc_params
.iter()
.filter_map(|p| self.params.get_named(strip_dollar_prefix(&p.name)).cloned())
.collect();
if named.len() != proc_params.len() {
return Err(TypeError::Function {
kind: FunctionErrorKind::ArityMismatch {
expected: proc_params.len(),
actual: named.len(),
},
message: format!(
"procedure {} expects {} arguments, got {} named",
name.text(),
proc_params.len(),
named.len()
),
fragment: name.clone(),
}
.into());
}
named
} else {
args
};
if args.len() != proc_params.len() {
return Err(TypeError::Function {
kind: FunctionErrorKind::ArityMismatch {
expected: proc_params.len(),
actual: args.len(),
},
message: format!(
"procedure {} expects {} arguments, got {}",
name.text(),
proc_params.len(),
args.len()
),
fragment: name.clone(),
}
.into());
}
let saved_ip = self.ip;
self.symbols.enter_scope();
for (param_def, arg) in proc_params.iter().zip(args) {
let bare_name = strip_dollar_prefix(¶m_def.name);
let arg = self.coerce_value(
arg,
¶m_def.param_type.get_type(),
Fragment::internal(bare_name),
)?;
self.symbols.set(param_def.name.clone(), Variable::scalar_named(bare_name, arg), true)?;
}
let mut proc_result = Vec::new();
for compiled in compiled_list.iter() {
self.ip = 0;
self.run(ctx.services, ctx.tx, &compiled.instructions, &mut proc_result)?;
if !self.control_flow.is_normal() {
break;
}
}
let stack_value = collect_call_result(self, &mut proc_result);
self.ip = saved_ip;
let _ = self.symbols.exit_scope();
self.stack.push(stack_value);
Ok(())
}
#[cfg(not(reifydb_single_threaded))]
fn call_remote_procedure(
&mut self,
services: &Arc<Services>,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
address: &str,
token: Option<&str>,
) -> Result<()> {
if let Some(ref registry) = services.remote_registry {
let param_refs: Vec<String> = (1..=args.len()).map(|i| format!("${}", i)).collect();
let remote_rql = format!("CALL {}({})", func_name, param_refs.join(", "));
let frames = registry.forward_query(
address,
&remote_rql,
Params::Positional(Arc::new(args)),
token,
)?;
if let Some(frame) = frames.into_iter().next() {
self.stack.push(Variable::columns(frame.batch));
} else {
self.stack.push(Variable::scalar(Value::none()));
}
Ok(())
} else {
Err(TypeError::Procedure {
kind: ProcedureErrorKind::UndefinedProcedure {
name: func_name.to_string(),
},
message: format!("Unknown procedure: {}", func_name),
fragment: name.clone(),
}
.into())
}
}
fn call_builtin_or_error(
&mut self,
ctx: CallContext<'_, '_>,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
is_procedure_call: bool,
) -> Result<()> {
if let Some(routine) = ctx.services.routines.get_procedure(func_name) {
return self.call_procedure_routine(ctx, routine, args, name, func_name);
}
if let Some(generator) = ctx.services.routines.get_generator_function(func_name) {
return self.call_generator_function(ctx, generator, args, name);
}
if is_procedure_call {
return Err(unknown_procedure_error(func_name, name));
}
self.call_plain_function(ctx, args, name, func_name)
}
#[inline]
fn call_procedure_routine(
&mut self,
ctx: CallContext<'_, '_>,
routine: Arc<dyn RoutineProcedure>,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
) -> Result<()> {
let call_params = Params::Positional(Arc::new(args));
let attach_metadata = routine.attaches_row_metadata();
let columns = invoke_procedure_routine(
ctx.services,
&self.symbols,
ctx.tx,
ProcedureCall {
routine: &routine,
fragment: name,
target: func_name,
params: &call_params,
},
CallSite::Named,
)?;
let columns = if attach_metadata {
assign_row_numbers_if_absent(columns)?
} else {
columns
};
if func_name == "identity::inject"
&& let Some((field, array)) = user_columns(&columns).next()
&& let Value::IdentityId(id) = ColumnView::try_from((array, field.as_ref()))?.get_value(0)
{
ctx.tx.set_identity(id);
}
self.stack.push(Variable::columns(columns));
Ok(())
}
#[inline]
fn call_generator_function(
&mut self,
ctx: CallContext<'_, '_>,
generator: Arc<dyn RoutineFunction>,
args: Vec<Value>,
name: &Fragment,
) -> Result<()> {
generator.arity().check(name, args.len())?;
let arg_columns: Vec<(FieldRef, ArrayRef)> = args
.into_iter()
.enumerate()
.map(|(i, v)| {
let mut data = ColumnBuilder::with_capacity(v.get_type(), 1);
data.push_value(v);
data.finish(&format!("arg{}", i))
})
.collect();
let identity = ctx.tx.identity();
let mut fn_ctx = RoutineFunctionContext {
fragment: name.clone(),
identity,
row_count: arg_columns.first().map_or(0, |c| c.1.len()),
runtime_context: &ctx.services.runtime_context,
};
let column =
generator.call(&mut fn_ctx, &arg_columns).map_err(|e| e.with_context(name.clone(), false))?;
self.stack.push(Variable::columns(batch(vec![column])?));
Ok(())
}
#[inline]
fn call_plain_function(
&mut self,
ctx: CallContext<'_, '_>,
args: Vec<Value>,
name: &Fragment,
func_name: &str,
) -> Result<()> {
let function = ctx.services.routines.get_function(func_name).ok_or_else(|| {
ReifyError::from(EngineError::UnknownCallable {
name: func_name.to_string(),
fragment: name.clone(),
})
})?;
function.arity().check(name, args.len())?;
let arg_columns: Vec<(FieldRef, ArrayRef)> = args
.into_iter()
.enumerate()
.map(|(i, v)| {
let mut data = ColumnBuilder::with_capacity(v.get_type(), 1);
data.push_value(v);
data.finish(&format!("arg{}", i))
})
.collect();
let identity = ctx.tx.identity();
let mut fn_ctx = RoutineFunctionContext {
fragment: name.clone(),
identity,
row_count: 1,
runtime_context: &ctx.services.runtime_context,
};
let result_column =
function.call(&mut fn_ctx, &arg_columns).map_err(|e| e.with_context(name.clone(), false))?;
let value = if result_column.1.is_empty() {
Value::none()
} else {
ColumnView::try_from(&result_column)?.get_value(0)
};
self.stack.push(Variable::scalar(value));
Ok(())
}
pub(crate) fn exec_define_function(&mut self, node: &CompiledFunction) {
let func_name = node.name.text().to_string();
self.symbols.define_function(func_name, node.clone());
}
pub(crate) fn exec_return_value(&mut self) -> Result<()> {
let cols = self.pop_as_columns()?;
self.control_flow = ControlFlow::Return(Some(cols));
Ok(())
}
pub(crate) fn exec_return_void(&mut self) {
self.control_flow = ControlFlow::Return(None);
}
pub(crate) fn exec_define_closure(&mut self, closure_def: &CompiledClosure) {
let mut captured = HashMap::new();
for cap_name in &closure_def.captures {
let stripped = strip_dollar_prefix(cap_name.text()).to_string();
if let Some(var) = self.symbols.get(&stripped) {
captured.insert(stripped, var.clone());
}
}
self.stack.push(Variable::Closure(ClosureValue {
def: closure_def.clone(),
captured,
}));
}
}