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
}
}