polydat_core/iteration/comprehension/optimize/

mod.rs

// Copyright 2024-2026 Jonathan Shook
// SPDX-License-Identifier: Apache-2.0

//! Post-parse optimizer — comprehension_forms.md §10.
//!
//! Required pass upstream of compilation. Takes an AST and
//! produces a canonical, push-down form with these properties
//! (§10.6):
//!
//! 1. **Semantic-preserving.** Output produces the same
//!    dispense sequence (per §9.2).
//! 2. **Idempotent.** `optimize(optimize(C)) == optimize(C)`.
//! 3. **Decidable termination.** Each rewrite strictly
//!    decreases a metadata-derived measure or leaves the AST
//!    unchanged.
//! 4. **Bounds-improving.** Peak memory never grows.
//! 5. **No rejections.** Validity is decided on the tree as written,
//!    before any rewrite, and every rewrite keeps a valid tree valid.
//!
//! ## R-rule catalog
//!
//! Priority order: R0a → R0b → R1 → R2 → R3 → R4 → R5 → R6 →
//! R7 (§10.10.5).
//!
//! - **R0a — identity elimination** (I1–I5): singleton
//!   combinators, trivially-true filter, `order(Lex, None)`.
//! - **R0b — associativity flattening** (A1, A2): nested
//!   union / cartesian collapse to n-ary form.
//! - **R1 — `order(Lex)` → `ORDER_STREAMING`**: the IR
//!   compiler emits `ORDER_STREAMING` for `order(Lex, _)` and
//!   `ORDER_MATERIALIZE` otherwise, carrying the input's
//!   `metadata.index_addressable`. Not an AST rewrite; recorded
//!   in the reducibility catalog as an IR-compilation
//!   eligibility.
//! - **R2 — `order(c, strategy, Some(n))` → `indexed_order`**:
//!   metadata-driven. The working set is the selection, sized by
//!   `strategy_working_set` in `metadata.rs`'s propagation rule,
//!   and `ORDER_MATERIALIZE` selects positions from its input's
//!   shape when it evaluates.
//! - **R3 — `order(filter, Lex, None)` → `filter(order, Lex, None)`**:
//!   AST rewrite. Commute when un-truncated.
//! - **R4 — `filter(union(...), p)` → `union(filter(...))`**:
//!   AST rewrite. Distribute filter into each union child, except
//!   under a non-`Lex` order that ranks the filter's survivors.
//! - **R5 — per-axis filter pushdown**: AST rewrite. Consults
//!   the predicate analyzer (§10.9) for factorization; when
//!   `factorization = PerAxis` and every per-axis sub-predicate is
//!   total over its axis, splits the filter into per-axis filters
//!   wrapping each cartesian child, except under a non-`Lex` order
//!   that ranks the filter's survivors.
//! - **R6 — chained filter folding** (F1): AST rewrite.
//!   `filter(filter(c, p), q)` → `filter(c, p && q)`.
//! - **R7 — order chain folding** (O1): AST rewrite.
//!   `order(order(c, s1, None), s2, t)` → `order(c, s2, t)` when `s2`
//!   selects from its input's shape.
//!
//! ## Module layout
//!
//! - [`finding`] — `ReducibilityFinding`, `Reduction`,
//!   `ComplexityDelta`.
//! - [`r0a_identity`] — I1–I5 elimination.
//! - [`r0b_flatten`] — A1, A2 flattening.
//! - [`r3_commute`] — Lex/filter commute.
//! - [`r4_distribute`] — filter over union.
//! - [`r5_factorize`] — per-axis filter pushdown.
//! - [`r6_filter_fold`] — chained filter folding.
//! - [`r7_order_fold`] — order chain folding.

use super::ast::Comprehension;
use super::predicate::CoordSet;
use super::strategies::for_name;
use super::strategy::StrategyName;
use crate::iteration::comprehension::metadata::Metadata;

pub mod finding;
pub mod r0a_identity;
pub mod r0b_flatten;
pub mod r3_commute;
pub mod r4_distribute;
pub mod r5_factorize;
pub mod r6_filter_fold;
pub mod r7_order_fold;

pub use finding::{
    ComplexityDelta, Ordering as ComplexityOrdering, ReducibilityFinding, Reduction, RuleId,
};

/// Top-level optimizer entry. Applies the R-rule catalog to a
/// fixed point and returns the optimized AST.
///
/// Per comprehension_forms.md §10.6 the function is total — it
/// never rejects. Validation (V1–V9) runs on the tree as written,
/// before this; the optimizer assumes its input is well-formed and
/// keeps it so.
///
/// The optimizer is a thin loop over the reducibility analyzer
/// (§10.10): ask `analyze_reducibility` for a finding; apply
/// its witness if non-empty; repeat. The empty finding ends
/// the loop.
pub fn optimize(ast: Comprehension) -> Comprehension {
    let mut current = ast;
    let mut steps_remaining = max_steps(&current);
    while steps_remaining > 0 {
        match analyze_reducibility(&current) {
            ReducibilityFinding {
                reduction: Some(Reduction::Rewrite { witness, .. }),
                ..
            } => {
                current = witness;
            }
            ReducibilityFinding {
                reduction: Some(Reduction::Replace { with }),
                ..
            } => {
                current = with;
            }
            _ => break,
        }
        steps_remaining -= 1;
    }
    current
}

/// Reducibility analyzer entry — comprehension_forms.md §10.10.
///
/// Walks the AST bottom-up trying each R-rule in priority
/// order. Returns the first non-empty finding; returns
/// the empty finding when no rule fires.
pub fn analyze_reducibility(ast: &Comprehension) -> ReducibilityFinding {
    analyze_at(ast, Place::Free)
}

/// Where a node sits relative to the orders above it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Place {
    /// No non-`Lex` order ranks this node's tuples by position.
    Free,
    /// A non-`Lex` order ranks this node's tuples by their positions in
    /// its input (comprehension_forms.md §5 V5): a filter here keeps
    /// its input's shape, so R4 and R5, which reshape that input, do
    /// not fire. `through_orders` holds when the ranking order selects
    /// from the shape and so reads through an untruncated order here
    /// (§7.4 O1).
    Ranked { through_orders: bool },
}

impl Place {
    /// The place of a child of `node`, which sits at `self`.
    fn of_child(self, node: &Comprehension) -> Place {
        // The input of a ranked filter keeps its shape too: an order
        // that selects from the shape ranks the survivors by their
        // positions beneath the untruncated orders there.
        if let (Comprehension::Filter { .. }, Place::Ranked { through_orders }) = (node, self) {
            return Place::Ranked { through_orders };
        }
        let Comprehension::Order {
            strategy,
            truncation,
            ..
        } = node
        else {
            return Place::Free;
        };
        let inherited = truncation.is_none()
            && self
                == Place::Ranked {
                    through_orders: true,
                };
        if *strategy != StrategyName::Lex {
            Place::Ranked {
                through_orders: inherited || for_name(*strategy).selects_from_shape(),
            }
        } else if inherited {
            self
        } else {
            Place::Free
        }
    }
}

fn analyze_at(ast: &Comprehension, place: Place) -> ReducibilityFinding {
    // Bottom-up: try to rewrite each child first.
    // Rewriting a child returns a new parent that wraps the
    // rewritten child; subsequent rule attempts then see the
    // updated subtree on the next outer-loop iteration.
    if let Some(finding) = try_rewrite_child_first(ast, place) {
        return finding;
    }
    // No rewrite in a child — try rules at this node.
    try_rules_at_node(ast, place)
}

/// Attempt to rewrite a child; return a finding that wraps
/// the rewritten subtree in this node's variant.
fn try_rewrite_child_first(ast: &Comprehension, place: Place) -> Option<ReducibilityFinding> {
    let children: Vec<Comprehension> = ast.children().cloned().collect();
    for (i, child) in children.iter().enumerate() {
        let child_finding = analyze_at(child, place.of_child(ast));
        let rewritten = match child_finding.reduction {
            Some(Reduction::Rewrite { witness, .. }) => witness,
            Some(Reduction::Replace { with }) => with,
            None => continue,
        };
        // Re-build this node with the rewritten child at position i.
        let new_ast = replace_child_at(ast, i, rewritten);
        return Some(ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: child_finding.rule.unwrap_or(RuleId::R0a),
                witness: new_ast,
            }),
            rule: child_finding.rule,
            improvement: child_finding.improvement,
        });
    }
    None
}

/// Try every R-rule at this node in priority order.
/// First fire wins.
fn try_rules_at_node(ast: &Comprehension, place: Place) -> ReducibilityFinding {
    let reshapes = place == Place::Free;
    // R0a — identity elimination
    if let Some(witness) = r0a_identity::apply(ast) {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R0a,
                witness,
            }),
            rule: Some(RuleId::R0a),
            improvement: ComplexityDelta::less_compute(),
        };
    }
    // R0b — associativity flattening
    if let Some(witness) = r0b_flatten::apply(ast) {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R0b,
                witness,
            }),
            rule: Some(RuleId::R0b),
            improvement: ComplexityDelta::less_compute(),
        };
    }
    // R3 — Lex/filter commute
    if let Some(witness) = r3_commute::apply(ast) {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R3,
                witness,
            }),
            rule: Some(RuleId::R3),
            improvement: ComplexityDelta::less_memory(),
        };
    }
    // R4 — filter distributes over union
    if let Some(witness) = r4_distribute::apply(ast).filter(|_| reshapes) {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R4,
                witness,
            }),
            rule: Some(RuleId::R4),
            improvement: ComplexityDelta::less_memory(),
        };
    }
    // R5 — per-axis filter pushdown
    if let Some(witness) =
        r5_factorize::apply(ast, &|p, c| super::predicate::analyze(p, c)).filter(|_| reshapes)
    {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R5,
                witness,
            }),
            rule: Some(RuleId::R5),
            improvement: ComplexityDelta::less_both(),
        };
    }
    // R6 — chained filter folding
    if let Some(witness) = r6_filter_fold::apply(ast) {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R6,
                witness,
            }),
            rule: Some(RuleId::R6),
            improvement: ComplexityDelta::less_compute(),
        };
    }
    // R7 — order chain folding
    if let Some(witness) = r7_order_fold::apply(ast) {
        return ReducibilityFinding {
            reduction: Some(Reduction::Rewrite {
                rule: RuleId::R7,
                witness,
            }),
            rule: Some(RuleId::R7),
            improvement: ComplexityDelta::less_both(),
        };
    }
    // No rule fires.
    ReducibilityFinding {
        reduction: None,
        rule: None,
        improvement: ComplexityDelta::equal(),
    }
}

/// Replace the i-th child of `ast` with `replacement`. Used by
/// the bottom-up walker to plumb child rewrites back into the
/// parent node.
fn replace_child_at(ast: &Comprehension, i: usize, replacement: Comprehension) -> Comprehension {
    match ast {
        Comprehension::Clause { .. } => unreachable!("clause has no children"),
        Comprehension::Cartesian { children } => {
            let mut new_children = children.clone();
            new_children[i] = replacement;
            Comprehension::Cartesian {
                children: new_children,
            }
        }
        Comprehension::Zip { children, mode } => {
            let mut new_children = children.clone();
            new_children[i] = replacement;
            Comprehension::Zip {
                children: new_children,
                mode: *mode,
            }
        }
        Comprehension::Union { children } => {
            let mut new_children = children.clone();
            new_children[i] = replacement;
            Comprehension::Union {
                children: new_children,
            }
        }
        Comprehension::Filter { predicate, .. } => Comprehension::Filter {
            child: Box::new(replacement),
            predicate: predicate.clone(),
        },
        Comprehension::Order {
            strategy,
            truncation,
            seed,
            ..
        } => Comprehension::Order {
            child: Box::new(replacement),
            strategy: *strategy,
            truncation: *truncation,
            seed: *seed,
        },
    }
}

/// Bound on optimizer iterations. Per comprehension_forms.md §10.6
/// (property 3) the optimizer halts because each rewrite strictly
/// decreases a well-founded measure. Iterations are bounded as
/// `node_count^2` to guard against any bug in a rule that
/// would otherwise loop.
fn max_steps(ast: &Comprehension) -> usize {
    let n = ast.node_count();
    n.saturating_mul(n).saturating_add(16)
}

/// Convenience: build a `CoordSet` from a comprehension's
/// coordinate names and its computed metadata. R5 uses this
/// when invoking the predicate analyzer.
pub fn coord_set_for(ast: &Comprehension) -> CoordSet {
    let names = ast.coordinate_names();
    let metadata = ast.metadata();
    coord_set_from(&names, &metadata)
}

fn coord_set_from(names: &[String], metadata: &Metadata) -> CoordSet {
    CoordSet::from_metadata(names, metadata)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::iteration::comprehension::source::{LiteralValue, Source};
    use crate::iteration::comprehension::strategy::StrategyName;

    fn clause(name: &str, vs: &[i64]) -> Comprehension {
        Comprehension::clause(
            name,
            Source::Literal {
                values: vs.iter().map(|n| LiteralValue::Int(*n)).collect(),
            },
        )
    }

    #[test]
    fn optimize_well_formed_ast_does_not_panic() {
        let ast = Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
        let _ = optimize(ast);
    }

    #[test]
    fn optimize_singleton_cartesian_eliminates() {
        // R0a I2: singleton cartesian → its only child.
        let ast = Comprehension::cartesian(vec![clause("k", &[1, 2, 3])]);
        let optimized = optimize(ast);
        assert!(matches!(optimized, Comprehension::Clause { .. }));
    }

    #[test]
    fn optimize_lex_none_eliminates() {
        // R0a I5: order(c, Lex, None) → c.
        let inner = clause("k", &[1, 2, 3]);
        let ast = Comprehension::order(inner.clone(), StrategyName::Lex, None);
        let optimized = optimize(ast);
        assert_eq!(optimized, inner);
    }

    #[test]
    fn optimize_is_idempotent() {
        let ast = Comprehension::cartesian(vec![
            Comprehension::cartesian(vec![clause("a", &[1])]),
            clause("b", &[2]),
        ]);
        let once = optimize(ast.clone());
        let twice = optimize(once.clone());
        assert_eq!(once, twice);
    }
}