use arrow_array::ArrayRef;
use arrow_schema::FieldRef;
use reifydb_core::{
expression::PrefixOperator,
value::{
batch::batch,
column::arith::{
add::add_columns, context::ArithContext, div::div_columns, length_mismatch, mul::mul_columns,
rem::rem_columns, sub::sub_columns,
},
},
};
use reifydb_evaluate::{expression::prefix::prefix_apply, stack::Variable};
use reifydb_value::fragment::Fragment;
use crate::{Result, vm::vm::Vm};
impl<'a> Vm<'a> {
fn exec_binary_column_op<F>(&mut self, op: F, frag: fn() -> Fragment) -> Result<()>
where
F: FnOnce(
&ArithContext,
&(FieldRef, ArrayRef),
&(FieldRef, ArrayRef),
fn() -> Fragment,
) -> Result<(FieldRef, ArrayRef)>,
{
let right = self.pop_as_column()?;
let left = self.pop_as_column()?;
if left.1.len() != right.1.len() {
return Err(length_mismatch(left.1.len(), right.1.len(), &frag()).into());
}
let ctx = self.eval_ctx();
let result = op(&ctx.arith(), &left, &right, frag)?;
self.stack.push(Variable::columns(batch(vec![result])?));
Ok(())
}
pub(crate) fn exec_add(&mut self) -> Result<()> {
self.exec_binary_column_op(add_columns, || Fragment::internal("vm_add"))
}
pub(crate) fn exec_sub(&mut self) -> Result<()> {
self.exec_binary_column_op(sub_columns, || Fragment::internal("vm_sub"))
}
pub(crate) fn exec_mul(&mut self) -> Result<()> {
self.exec_binary_column_op(mul_columns, || Fragment::internal("vm_mul"))
}
pub(crate) fn exec_div(&mut self) -> Result<()> {
self.exec_binary_column_op(div_columns, || Fragment::internal("vm_div"))
}
pub(crate) fn exec_rem(&mut self) -> Result<()> {
self.exec_binary_column_op(rem_columns, || Fragment::internal("vm_rem"))
}
pub(crate) fn exec_negate(&mut self) -> Result<()> {
let col = self.pop_as_column()?;
let frag = Fragment::internal("vm_negate");
let result = prefix_apply(&col, &PrefixOperator::Minus(frag.clone()), &frag)?;
self.stack.push(Variable::columns(batch(vec![result])?));
Ok(())
}
}