use super::{add_subtract_columns, DecimalProofExpr, DynProofExpr, ProofExpr};
use crate::{
base::{
database::{
try_add_subtract_column_types, Column, ColumnRef, ColumnType, LiteralValue, Table,
},
map::{IndexMap, IndexSet},
proof::{PlaceholderResult, ProofError},
scalar::Scalar,
},
sql::{
proof::{FinalRoundBuilder, VerificationBuilder},
AnalyzeError, AnalyzeResult,
},
utils::log,
};
use alloc::{boxed::Box, string::ToString};
use bumpalo::Bump;
use serde::{Deserialize, Serialize};
use sqlparser::ast::Ident;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SubtractExpr {
lhs: Box<DynProofExpr>,
rhs: Box<DynProofExpr>,
}
impl SubtractExpr {
pub fn try_new(lhs: Box<DynProofExpr>, rhs: Box<DynProofExpr>) -> AnalyzeResult<Self> {
let left_datatype = lhs.data_type();
let right_datatype = rhs.data_type();
try_add_subtract_column_types(left_datatype, right_datatype)
.map(|_| Self { lhs, rhs })
.map_err(|_| AnalyzeError::DataTypeMismatch {
left_type: left_datatype.to_string(),
right_type: right_datatype.to_string(),
})
}
pub fn lhs(&self) -> &DynProofExpr {
&self.lhs
}
pub fn rhs(&self) -> &DynProofExpr {
&self.rhs
}
}
impl ProofExpr for SubtractExpr {
fn data_type(&self) -> ColumnType {
try_add_subtract_column_types(self.lhs.data_type(), self.rhs.data_type())
.expect("Failed to add/subtract column types")
}
fn first_round_evaluate<'a, S: Scalar>(
&self,
alloc: &'a Bump,
table: &Table<'a, S>,
params: &[LiteralValue],
) -> PlaceholderResult<Column<'a, S>> {
let lhs_column: Column<'a, S> = self.lhs.first_round_evaluate(alloc, table, params)?;
let rhs_column: Column<'a, S> = self.rhs.first_round_evaluate(alloc, table, params)?;
let res = add_subtract_columns(lhs_column, rhs_column, alloc, true);
Ok(Column::Decimal75(self.precision(), self.scale(), res))
}
#[tracing::instrument(
name = "proofs.sql.ast.subtract_expr.final_round_evaluate",
level = "info",
skip_all
)]
fn final_round_evaluate<'a, S: Scalar>(
&self,
builder: &mut FinalRoundBuilder<'a, S>,
alloc: &'a Bump,
table: &Table<'a, S>,
params: &[LiteralValue],
) -> PlaceholderResult<Column<'a, S>> {
log::log_memory_usage("Start");
let lhs_column: Column<'a, S> = self
.lhs
.final_round_evaluate(builder, alloc, table, params)?;
let rhs_column: Column<'a, S> = self
.rhs
.final_round_evaluate(builder, alloc, table, params)?;
let res = add_subtract_columns(lhs_column, rhs_column, alloc, true);
log::log_memory_usage("End");
Ok(Column::Decimal75(self.precision(), self.scale(), res))
}
fn verifier_evaluate<S: Scalar>(
&self,
builder: &mut impl VerificationBuilder<S>,
accessor: &IndexMap<Ident, S>,
chi_eval: S,
params: &[LiteralValue],
) -> Result<S, ProofError> {
let lhs_eval = self
.lhs
.verifier_evaluate(builder, accessor, chi_eval, params)?;
let rhs_eval = self
.rhs
.verifier_evaluate(builder, accessor, chi_eval, params)?;
Ok(lhs_eval - rhs_eval)
}
fn get_column_references(&self, columns: &mut IndexSet<ColumnRef>) {
self.lhs.get_column_references(columns);
self.rhs.get_column_references(columns);
}
}
impl DecimalProofExpr for SubtractExpr {}