reifydb_engine/vm/volcano/

assert.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

use std::sync::Arc;

use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns, headers::ColumnHeaders};
use reifydb_evaluate::expression::{context::EvalContext, eval::evaluate};
use reifydb_rql::expression::{Expression, name::display_label};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::reifydb_assertions;
use tracing::instrument;

use crate::{
	Result,
	error::EngineError,
	vm::volcano::query::{QueryContext, QueryNode, eval_context_from_query},
};

pub(crate) struct AssertNode {
	input: Box<dyn QueryNode>,
	expressions: Vec<Expression>,
	message: Option<String>,
	context: Option<Arc<QueryContext>>,
}

impl AssertNode {
	pub fn new(input: Box<dyn QueryNode>, expressions: Vec<Expression>, message: Option<String>) -> Self {
		Self {
			input,
			expressions,
			message,
			context: None,
		}
	}

	#[instrument(level = "trace", skip_all, name = "volcano::assert::eval")]
	fn eval(
		session: &EvalContext<'_>,
		columns: &Columns,
		row_count: usize,
		assert_expr: &Expression,
	) -> Result<ColumnWithName> {
		let eval_ctx = session.with_eval(columns.clone(), row_count);
		evaluate(&eval_ctx, assert_expr)
	}

	#[instrument(level = "trace", skip_all, name = "volcano::assert::verify")]
	fn verify(&self, data: &ColumnBuffer, row_count: usize, assert_expr: &Expression) -> Result<()> {
		let frag = assert_expr.full_fragment_owned();
		let label = display_label(assert_expr);
		match data {
			ColumnBuffer::Bool(container) => {
				for i in 0..row_count {
					let valid = container.is_defined(i);
					let value = container.data().get(i);
					if !valid || !value {
						return Err(EngineError::AssertionFailed {
							fragment: frag.clone(),
							message: self.message.clone().unwrap_or_default(),
							expression: Some(label.text().to_string()),
						}
						.into());
					}
				}
			}
			ColumnBuffer::Option {
				inner,
				bitvec,
			} => match inner.as_ref() {
				ColumnBuffer::Bool(container) => {
					for i in 0..row_count {
						let defined = i < bitvec.len() && bitvec.get(i);
						let valid = defined && container.is_defined(i);
						let value = valid && container.data().get(i);
						if !value {
							return Err(EngineError::AssertionFailed {
								fragment: frag.clone(),
								message: self.message.clone().unwrap_or_default(),
								expression: Some(label.text().to_string()),
							}
							.into());
						}
					}
				}
				_ => {
					return Err(EngineError::AssertionFailed {
						fragment: frag.clone(),
						message: "assert expression must evaluate to a boolean".to_string(),
						expression: Some(label.text().to_string()),
					}
					.into());
				}
			},
			_ => {
				return Err(EngineError::AssertionFailed {
					fragment: frag.clone(),
					message: "assert expression must evaluate to a boolean".to_string(),
					expression: Some(label.text().to_string()),
				}
				.into());
			}
		}
		Ok(())
	}
}

impl QueryNode for AssertNode {
	#[instrument(level = "trace", skip_all, name = "volcano::assert::initialize")]
	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
		self.context = Some(Arc::new(ctx.clone()));
		self.input.initialize(rx, ctx)?;
		Ok(())
	}

	#[instrument(level = "trace", skip_all, name = "volcano::assert::next")]
	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<Columns>> {
		reifydb_assertions! {
			assert!(self.context.is_some(), "AssertNode::next() called before initialize()");
		}
		let stored_ctx = self.context.as_ref().unwrap();

		if let Some(columns) = self.input.next(rx, ctx)? {
			let row_count = columns.row_count();
			let session = eval_context_from_query(stored_ctx);

			for assert_expr in &self.expressions {
				let result = Self::eval(&session, &columns, row_count, assert_expr)?;
				self.verify(result.data(), row_count, assert_expr)?;
			}

			Ok(Some(columns))
		} else {
			Ok(None)
		}
	}

	fn headers(&self) -> Option<ColumnHeaders> {
		self.input.headers()
	}
}

pub(crate) struct AssertWithoutInputNode {
	expressions: Vec<Expression>,
	message: Option<String>,
	context: Option<Arc<QueryContext>>,
	done: bool,
}

impl AssertWithoutInputNode {
	pub fn new(expressions: Vec<Expression>, message: Option<String>) -> Self {
		Self {
			expressions,
			message,
			context: None,
			done: false,
		}
	}

	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::eval")]
	fn eval(session: &EvalContext<'_>, assert_expr: &Expression) -> Result<ColumnWithName> {
		let eval_ctx = session.with_eval_empty();
		evaluate(&eval_ctx, assert_expr)
	}

	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::verify")]
	fn verify(&self, data: &ColumnBuffer, assert_expr: &Expression) -> Result<()> {
		let frag = assert_expr.full_fragment_owned();
		let label = display_label(assert_expr);
		match data {
			ColumnBuffer::Bool(container) => {
				let valid = container.is_defined(0);
				let value = container.data().get(0);
				if !valid || !value {
					return Err(EngineError::AssertionFailed {
						fragment: frag.clone(),
						message: self.message.clone().unwrap_or_default(),
						expression: Some(label.text().to_string()),
					}
					.into());
				}
			}
			ColumnBuffer::Option {
				inner,
				bitvec,
			} => match inner.as_ref() {
				ColumnBuffer::Bool(container) => {
					let defined = !bitvec.is_empty() && bitvec.get(0);
					let valid = defined && container.is_defined(0);
					let value = valid && container.data().get(0);
					if !value {
						return Err(EngineError::AssertionFailed {
							fragment: frag.clone(),
							message: self.message.clone().unwrap_or_default(),
							expression: Some(label.text().to_string()),
						}
						.into());
					}
				}
				_ => {
					return Err(EngineError::AssertionFailed {
						fragment: frag.clone(),
						message: "assert expression must evaluate to a boolean".to_string(),
						expression: Some(label.text().to_string()),
					}
					.into());
				}
			},
			_ => {
				return Err(EngineError::AssertionFailed {
					fragment: frag.clone(),
					message: "assert expression must evaluate to a boolean".to_string(),
					expression: Some(label.text().to_string()),
				}
				.into());
			}
		}
		Ok(())
	}
}

impl QueryNode for AssertWithoutInputNode {
	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::initialize")]
	fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
		self.context = Some(Arc::new(ctx.clone()));
		Ok(())
	}

	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::next")]
	fn next<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<Columns>> {
		if self.done {
			return Ok(None);
		}
		self.done = true;

		reifydb_assertions! {
			assert!(self.context.is_some(), "AssertWithoutInputNode::next() called before initialize()");
		}
		let stored_ctx = self.context.as_ref().unwrap();
		let session = eval_context_from_query(stored_ctx);

		for assert_expr in &self.expressions {
			let result = Self::eval(&session, assert_expr)?;
			self.verify(result.data(), assert_expr)?;
		}

		Ok(None)
	}

	fn headers(&self) -> Option<ColumnHeaders> {
		None
	}
}