llkv_compute/

program.rs

//! Bytecode-style programs for predicate evaluation and domain analysis.

use std::ops::Bound;
use std::sync::Arc;

use llkv_expr::{CompareOp, Expr, Filter, Operator, ScalarExpr, literal::Literal};
use llkv_result::{Error, Result as LlkvResult};
use llkv_types::FieldId;
use rustc_hash::{FxHashMap, FxHashSet};

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
struct ExprKey(*const ());

impl ExprKey {
    #[inline]
    fn new(expr: &Expr<'_, FieldId>) -> Self {
        Self(expr as *const _ as *const ())
    }
}

#[derive(Debug)]
pub struct ProgramSet<'expr> {
    eval: EvalProgram,
    domains: DomainRegistry,
    _root_expr: Arc<Expr<'expr, FieldId>>,
}

impl<'expr> ProgramSet<'expr> {
    pub fn eval_ops(&self) -> &[EvalOp] {
        &self.eval.ops
    }

    pub fn domain(&self, id: DomainProgramId) -> Option<&DomainProgram> {
        self.domains.domain(id)
    }

    pub fn root_domain(&self) -> Option<DomainProgramId> {
        self.domains.root
    }
}

#[derive(Debug, Default)]
struct EvalProgram {
    ops: Vec<EvalOp>,
}

#[derive(Debug)]
pub enum EvalOp {
    PushPredicate(OwnedFilter),
    PushCompare {
        left: ScalarExpr<FieldId>,
        right: ScalarExpr<FieldId>,
        op: CompareOp,
    },
    PushInList {
        expr: ScalarExpr<FieldId>,
        list: Vec<ScalarExpr<FieldId>>,
        negated: bool,
    },
    PushIsNull {
        expr: ScalarExpr<FieldId>,
        negated: bool,
    },
    PushLiteral(bool),
    FusedAnd {
        field_id: FieldId,
        filters: Vec<OwnedFilter>,
    },
    And {
        child_count: usize,
    },
    Or {
        child_count: usize,
    },
    Not {
        domain: DomainProgramId,
    },
}

#[derive(Debug, Clone)]
pub struct OwnedFilter {
    pub field_id: FieldId,
    pub op: OwnedOperator,
}

#[derive(Debug, Clone)]
pub enum OwnedOperator {
    Equals(Literal),
    Range {
        lower: Bound<Literal>,
        upper: Bound<Literal>,
    },
    GreaterThan(Literal),
    GreaterThanOrEquals(Literal),
    LessThan(Literal),
    LessThanOrEquals(Literal),
    In(Vec<Literal>),
    StartsWith {
        pattern: String,
        case_sensitive: bool,
    },
    EndsWith {
        pattern: String,
        case_sensitive: bool,
    },
    Contains {
        pattern: String,
        case_sensitive: bool,
    },
    IsNull,
    IsNotNull,
}

impl OwnedOperator {
    pub fn to_operator(&self) -> Operator<'_> {
        match self {
            Self::Equals(lit) => Operator::Equals(lit.clone()),
            Self::Range { lower, upper } => Operator::Range {
                lower: lower.clone(),
                upper: upper.clone(),
            },
            Self::GreaterThan(lit) => Operator::GreaterThan(lit.clone()),
            Self::GreaterThanOrEquals(lit) => Operator::GreaterThanOrEquals(lit.clone()),
            Self::LessThan(lit) => Operator::LessThan(lit.clone()),
            Self::LessThanOrEquals(lit) => Operator::LessThanOrEquals(lit.clone()),
            Self::In(values) => Operator::In(values.as_slice()),
            Self::StartsWith {
                pattern,
                case_sensitive,
            } => Operator::StartsWith {
                pattern: pattern.clone(),
                case_sensitive: *case_sensitive,
            },
            Self::EndsWith {
                pattern,
                case_sensitive,
            } => Operator::EndsWith {
                pattern: pattern.clone(),
                case_sensitive: *case_sensitive,
            },
            Self::Contains {
                pattern,
                case_sensitive,
            } => Operator::Contains {
                pattern: pattern.clone(),
                case_sensitive: *case_sensitive,
            },
            Self::IsNull => Operator::IsNull,
            Self::IsNotNull => Operator::IsNotNull,
        }
    }
}

impl<'a> From<&'a Operator<'a>> for OwnedOperator {
    fn from(op: &'a Operator<'a>) -> Self {
        match op {
            Operator::Equals(lit) => Self::Equals(lit.clone()),
            Operator::Range { lower, upper } => Self::Range {
                lower: lower.clone(),
                upper: upper.clone(),
            },
            Operator::GreaterThan(lit) => Self::GreaterThan(lit.clone()),
            Operator::GreaterThanOrEquals(lit) => Self::GreaterThanOrEquals(lit.clone()),
            Operator::LessThan(lit) => Self::LessThan(lit.clone()),
            Operator::LessThanOrEquals(lit) => Self::LessThanOrEquals(lit.clone()),
            Operator::In(values) => Self::In((*values).to_vec()),
            Operator::StartsWith {
                pattern,
                case_sensitive,
            } => Self::StartsWith {
                pattern: (*pattern).to_string(),
                case_sensitive: *case_sensitive,
            },
            Operator::EndsWith {
                pattern,
                case_sensitive,
            } => Self::EndsWith {
                pattern: (*pattern).to_string(),
                case_sensitive: *case_sensitive,
            },
            Operator::Contains {
                pattern,
                case_sensitive,
            } => Self::Contains {
                pattern: (*pattern).to_string(),
                case_sensitive: *case_sensitive,
            },
            Operator::IsNull => Self::IsNull,
            Operator::IsNotNull => Self::IsNotNull,
        }
    }
}

impl<'a> From<&'a Filter<'a, FieldId>> for OwnedFilter {
    fn from(filter: &'a Filter<'a, FieldId>) -> Self {
        Self {
            field_id: filter.field_id,
            op: OwnedOperator::from(&filter.op),
        }
    }
}

pub type DomainProgramId = u32;

#[derive(Debug, Default)]
pub struct DomainProgram {
    ops: Vec<DomainOp>,
}

impl DomainProgram {
    pub fn ops(&self) -> &[DomainOp] {
        &self.ops
    }
}

#[derive(Debug)]
pub enum DomainOp {
    PushFieldAll(FieldId),
    PushCompareDomain {
        left: ScalarExpr<FieldId>,
        right: ScalarExpr<FieldId>,
        op: CompareOp,
        fields: Vec<FieldId>,
    },
    PushInListDomain {
        expr: ScalarExpr<FieldId>,
        list: Vec<ScalarExpr<FieldId>>,
        fields: Vec<FieldId>,
        negated: bool,
    },
    PushIsNullDomain {
        expr: ScalarExpr<FieldId>,
        fields: Vec<FieldId>,
        negated: bool,
    },
    PushAllRows,
    Union {
        child_count: usize,
    },
    Intersect {
        child_count: usize,
    },
}

#[derive(Debug, Default)]
struct DomainRegistry {
    programs: Vec<DomainProgram>,
    index: FxHashMap<ExprKey, DomainProgramId>,
    root: Option<DomainProgramId>,
}

impl DomainRegistry {
    fn domain(&self, id: DomainProgramId) -> Option<&DomainProgram> {
        self.programs.get(id as usize)
    }

    fn ensure(&mut self, expr: &Expr<'_, FieldId>) -> DomainProgramId {
        let key = ExprKey::new(expr);
        if let Some(existing) = self.index.get(&key) {
            return *existing;
        }
        let id = self.programs.len() as DomainProgramId;
        let program = compile_domain(expr);
        self.programs.push(program);
        self.index.insert(key, id);
        id
    }
}

#[derive(Debug)]
pub struct ProgramCompiler<'expr> {
    root: Arc<Expr<'expr, FieldId>>,
}

impl<'expr> ProgramCompiler<'expr> {
    pub fn new(root: Arc<Expr<'expr, FieldId>>) -> Self {
        Self { root }
    }

    pub fn compile(self) -> LlkvResult<ProgramSet<'expr>> {
        let ProgramCompiler { root } = self;
        let mut domains = DomainRegistry::default();
        let eval_ops = {
            let root_ref = root.as_ref();
            compile_eval(root_ref, &mut domains)?
        };
        let root_domain = {
            let root_ref = root.as_ref();
            domains.ensure(root_ref)
        };
        domains.root = Some(root_domain);
        Ok(ProgramSet {
            eval: EvalProgram { ops: eval_ops },
            domains,
            _root_expr: root,
        })
    }
}

pub fn normalize_predicate<'expr>(expr: Expr<'expr, FieldId>) -> Expr<'expr, FieldId> {
    llkv_expr::normalization::normalize_predicate(expr)
}

#[derive(Clone, Copy)]
enum EvalVisit<'expr> {
    Enter(&'expr Expr<'expr, FieldId>),
    Exit(&'expr Expr<'expr, FieldId>),
    EmitFused { key: ExprKey, field_id: FieldId },
}

type PredicateVec = Vec<OwnedFilter>;

fn compile_eval<'expr>(
    expr: &'expr Expr<'expr, FieldId>,
    domains: &mut DomainRegistry,
) -> LlkvResult<Vec<EvalOp>> {
    let mut ops = Vec::new();
    let mut stack = vec![EvalVisit::Enter(expr)];
    let mut fused: FxHashMap<ExprKey, PredicateVec> = FxHashMap::default();

    while let Some(frame) = stack.pop() {
        match frame {
            EvalVisit::Enter(node) => match node {
                Expr::And(children) => {
                    if children.is_empty() {
                        return Err(Error::InvalidArgumentError(
                            "AND expression requires at least one predicate".into(),
                        ));
                    }
                    if let Some((field_id, filters)) = gather_fused(children) {
                        let key = ExprKey::new(node);
                        fused.insert(key, filters);
                        stack.push(EvalVisit::EmitFused { key, field_id });
                    } else {
                        stack.push(EvalVisit::Exit(node));
                        for child in children.iter().rev() {
                            stack.push(EvalVisit::Enter(child));
                        }
                    }
                }
                Expr::Or(children) => {
                    if children.is_empty() {
                        return Err(Error::InvalidArgumentError(
                            "OR expression requires at least one predicate".into(),
                        ));
                    }
                    stack.push(EvalVisit::Exit(node));
                    for child in children.iter().rev() {
                        stack.push(EvalVisit::Enter(child));
                    }
                }
                Expr::Not(inner) => {
                    stack.push(EvalVisit::Exit(node));
                    stack.push(EvalVisit::Enter(inner));
                }
                _ => stack.push(EvalVisit::Exit(node)),
            },
            EvalVisit::Exit(node) => match node {
                Expr::Pred(filter) => {
                    ops.push(EvalOp::PushPredicate(OwnedFilter::from(filter)));
                }
                Expr::Compare { left, op, right } => {
                    ops.push(EvalOp::PushCompare {
                        left: left.clone(),
                        right: right.clone(),
                        op: *op,
                    });
                }
                Expr::InList {
                    expr,
                    list,
                    negated,
                } => {
                    ops.push(EvalOp::PushInList {
                        expr: expr.clone(),
                        list: list.clone(),
                        negated: *negated,
                    });
                }
                Expr::IsNull { expr, negated } => {
                    ops.push(EvalOp::PushIsNull {
                        expr: expr.clone(),
                        negated: *negated,
                    });
                }
                Expr::Literal(value) => ops.push(EvalOp::PushLiteral(*value)),
                Expr::And(children) => ops.push(EvalOp::And {
                    child_count: children.len(),
                }),
                Expr::Or(children) => ops.push(EvalOp::Or {
                    child_count: children.len(),
                }),
                Expr::Not(inner) => {
                    let id = domains.ensure(inner);
                    ops.push(EvalOp::Not { domain: id });
                }
                Expr::Exists(_) => {
                    return Err(Error::InvalidArgumentError(
                        "EXISTS predicates are not supported in storage evaluation".into(),
                    ));
                }
            },
            EvalVisit::EmitFused { key, field_id } => {
                let filters = fused
                    .remove(&key)
                    .ok_or_else(|| Error::Internal("missing fused predicate metadata".into()))?;
                ops.push(EvalOp::FusedAnd { field_id, filters });
            }
        }
    }

    Ok(ops)
}

fn gather_fused<'expr>(
    children: &'expr [Expr<'expr, FieldId>],
) -> Option<(FieldId, Vec<OwnedFilter>)> {
    if children.is_empty() {
        return None;
    }
    let mut field: Option<FieldId> = None;
    let mut out: Vec<OwnedFilter> = Vec::with_capacity(children.len());
    for child in children {
        match child {
            Expr::Pred(filter) => {
                if let Some(expected) = field {
                    if expected != filter.field_id {
                        return None;
                    }
                } else {
                    field = Some(filter.field_id);
                }
                out.push(OwnedFilter::from(filter));
            }
            _ => return None,
        }
    }
    field.map(|fid| (fid, out))
}

#[derive(Clone, Copy)]
enum DomainVisit<'expr> {
    Enter(&'expr Expr<'expr, FieldId>),
    Exit(&'expr Expr<'expr, FieldId>),
}

fn compile_domain(expr: &Expr<'_, FieldId>) -> DomainProgram {
    let mut ops = Vec::new();
    let mut stack = vec![DomainVisit::Enter(expr)];

    while let Some(frame) = stack.pop() {
        match frame {
            DomainVisit::Enter(node) => match node {
                Expr::And(children) | Expr::Or(children) => {
                    stack.push(DomainVisit::Exit(node));
                    for child in children.iter().rev() {
                        stack.push(DomainVisit::Enter(child));
                    }
                }
                Expr::Not(inner) => {
                    stack.push(DomainVisit::Exit(node));
                    stack.push(DomainVisit::Enter(inner));
                }
                _ => stack.push(DomainVisit::Exit(node)),
            },
            DomainVisit::Exit(node) => match node {
                Expr::Pred(filter) => ops.push(DomainOp::PushFieldAll(filter.field_id)),
                Expr::Compare { left, op, right } => ops.push(DomainOp::PushCompareDomain {
                    left: left.clone(),
                    right: right.clone(),
                    op: *op,
                    fields: collect_fields([left, right]),
                }),
                Expr::InList {
                    expr,
                    list,
                    negated,
                } => {
                    let mut exprs: Vec<&ScalarExpr<FieldId>> = Vec::with_capacity(list.len() + 1);
                    exprs.push(expr);
                    exprs.extend(list.iter());
                    ops.push(DomainOp::PushInListDomain {
                        expr: expr.clone(),
                        list: list.clone(),
                        fields: collect_fields(exprs),
                        negated: *negated,
                    });
                }
                Expr::IsNull { expr, negated } => {
                    ops.push(DomainOp::PushIsNullDomain {
                        expr: expr.clone(),
                        fields: collect_fields([expr]),
                        negated: *negated,
                    });
                }
                Expr::Literal(_) => {
                    // Boolean literals are defined for every row regardless of value.
                    ops.push(DomainOp::PushAllRows);
                }
                Expr::And(children) => {
                    if children.len() > 1 {
                        ops.push(DomainOp::Intersect {
                            child_count: children.len(),
                        });
                    }
                }
                Expr::Or(children) => {
                    if children.len() > 1 {
                        ops.push(DomainOp::Union {
                            child_count: children.len(),
                        });
                    }
                }
                Expr::Not(_) => {
                    // Domain equals child domain; no-op.
                }
                Expr::Exists(_) => {
                    panic!("EXISTS predicates should not reach storage domain evaluation stage");
                }
            },
        }
    }

    DomainProgram { ops }
}

fn collect_fields<'expr>(
    exprs: impl IntoIterator<Item = &'expr ScalarExpr<FieldId>>,
) -> Vec<FieldId> {
    let mut seen: FxHashSet<FieldId> = FxHashSet::default();
    let mut stack: Vec<&'expr ScalarExpr<FieldId>> = exprs.into_iter().collect();
    while let Some(expr) = stack.pop() {
        match expr {
            ScalarExpr::Column(fid) => {
                seen.insert(*fid);
            }
            ScalarExpr::Literal(_) => {}
            ScalarExpr::Binary { left, right, .. } => {
                stack.push(left);
                stack.push(right);
            }
            ScalarExpr::Compare { left, right, .. } => {
                stack.push(left);
                stack.push(right);
            }
            ScalarExpr::Aggregate(agg) => match agg {
                llkv_expr::expr::AggregateCall::CountStar => {}
                llkv_expr::expr::AggregateCall::Count { expr, .. }
                | llkv_expr::expr::AggregateCall::Sum { expr, .. }
                | llkv_expr::expr::AggregateCall::Total { expr, .. }
                | llkv_expr::expr::AggregateCall::Avg { expr, .. }
                | llkv_expr::expr::AggregateCall::Min(expr)
                | llkv_expr::expr::AggregateCall::Max(expr)
                | llkv_expr::expr::AggregateCall::CountNulls(expr)
                | llkv_expr::expr::AggregateCall::GroupConcat { expr, .. } => {
                    stack.push(expr.as_ref());
                }
            },
            ScalarExpr::GetField { base, .. } => {
                stack.push(base);
            }
            ScalarExpr::Cast { expr, .. } => {
                stack.push(expr.as_ref());
            }
            ScalarExpr::Not(expr) => {
                stack.push(expr.as_ref());
            }
            ScalarExpr::IsNull { expr, .. } => {
                stack.push(expr.as_ref());
            }
            ScalarExpr::Case {
                operand,
                branches,
                else_expr,
            } => {
                if let Some(inner) = operand.as_deref() {
                    stack.push(inner);
                }
                for (when_expr, then_expr) in branches {
                    stack.push(when_expr);
                    stack.push(then_expr);
                }
                if let Some(inner) = else_expr.as_deref() {
                    stack.push(inner);
                }
            }
            ScalarExpr::Coalesce(items) => {
                for item in items {
                    stack.push(item);
                }
            }
            ScalarExpr::Random => {
                // Random does not reference any fields
            }
            ScalarExpr::ScalarSubquery(_) => {
                // Scalar subqueries are resolved separately at planning time
            }
        }
    }
    let mut fields: Vec<FieldId> = seen.into_iter().collect();
    fields.sort_unstable();
    fields
}