limbo/simulator/model/query/predicate.rs

use std::fmt::Display;

use limbo_sqlite3_parser::{ast, to_sql_string::ToSqlString};
use serde::{Deserialize, Serialize};

use crate::model::{
    query::EmptyContext,
    table::{Table, Value},
};

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Predicate(pub ast::Expr);

impl Predicate {
    pub(crate) fn true_() -> Self {
        Self(ast::Expr::Literal(ast::Literal::Numeric("1".to_string())))
    }

    pub(crate) fn false_() -> Self {
        Self(ast::Expr::Literal(ast::Literal::Numeric("0".to_string())))
    }

    pub(crate) fn test(&self, row: &[Value], table: &Table) -> bool {
        let value = expr_to_value(&self.0, row, table);
        value.map_or(false, |value| value.into_bool())
    }
}

// TODO: In the future pass a Vec<Table> to support resolving a value from another table
// This function attempts to convert an simpler easily computable expression into values
// TODO: In the future, we can try to expand this computation if we want to support harder properties that require us
// to already know more values before hand
fn expr_to_value(expr: &ast::Expr, row: &[Value], table: &Table) -> Option<Value> {
    match expr {
        ast::Expr::DoublyQualified(_, _, ast::Name(col_name))
        | ast::Expr::Qualified(_, ast::Name(col_name))
        | ast::Expr::Id(ast::Id(col_name)) => table
            .columns
            .iter()
            .zip(row.iter())
            .find(|(column, _)| column.name == *col_name)
            .map(|(_, value)| value)
            .cloned(),
        ast::Expr::Literal(literal) => Some(literal.into()),
        ast::Expr::Binary(lhs, op, rhs) => {
            let lhs = expr_to_value(lhs, row, table);
            let rhs = expr_to_value(rhs, row, table);
            match (lhs, rhs) {
                (Some(lhs), Some(rhs)) => Some(lhs.binary_compare(&rhs, *op)),
                _ => None,
            }
        }
        ast::Expr::Like {
            lhs,
            not,
            op,
            rhs,
            escape: _, // TODO: support escape
        } => {
            let lhs = expr_to_value(lhs, row, table);
            let rhs = expr_to_value(rhs, row, table);
            let res = match (lhs, rhs) {
                (Some(lhs), Some(rhs)) => lhs.like_compare(&rhs, *op),
                _ => return None,
            };
            let value: Value = if *not { !res } else { res }.into();
            Some(value)
        }
        // TODO: add unary
        _ => unreachable!("{:?}", expr),
    }
}

impl Display for Predicate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0.to_sql_string(&EmptyContext))
    }
}