uqa-planner 0.5.1

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Payload-safe algebraic rewrites for operator trees.

use uqa_operators::OperatorTree;

use super::{tree_map::map_operator_children, QueryOptimizer};

impl QueryOptimizer {
    // ---------------------------------------------------------------
    // 1. Algebraic simplification
    // ---------------------------------------------------------------

    pub(super) fn simplify_algebra(&self, op: OperatorTree) -> OperatorTree {
        // Recurse first (bottom-up).
        let op = self.recurse_simplify(op);
        match op {
            OperatorTree::Intersect(operands) => {
                // Empty elimination: any empty child collapses the
                // intersection.
                for child in &operands {
                    if child.is_empty() {
                        return OperatorTree::Intersect(Vec::new());
                    }
                }
                // Idempotence is valid only for membership-only operands.
                // Posting-list merges add scores, so structurally equal scored
                // terms must remain distinct.
                let mut seen: Vec<OperatorTree> = Vec::new();
                'outer: for child in operands {
                    for s in &seen {
                        if same_membership_term(s, &child) {
                            continue 'outer;
                        }
                    }
                    seen.push(child);
                }
                let operands = seen;
                // Absorption: drop Union(A, ...) when A also appears
                // in the intersection.
                let mut absorbed: Vec<OperatorTree> = Vec::new();
                for (child_index, child) in operands.iter().enumerate() {
                    if let OperatorTree::Union(union_operands) = child {
                        let drop = is_membership_only(child)
                            && operands.iter().enumerate().any(|(other_index, other)| {
                                other_index != child_index
                                    && union_operands
                                        .iter()
                                        .any(|union_term| same_membership_term(union_term, other))
                            });
                        if drop {
                            continue;
                        }
                    }
                    absorbed.push(child.clone());
                }
                if absorbed.len() == 1 {
                    if let Some(only) = absorbed.pop() {
                        return only;
                    }
                }
                OperatorTree::Intersect(absorbed)
            }
            OperatorTree::Union(operands) => {
                // Drop empty children.
                let mut kept: Vec<OperatorTree> =
                    operands.into_iter().filter(|c| !c.is_empty()).collect();
                // Idempotence is valid only for membership-only operands.
                let mut seen: Vec<OperatorTree> = Vec::new();
                'outer: for child in kept.drain(..) {
                    for s in &seen {
                        if same_membership_term(s, &child) {
                            continue 'outer;
                        }
                    }
                    seen.push(child);
                }
                let operands = seen;
                // Absorption: drop Intersect(A, ...) when A also appears
                // in the union.
                let mut absorbed: Vec<OperatorTree> = Vec::new();
                for (child_index, child) in operands.iter().enumerate() {
                    if let OperatorTree::Intersect(int_operands) = child {
                        let drop = is_membership_only(child)
                            && operands.iter().enumerate().any(|(other_index, other)| {
                                other_index != child_index
                                    && int_operands.iter().any(|intersect_term| {
                                        same_membership_term(intersect_term, other)
                                    })
                            });
                        if drop {
                            continue;
                        }
                    }
                    absorbed.push(child.clone());
                }
                if absorbed.len() == 1 {
                    if let Some(only) = absorbed.pop() {
                        return only;
                    }
                }
                if absorbed.is_empty() {
                    return OperatorTree::Union(Vec::new());
                }
                OperatorTree::Union(absorbed)
            }
            other => other,
        }
    }

    pub(super) fn recurse_simplify(&self, op: OperatorTree) -> OperatorTree {
        map_operator_children(op, |child| self.simplify_algebra(child))
    }
}

/// Whether Boolean composition observes only membership for this subtree.
///
/// `PostingList::merge_union` and `PostingList::merge_intersection` add scores when the same
/// document appears on both sides. They may also carry operator-specific
/// fields. The algebraic identities are therefore safe only for the explicit
/// subset classified exhaustively by `OperatorTree`, whose execution produces
/// default payloads. A new variant opts out until its payload effect has been
/// reviewed.
fn is_membership_only(op: &OperatorTree) -> bool {
    op.is_membership_only()
}

/// Address-independent structural equivalence for operands on which Boolean
/// idempotence and absorption preserve the complete posting-list payload.
fn same_membership_term(left: &OperatorTree, right: &OperatorTree) -> bool {
    if !is_membership_only(left) || !is_membership_only(right) {
        return false;
    }

    // All three zero-child composition forms execute as the same empty set.
    if left.is_empty() || right.is_empty() {
        return left.is_empty() && right.is_empty();
    }

    match (left, right) {
        (
            OperatorTree::Filter {
                field: left_field,
                predicate: left_predicate,
                source: left_source,
            },
            OperatorTree::Filter {
                field: right_field,
                predicate: right_predicate,
                source: right_source,
            },
        ) => {
            left_field == right_field
                && left_predicate == right_predicate
                && same_optional_membership_source(left_source.as_deref(), right_source.as_deref())
        }
        (
            OperatorTree::IndexScan {
                index_name: left_index,
                field: left_field,
                predicate: left_predicate,
            },
            OperatorTree::IndexScan {
                index_name: right_index,
                field: right_field,
                predicate: right_predicate,
            },
        ) => {
            left_index == right_index
                && left_field == right_field
                && left_predicate == right_predicate
        }
        (OperatorTree::Intersect(left), OperatorTree::Intersect(right))
        | (OperatorTree::Union(left), OperatorTree::Union(right)) => {
            same_membership_multiset(left, right)
        }
        (OperatorTree::Complement(left), OperatorTree::Complement(right)) => {
            same_membership_term(left, right)
        }
        (OperatorTree::Composed(left), OperatorTree::Composed(right)) => {
            same_membership_sequence(left, right)
        }
        (
            OperatorTree::VectorExclusion {
                positive: left_positive,
                negative: left_negative,
            },
            OperatorTree::VectorExclusion {
                positive: right_positive,
                negative: right_negative,
            },
        ) => {
            same_membership_term(left_positive, right_positive)
                && same_membership_term(left_negative, right_negative)
        }
        _ => false,
    }
}

fn same_optional_membership_source(
    left: Option<&OperatorTree>,
    right: Option<&OperatorTree>,
) -> bool {
    match (left, right) {
        (None, None) => true,
        (Some(left), Some(right)) => same_membership_term(left, right),
        (None, Some(_)) | (Some(_), None) => false,
    }
}

fn same_membership_sequence(left: &[OperatorTree], right: &[OperatorTree]) -> bool {
    left.len() == right.len()
        && left
            .iter()
            .zip(right)
            .all(|(left, right)| same_membership_term(left, right))
}

fn same_membership_multiset(left: &[OperatorTree], right: &[OperatorTree]) -> bool {
    if left.len() != right.len() {
        return false;
    }

    let mut matched = vec![false; right.len()];
    for left_term in left {
        let Some(index) = right.iter().enumerate().position(|(index, right_term)| {
            !matched[index] && same_membership_term(left_term, right_term)
        }) else {
            return false;
        };
        matched[index] = true;
    }
    true
}