surrealdb-core 3.3.1

//! Shared helpers for the plan-or-compute fallback pattern.
//!
//! Operators that use deferred planning (foreach, ifelse, sequence, block) share
//! a common pattern: try to plan an expression with the streaming engine, and if
//! the planner returns `PlannerUnsupported` or `PlannerUnimplemented`, fall back
//! to the legacy `Expr::compute()` path.
//!
//! This module centralises that logic so each operator does not need its own copy.

use std::sync::Arc;

use futures::StreamExt;
use reblessive::tree::TreeStack;

use crate::ctx::FrozenContext;
use crate::err::Error;
use crate::exec::context::{ContextLevel, ExecutionContext};
use crate::exec::planner::try_plan_expr;
use crate::exec::{Error as ExecError, FlowResult, ValueBatchStream};
use crate::expr::part::{Part, RecurseInstruction};
use crate::expr::statements::InfoStatement;
use crate::expr::{Base, Block, ControlFlow, ControlFlowExt, Expr, Literal};
use crate::val::Value;

// ============================================================================
// Legacy Context Helpers
// ============================================================================

/// Extract the `Options` and `FrozenContext` needed for legacy `Expr::compute()`.
///
/// The `ExecutionContext`'s `FrozenContext` is the single source of truth for
/// parameters, transactions, capabilities, and all legacy context fields.
pub(crate) fn get_legacy_context(
	exec_ctx: &ExecutionContext,
) -> Result<(crate::dbs::Options, FrozenContext), Error> {
	let options = exec_ctx.options().ok_or_else(|| {
		ExecError::Thrown("Options not available for legacy compute fallback".into())
	})?;
	let options = legacy_fallback_options(
		options,
		exec_ctx.root().skip_fetch_perms,
		exec_ctx.root().computing_field,
	);
	Ok((options, Arc::clone(exec_ctx.ctx())))
}

/// Derive the `Options` to use for a legacy compute fallback.
///
/// The streaming engine plans no writes of its own, so every write in an
/// expression reaches the legacy path through here. Two frames must not be
/// allowed to take it:
///
/// - a `PERMISSIONS` predicate (signalled by `skip_fetch_perms`), which is evaluated with
///   permission enforcement disabled (GHSA-66r2-5gwj-gxm2);
/// - a `COMPUTED` field body (signalled by `computing_field`), which is evaluated on every read of
///   the field and under the definer's auth.
///
/// The permission frame also disables `perms`; the computed frame does not,
/// because a computed body's own reads stay permission-checked.
/// The two signals are taken as arguments rather than read from `exec_ctx`
/// because the permission frame is tracked per `EvalContext` at some call sites
/// and per root at others. Naming them at each site keeps this the only place
/// that decides what they mean.
pub(crate) fn legacy_fallback_options(
	options: &crate::dbs::Options,
	permission_predicate: bool,
	computing_field: bool,
) -> crate::dbs::Options {
	if permission_predicate {
		options.new_for_permission_predicate()
	} else if computing_field {
		options.new_for_computed_field()
	} else {
		options.clone()
	}
}

/// Extract the `Options` and `FrozenContext` for legacy fallback, adding a loop
/// variable to the context.
///
/// Used by the `ForeachPlan` operator to inject the current iteration value.
pub(crate) fn get_legacy_context_with_param(
	exec_ctx: &ExecutionContext,
	param_name: &str,
	param_value: &Value,
) -> Result<(crate::dbs::Options, FrozenContext), Error> {
	let options = exec_ctx.options().ok_or_else(|| {
		ExecError::Thrown("Options not available for legacy compute fallback".into())
	})?;
	let options = legacy_fallback_options(
		options,
		exec_ctx.root().skip_fetch_perms,
		exec_ctx.root().computing_field,
	);

	let mut ctx = crate::ctx::Context::new_child(exec_ctx.ctx());
	ctx.add_value(param_name.to_string(), std::sync::Arc::new(param_value.clone()));

	Ok((options, ctx.freeze()))
}

// ============================================================================
// Legacy Compute
// ============================================================================

/// Execute an expression using the legacy `Expr::compute()` path.
///
/// This is the **single call site** for `Expr::compute()` in the streaming
/// executor.  All operators that need legacy compute fallback should call this
/// function rather than invoking `.compute()` directly.
pub(crate) async fn legacy_compute(
	expr: &Expr,
	frozen: &FrozenContext,
	opt: &crate::dbs::Options,
	doc: Option<&crate::doc::CursorDoc>,
) -> crate::expr::FlowResult<Value> {
	let mut stack = TreeStack::new();
	stack.enter(|stk| crate::legacy::expr_compute(expr, stk, frozen, opt, doc)).finish().await
}

// ============================================================================
// Plan-or-Compute Evaluation
// ============================================================================

/// The transaction the deferred-planning operators need to re-plan an
/// expression at execute time.
///
/// Planning resolves table and index definitions through the transaction, so a
/// context without one cannot plan at all. Reported as an error rather than
/// read through [`ExecutionContext::txn`], whose missing-transaction case
/// panics.
pub(crate) fn planning_txn(ctx: &ExecutionContext) -> Result<Arc<crate::kvs::Transaction>, Error> {
	ctx.try_txn().ok_or_else(|| {
		Error::unreachable("Deferred planning requires a transaction, but the context has none")
	})
}

/// Plan and evaluate an expression, falling back to legacy compute if the
/// planner returns `PlannerUnsupported` or `PlannerUnimplemented`.
///
/// This is the simple variant used when no context mutation is needed
/// (e.g. evaluating a FOR range, an IF condition, or an IF/ELSE branch body).
/// Plan and evaluate an expression, falling back to legacy compute if the
/// planner returns `PlannerUnsupported` or `PlannerUnimplemented`, seeding the
/// planner's expression-nesting depth (see
/// [`crate::exec::planner::Planner::with_depth`]).
///
/// `depth` is 0 for an ordinary top-level evaluation. It is non-zero when
/// re-planning a nested query at runtime — chiefly `eval::*` re-planning its
/// query string, and the deferred control-flow operators (IfElse/Foreach/
/// Sequence) re-planning their bodies — where passing the depth recorded on the
/// triggering node continues the count toward `max_computation_depth` instead of
/// resetting it. That is what bounds nested-`eval`/UDF recursion the same way the
/// legacy `compute` path bounds it via `Options::dive`; the legacy fallback below
/// mirrors it by handing the remaining budget (`max - depth`) to `compute`.
pub(crate) async fn evaluate_expr_at_depth(
	expr: &Expr,
	ctx: &ExecutionContext,
	depth: u32,
) -> crate::expr::FlowResult<Value> {
	let auth = ctx.options().map(|o| Arc::clone(&o.auth));
	// The transaction lookup is fallible and, per `try_plan_expr!`, runs only if
	// the macro takes its planning branch.
	match try_plan_expr!(
		expr,
		ctx.ctx(),
		ctx.function_registry(),
		planning_txn(ctx).map_err(|e| ControlFlow::Err(e.into()))?,
		auth,
		depth
	) {
		Ok(plan) => {
			let shape = plan.output_shape();
			let stream = plan.execute(ctx)?;
			collect_shaped_value(stream, shape).await
		}
		Err(Error::Exec(
			e @ (ExecError::PlannerUnsupported(_) | ExecError::PlannerUnimplemented(_)),
		)) => {
			match &e {
				ExecError::PlannerUnimplemented(msg) => {
					tracing::warn!("PlannerUnimplemented fallback in evaluate_expr: {msg}");
				}
				ExecError::PlannerUnsupported(msg) => {
					tracing::debug!("PlannerUnsupported fallback in evaluate_expr: {msg}",);
				}
				_ => {}
			}
			let (opt, frozen) =
				get_legacy_context(ctx).context("Legacy compute fallback context unavailable")?;
			// Continue the depth count into the legacy path: a DML/DDL body reached
			// via `eval` keeps counting from `depth` rather than a fresh budget.
			let opt = opt.with_dive_consumed(depth);
			legacy_compute(expr, &frozen, &opt, None).await
		}
		Err(e) => Err(ControlFlow::Err(e.into())),
	}
}

/// Plan and evaluate a body expression that may mutate the execution context
/// (e.g. a LET statement inside a FOR loop body).
///
/// When the planned operator has `mutates_context() == true`, the context is
/// updated via `output_context()`. The legacy fallback for loop bodies injects
/// the loop variable into the context before calling `Expr::compute()`.
pub(crate) async fn evaluate_body_expr(
	expr: &Expr,
	ctx: &mut ExecutionContext,
	param_name: &str,
	param_value: &Value,
	depth: u32,
) -> crate::expr::FlowResult<Value> {
	let frozen_ctx = Arc::clone(ctx.ctx());
	let auth = ctx.options().map(|o| Arc::clone(&o.auth));

	// Fallible transaction lookup, evaluated only on the planning branch.
	match try_plan_expr!(
		expr,
		&frozen_ctx,
		ctx.function_registry(),
		planning_txn(ctx).map_err(|e| ControlFlow::Err(e.into()))?,
		auth,
		depth
	) {
		Ok(plan) => {
			if plan.mutates_context() {
				// A loop signal from the bound expression is the caller's to act
				// on — `ForeachPlan` breaks or continues its loop on it.
				*ctx = plan.output_context(ctx).await?;
				Ok(Value::None)
			} else {
				let shape = plan.output_shape();
				let stream = plan.execute(ctx)?;
				collect_shaped_value(stream, shape).await
			}
		}
		Err(Error::Exec(
			e @ (ExecError::PlannerUnsupported(_) | ExecError::PlannerUnimplemented(_)),
		)) => {
			match &e {
				ExecError::PlannerUnimplemented(msg) => {
					tracing::warn!("PlannerUnimplemented fallback in evaluate_body_expr: {msg}");
				}
				ExecError::PlannerUnsupported(msg) => {
					tracing::debug!("PlannerUnsupported fallback in evaluate_body_expr: {msg}",);
				}
				_ => {}
			}
			let (opt, frozen) = get_legacy_context_with_param(ctx, param_name, param_value)
				.context("Legacy compute fallback context unavailable")?;
			// Continue the depth count into the legacy path from `depth`.
			let opt = opt.with_dive_consumed(depth);

			if let Expr::Let(set_stmt) = expr {
				if set_stmt.is_protected_set() {
					return Err(ExecError::InvalidParam {
						name: set_stmt.name.to_string(),
					}
					.into());
				}

				let value = legacy_compute(&set_stmt.what, &frozen, &opt, None).await?;

				let value = if let Some(kind) = &set_stmt.kind {
					value.coerce_to_kind(kind).map_err(|e| ExecError::SetCoerce {
						name: set_stmt.name.to_string(),
						error: Box::new(e),
					})?
				} else {
					value
				};

				*ctx = ctx.with_param(set_stmt.name.clone(), value);
				Ok(Value::None)
			} else {
				legacy_compute(expr, &frozen, &opt, None).await
			}
		}
		Err(e) => Err(ControlFlow::Err(e.into())),
	}
}

// ============================================================================
// Stream Collection Helpers
// ============================================================================

/// Collect values from a stream into the single value an expression evaluates
/// to, shaped by the plan's declared [`OutputShape`] rather than how many rows
/// it happened to produce.
///
/// - `OutputShape::Scalar` (e.g. `SELECT ONLY`, or any non-query expression) → the one emitted
///   value, or `Value::None` if the plan emitted none.
/// - `OutputShape::Rows` (e.g. an ordinary `SELECT`) → always `Value::Array`, regardless of row
///   count. A caller using the result as a `FOR` range, or any other array-shaped context, must see
///   `[]` for zero rows and `[x]` for one — never `NONE` or the bare element — the same contract
///   [`ScalarSubquery`](crate::exec::physical_expr::subquery::ScalarSubquery) already applies to a
///   subquery expression.
///
/// Propagates control flow signals (BREAK, CONTINUE, RETURN, errors).
pub(crate) async fn collect_shaped_value(
	mut stream: ValueBatchStream,
	shape: crate::exec::OutputShape,
) -> crate::expr::FlowResult<Value> {
	let mut values = Vec::new();

	while let Some(batch_result) = stream.next().await {
		match batch_result {
			Ok(batch) => values.extend(batch.into_values()),
			Err(ctrl) => return Err(ctrl),
		}
	}

	if shape.is_scalar() {
		Ok(values.pop().unwrap_or(Value::None))
	} else {
		Ok(Value::Array(crate::val::Array(values)))
	}
}

/// Collect all values from a stream into a `Vec`.
///
/// Propagates control flow signals directly.
pub(crate) async fn collect_stream(mut stream: ValueBatchStream) -> FlowResult<Vec<Value>> {
	let mut results = Vec::new();

	while let Some(batch_result) = stream.next().await {
		match batch_result {
			Ok(batch) => results.extend(batch.into_values()),
			Err(ctrl) => return Err(ctrl),
		}
	}

	Ok(results)
}

// ============================================================================
// Expression Context Level Analysis
// ============================================================================

/// Determine the minimum [`ContextLevel`] required to evaluate an [`Expr`].
///
/// This is used by deferred-planning operators ([`IfElsePlan`], [`SequencePlan`],
/// [`ForeachPlan`]) that store raw `Expr` values and plan them at runtime.
/// The function recursively inspects the expression tree to determine the
/// minimum context level needed, avoiding overly conservative hardcoded
/// `ContextLevel::Database` values.
pub(crate) fn expr_required_context(expr: &Expr) -> ContextLevel {
	match expr {
		// Pure values — no context needed
		Expr::Param(_) | Expr::Constant(_) | Expr::Mock(_) | Expr::Break | Expr::Continue => {
			ContextLevel::Root
		}

		// Literals may contain inner expressions (arrays, objects, record IDs)
		Expr::Literal(lit) => literal_required_context(lit),

		// Closure body may reference database-level constructs
		Expr::Closure(closure) => expr_required_context(&closure.body),

		// Sleep just does a time delay (duration is a concrete value, not an Expr)
		Expr::Sleep(_) => ContextLevel::Root,

		// Idiom parts may contain inner expressions (WHERE filters, method args,
		// graph lookups, etc.) that need their own context.
		Expr::Idiom(idiom) => idiom_required_context(idiom),

		// Table references are used in FROM clauses which need a database
		Expr::Table(_) => ContextLevel::Database,

		// Block: max of all contained expressions
		Expr::Block(block) => block_required_context(block),

		// Unary operators: delegate to inner expression
		Expr::Prefix {
			expr,
			..
		}
		| Expr::Postfix {
			expr,
			..
		}
		| Expr::Throw(expr) => expr_required_context(expr),

		// Binary: max of left and right
		Expr::Binary {
			left,
			right,
			..
		} => expr_required_context(left).max(expr_required_context(right)),

		// Function calls may be user-defined (stored in the database), so
		// conservatively require database context.
		Expr::FunctionCall(_) => ContextLevel::Database,

		// Return: delegate to inner expression
		Expr::Return(stmt) => expr_required_context(&stmt.what),

		// IfElse: max of all conditions and branches
		Expr::IfElse(stmt) => {
			let branches_ctx = stmt
				.exprs
				.iter()
				.flat_map(|(cond, body)| [expr_required_context(cond), expr_required_context(body)])
				.max()
				.unwrap_or(ContextLevel::Root);
			let else_ctx =
				stmt.close.as_ref().map(expr_required_context).unwrap_or(ContextLevel::Root);
			branches_ctx.max(else_ctx)
		}

		// DML statements need a database
		Expr::Select(_)
		| Expr::Create(_)
		| Expr::Update(_)
		| Expr::Upsert(_)
		| Expr::Delete(_)
		| Expr::Relate(_)
		| Expr::Insert(_) => ContextLevel::Database,

		// DDL statements need a database
		Expr::Define(_) | Expr::Remove(_) | Expr::Alter(_) | Expr::Rebuild(_) => {
			ContextLevel::Database
		}

		// Info: depends on the level
		Expr::Info(info) => info_stmt_required_context(info),

		// Foreach: max of range expression and body block
		Expr::Foreach(stmt) => {
			expr_required_context(&stmt.range).max(block_required_context(&stmt.block))
		}

		// Let: delegate to value expression
		Expr::Let(stmt) => expr_required_context(&stmt.what),

		// Explain: delegate to inner statement
		Expr::Explain {
			statement,
			..
		} => expr_required_context(statement),

		// GQL MATCH queries the datastore and needs a database.
		Expr::Match(_) => ContextLevel::Database,
	}
}

/// Determine the minimum [`ContextLevel`] required for a [`Block`].
pub(crate) fn block_required_context(block: &Block) -> ContextLevel {
	block.0.iter().map(expr_required_context).max().unwrap_or(ContextLevel::Root)
}

/// Determine the minimum [`ContextLevel`] required by an [`InfoStatement`].
fn info_stmt_required_context(info: &InfoStatement) -> ContextLevel {
	match info {
		InfoStatement::Root(_, _) => ContextLevel::Root,
		InfoStatement::Ns(_, _) => ContextLevel::Namespace,
		InfoStatement::Db(_, _) | InfoStatement::Tb(_, _, _) | InfoStatement::Index(_, _, _) => {
			ContextLevel::Database
		}
		InfoStatement::User(user_expr, base, _) => {
			let base_ctx = match base {
				Some(Base::Root) | None => ContextLevel::Root,
				Some(Base::Ns) => ContextLevel::Namespace,
				Some(Base::Db) => ContextLevel::Database,
			};
			base_ctx.max(expr_required_context(user_expr))
		}
	}
}

/// Determine the minimum [`ContextLevel`] required by a [`Literal`].
///
/// Most literals are pure values, but compound literals (arrays, objects,
/// record IDs) can contain inner [`Expr`]s that may reference the database.
fn literal_required_context(lit: &Literal) -> ContextLevel {
	match lit {
		Literal::Array(exprs) | Literal::Set(exprs) => {
			exprs.iter().map(expr_required_context).max().unwrap_or(ContextLevel::Root)
		}
		Literal::Object(entries) => entries
			.iter()
			.map(|e| expr_required_context(&e.value))
			.max()
			.unwrap_or(ContextLevel::Root),
		Literal::RecordId(rid) => record_id_key_required_context(&rid.key),
		_ => ContextLevel::Root,
	}
}

/// Determine the minimum [`ContextLevel`] required by a [`RecordIdKeyLit`].
fn record_id_key_required_context(key: &crate::expr::RecordIdKeyLit) -> ContextLevel {
	use crate::expr::RecordIdKeyLit;
	match key {
		RecordIdKeyLit::Array(exprs) => {
			exprs.iter().map(expr_required_context).max().unwrap_or(ContextLevel::Root)
		}
		RecordIdKeyLit::Object(entries) => entries
			.iter()
			.map(|e| expr_required_context(&e.value))
			.max()
			.unwrap_or(ContextLevel::Root),
		RecordIdKeyLit::Range(range) => {
			let start = match &range.start {
				std::ops::Bound::Included(k) | std::ops::Bound::Excluded(k) => {
					record_id_key_required_context(k)
				}
				std::ops::Bound::Unbounded => ContextLevel::Root,
			};
			let end = match &range.end {
				std::ops::Bound::Included(k) | std::ops::Bound::Excluded(k) => {
					record_id_key_required_context(k)
				}
				std::ops::Bound::Unbounded => ContextLevel::Root,
			};
			start.max(end)
		}
		_ => ContextLevel::Root,
	}
}

/// Determine the minimum [`ContextLevel`] required by an [`Idiom`].
///
/// An idiom's parts may contain inner expressions (WHERE filters, index
/// expressions, method arguments) and graph lookups that need database access.
fn idiom_required_context(idiom: &crate::expr::Idiom) -> ContextLevel {
	idiom.0.iter().map(part_required_context).max().unwrap_or(ContextLevel::Root)
}

/// Determine the minimum [`ContextLevel`] required by a single [`Part`].
fn part_required_context(part: &Part) -> ContextLevel {
	match part {
		// These parts contain inner expressions
		Part::Where(expr) | Part::Value(expr) | Part::Start(expr) => expr_required_context(expr),
		Part::Method(_, args) => {
			args.iter().map(expr_required_context).max().unwrap_or(ContextLevel::Root)
		}
		// Graph/reference lookups need database access
		Part::Lookup(_) => ContextLevel::Database,
		// Recurse: check the instruction for inner expressions
		Part::Recurse(_, _, instruction) => match instruction {
			Some(RecurseInstruction::Shortest {
				expects,
				..
			}) => expr_required_context(expects),
			_ => ContextLevel::Root,
		},
		// Pure structural parts — no context needed
		Part::All
		| Part::Flatten
		| Part::Last
		| Part::First
		| Part::Field(_)
		| Part::Destructure(_)
		| Part::Optional
		| Part::Doc
		| Part::RepeatRecurse => ContextLevel::Root,
	}
}