surrealdb-core 3.2.0

//! IfElse operator with deferred planning.
//!
//! The IfElsePlan operator evaluates IF/ELSE IF/ELSE conditional branches,
//! using deferred planning like SequencePlan. Each condition is evaluated
//! sequentially, and the first truthy branch's body is executed.

use std::sync::Arc;

use futures::stream;
use surrealdb_types::{SqlFormat, ToSql};

use crate::err::Error;
use crate::exec::context::{ContextLevel, ExecutionContext};
use crate::exec::plan_or_compute::{evaluate_expr_at_depth, expr_required_context};
use crate::exec::{
	AccessMode, CardinalityHint, ExecOperator, FlowResult, OperatorMetrics, ValueBatch,
	ValueBatchStream,
};
use crate::expr::{ControlFlow, Expr};
use crate::val::Value;

/// IfElse operator with deferred planning.
///
/// Stores the original condition-body pairs and optional else body.
/// Plans and executes each condition at runtime, executing the first
/// truthy branch's body.
///
/// Example:
/// ```surql
/// IF $x > 10 {
///     "large"
/// } ELSE IF $x > 5 {
///     "medium"
/// } ELSE {
///     "small"
/// }
/// ```
#[derive(Debug)]
pub struct IfElsePlan {
	/// Condition-body pairs: Vec<(condition_expr, body_expr)>
	pub branches: Vec<(Expr, Expr)>,
	/// Metrics for EXPLAIN ANALYZE
	pub(crate) metrics: Arc<OperatorMetrics>,
	/// Optional else body
	pub else_body: Option<Expr>,
	/// Expression-nesting depth recorded when this operator was planned. The
	/// deferred condition/body planning below is seeded with it so re-entry
	/// nodes (eval/UDF) keep counting toward `max_computation_depth`.
	plan_depth: u32,
}

impl IfElsePlan {
	pub(crate) fn new(
		branches: Vec<(Expr, Expr)>,
		else_body: Option<Expr>,
		plan_depth: u32,
	) -> Self {
		Self {
			branches,
			else_body,
			metrics: Arc::new(OperatorMetrics::new()),
			plan_depth,
		}
	}
}
impl ExecOperator for IfElsePlan {
	fn name(&self) -> &'static str {
		"IfElse"
	}

	fn attrs(&self) -> Vec<(String, String)> {
		let mut attrs = vec![("branches".to_string(), self.branches.len().to_string())];
		if self.else_body.is_some() {
			attrs.push(("has_else".to_string(), "true".to_string()));
		}
		attrs
	}

	fn required_context(&self) -> ContextLevel {
		// Derive the required context from condition and body expressions
		let branches_ctx = self
			.branches
			.iter()
			.flat_map(|(cond, body)| [expr_required_context(cond), expr_required_context(body)])
			.max()
			.unwrap_or(ContextLevel::Root);
		let else_ctx =
			self.else_body.as_ref().map(expr_required_context).unwrap_or(ContextLevel::Root);
		branches_ctx.max(else_ctx)
	}

	fn access_mode(&self) -> AccessMode {
		// Check if any branch requires write access
		let branches_read_only =
			self.branches.iter().all(|(cond, body)| cond.read_only() && body.read_only());
		let else_read_only = self.else_body.as_ref().map(|e| e.read_only()).unwrap_or(true);

		if branches_read_only && else_read_only {
			AccessMode::ReadOnly
		} else {
			AccessMode::ReadWrite
		}
	}

	fn cardinality_hint(&self) -> CardinalityHint {
		CardinalityHint::AtMostOne
	}

	fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
		let branches = self.branches.clone();
		let else_body = self.else_body.clone();
		// Branch bodies are re-planned one re-entry deeper than this operator, so
		// the depth count continues at `plan_depth + 1` toward `max_computation_depth`.
		let depth = self.plan_depth + 1;
		let ctx = ctx.clone();

		let stream =
			stream::once(async move { execute_ifelse(&branches, &else_body, &ctx, depth).await });

		Ok(Box::pin(stream))
	}

	fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
		// With deferred planning, we don't have pre-built children
		vec![]
	}

	fn metrics(&self) -> Option<&OperatorMetrics> {
		Some(&self.metrics)
	}

	fn is_scalar(&self) -> bool {
		// IF/ELSE expressions return a single value
		true
	}
}

/// Execute the IF/ELSE logic with deferred planning.
async fn execute_ifelse(
	branches: &[(Expr, Expr)],
	else_body: &Option<Expr>,
	ctx: &ExecutionContext,
	depth: u32,
) -> crate::expr::FlowResult<ValueBatch> {
	for (cond, body) in branches {
		if ctx.cancellation().is_cancelled() {
			return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
		}
		let cond_value = evaluate_expr_at_depth(cond, ctx, depth).await?;

		if cond_value.is_truthy() {
			let result = evaluate_expr_at_depth(body, ctx, depth).await?;
			return Ok(ValueBatch {
				values: vec![result],
			});
		}
	}

	// No branch matched - check for else body
	if let Some(else_expr) = else_body {
		let result = evaluate_expr_at_depth(else_expr, ctx, depth).await?;
		Ok(ValueBatch {
			values: vec![result],
		})
	} else {
		// No else - return NONE
		Ok(ValueBatch {
			values: vec![Value::None],
		})
	}
}

impl ToSql for IfElsePlan {
	fn fmt_sql(&self, f: &mut String, fmt: SqlFormat) {
		for (i, (cond, body)) in self.branches.iter().enumerate() {
			if i == 0 {
				f.push_str("IF ");
			} else {
				f.push_str(" ELSE IF ");
			}
			cond.fmt_sql(f, fmt);
			f.push(' ');
			body.fmt_sql(f, fmt);
		}
		if let Some(ref else_body) = self.else_body {
			f.push_str(" ELSE ");
			else_body.fmt_sql(f, fmt);
		}
	}
}