flowlog-build 0.3.4

//! Stratum planner that plans a stratum (a group of rules).

use std::collections::{HashMap, HashSet};
use tracing::{debug, trace};

use crate::catalog::Catalog;
use crate::common::{Config, SECTION_BAR, SUBSECTION_BAR};
use crate::optimizer::Optimizer;
use crate::parser::{AggregationOperator, FlowLogRule, HeadArg, LoopCondition};
use crate::profiler::{Profiler, with_profiler};
use crate::stratifier::Stratifier;

use crate::planner::{PlanError, RulePlanner, Transformation};

/// Planner for a single stratum (a group of parallel rules).
///
/// A stratum groups rules that can be evaluated parallelly together.
/// The planner owns per-rule planners, deduplicates the
/// generated transformation graphs, separates non-recursive (EDB-only) work from
/// recursive (IDB-dependent) work, and records metadata (enter/leave collections,
/// aggregations) that compiler/executor stages need to run the stratum efficiently.
#[derive(Debug, Default)]
pub struct StratumPlanner {
    /// One planner per rule; these own the raw transformation infos.
    rule_planners: Vec<RulePlanner>,

    /// Whether the stratum is recursive.
    is_recursive: bool,

    /// All deduplicated transformations before recursive/non-recursive separation.
    /// Prefer `non_recursive_transformations` and `recursive_transformations` afterwards.
    transformations: Vec<Transformation>,

    /// Transformations that depend only on EDB inputs; computed once.
    non_recursive_transformations: Vec<Transformation>,

    /// Transformations that touch IDB inputs; re-run during recursion.
    recursive_transformations: Vec<Transformation>,

    /// Fingerprints of collections that enter recursion.
    recursion_enter_collections: Vec<u64>,

    /// Fingerprints of accumulative recursive collections within recursion.
    ///
    /// These use `Variable::new` + concat-with-self feedback semantics (monotone growth).
    recursion_accumulate_recursive_collections: Vec<u64>,

    /// Fingerprints of iterative recursive collections within recursion.
    ///
    /// These use `Variable::new_from` + replace-only feedback semantics (replacement each step).
    recursion_iterative_recursive_collections: Vec<u64>,

    /// Fingerprints of collections that exit recursion.
    recursion_leave_collections: Vec<u64>,

    /// Map each IDB fingerprint to the per-rule head fingerprints that feed it.
    /// Enables the compiler to locate the materialized results per rule.
    idb_to_heads_map: HashMap<u64, Vec<u64>>,

    /// Reverse map: per-rule head fingerprint → IDB fingerprint.
    /// Used by the compiler to type-check rule outputs against their target IDB.
    head_to_idb_map: HashMap<u64, u64>,

    /// Aggregation metadata keyed by IDB fingerprint.
    /// Only populated for rules whose heads contain an aggregation argument.
    /// Values are `(AggregationOperator, output_position, output_arity)`.
    idb_to_aggregation_map: HashMap<u64, (AggregationOperator, usize, usize)>,

    /// Loop condition for this stratum, if it was declared as a loop block.
    /// `None` for plain strata (which run to fixpoint implicitly).
    loop_condition: Option<LoopCondition>,

    /// Atom fingerprints unioned across every rule's rhs. Computed once at
    /// stratum construction so codegen can tell which transformation inputs
    /// are named atoms without re-walking the rule planners.
    atom_fps: HashSet<u64>,
}

impl StratumPlanner {
    /// Build a stratum planner from a stratum.
    pub fn from_rules(
        config: &Config,
        stratum: &[FlowLogRule],
        optimizer: &mut Optimizer,
        profiler: &mut Option<Profiler>,
        stratifier: &Stratifier,
        stratum_idx: usize,
    ) -> Result<Self, PlanError> {
        let is_recursive = stratifier.is_recursive_stratum(stratum_idx);
        let mut catalogs = Vec::with_capacity(stratum.len());
        let mut rule_planners = Vec::with_capacity(stratum.len());

        trace!("New Stratum");

        // Phase 1: Initialize catalogs and run prepare phase
        // to apply local filters/semijoin/comparison before core join planning
        for (i, rule) in stratum.iter().enumerate() {
            trace!("rule[{i}] init:");
            let mut catalog = Catalog::from_rule(rule)?;

            let mut planner = RulePlanner::new(rule.clone());
            planner.prepare(&mut catalog)?;

            debug!(
                "rule[{i}] prepare: {} transformations",
                planner.transformation_infos().len()
            );
            trace!("rule[{i}] after prepare:\n{catalog}");

            catalogs.push(catalog);
            rule_planners.push(planner);
        }

        // Phase 2: Side Information Passing (SIP) optimization
        // to push down filters before the main join and reduce intermediate result size
        if config.sip_enabled() {
            for (i, planner) in rule_planners.iter_mut().enumerate() {
                trace!("rule[{i}] SIP");
                planner.apply_sip(&mut catalogs[i])?;
            }
        }

        // Phase 3: Core planning with optimizer guidance
        // this phase may introduce exponential blowup in intermediate results if not guided properly
        while !catalogs.iter().all(|c| c.is_planned()) {
            let join_decisions = optimizer.plan_stratum(&catalogs);

            for ((planner, catalog), join_decision) in rule_planners
                .iter_mut()
                .zip(catalogs.iter_mut())
                .zip(join_decisions)
            {
                if let Some(join_tuple_index) = join_decision {
                    planner.core(catalog, join_tuple_index)?;
                }
            }
        }

        // Phase 4: Fusion
        // to combine transformations and optimize execution
        for (planner, catalog) in rule_planners.iter_mut().zip(catalogs.iter()) {
            planner.fuse(catalog.original_atom_fingerprints())?;
        }

        // Phase 5: Post-processing
        // align final output to the rule head (vars and arithmetic)
        // to apply final adjustments after fusion, e.g. convert to row type
        for (planner, catalog) in rule_planners.iter_mut().zip(catalogs.iter_mut()) {
            planner.post(catalog)?;
        }

        // Phase 6: Materialize per-rule transformations, rewriting lineage
        // (rhs_id-laden) fingerprints to content-canonical ones so identical
        // operations dedup across rules.
        for planner in rule_planners.iter_mut() {
            planner.materialize();
        }

        // Debug info for per-rule plan trees
        rule_planners.iter().for_each(|rp| {
            debug!("{}", rp);
        });

        // Profiler: record rule logic profiles if enabled
        with_profiler(profiler, |profiler| {
            for rule_planner in rule_planners.iter() {
                profiler.insert_rule(
                    rule_planner.rule().to_string(),
                    rule_planner
                        .transformations()
                        .iter()
                        .map(|tx| {
                            let inputs = if tx.is_unary() {
                                (tx.unary_input().fingerprint(), None)
                            } else {
                                let (left, right) = tx.binary_input();
                                (left.fingerprint(), Some(right.fingerprint()))
                            };
                            (inputs, tx.output().fingerprint())
                        })
                        .collect(),
                );
            }
        });

        // Phase 7: Cross-rule sharing — dedup the per-rule transformations
        // by content fingerprint
        let atom_fps: HashSet<u64> = rule_planners
            .iter()
            .flat_map(RulePlanner::rhs_atom_fps)
            .collect();
        let mut stratum_planner = Self {
            rule_planners,
            is_recursive,
            recursion_accumulate_recursive_collections: stratifier
                .stratum_accumulate_recursive_relation(stratum_idx)
                .to_vec(),
            recursion_iterative_recursive_collections: stratifier
                .stratum_iterative_recursive_relation(stratum_idx)
                .to_vec(),
            recursion_leave_collections: stratifier.stratum_leave_relation(stratum_idx).to_vec(),
            loop_condition: stratifier.loop_condition(stratum_idx).cloned(),
            atom_fps,
            ..Self::default()
        };
        stratum_planner.deduplicate_transformations();

        // Phase 8: Recursive split and metadata mappings
        // this phase to factoring optimizations
        stratum_planner.build_idb_to_heads_map(&catalogs);
        stratum_planner.identify_recursive_transformations(is_recursive);
        stratum_planner
            .build_recursion_enter_collections(stratifier.stratum_available_relations(stratum_idx));
        stratum_planner.build_idb_to_aggregation_map(&catalogs, stratifier)?;

        // Debug info for non-recursive vs recursive transformations.
        debug!("\n{}", stratum_planner);

        Ok(stratum_planner)
    }
}

// =========================================================================
// Getters
// =========================================================================
impl StratumPlanner {
    /// Get non-recursive transformations that depend only on EDBs.
    /// These transformations can be computed once outside recursion.
    #[inline]
    pub(crate) fn non_recursive_transformations(&self) -> &[Transformation] {
        &self.non_recursive_transformations
    }

    /// Retain only the non-recursive transformations matching `f`. Used by
    /// the cross-stratum prune pass to drop transformations whose output
    /// fingerprint was already produced by an earlier stratum's prelude.
    pub(crate) fn retain_non_recursive_transformations<F>(&mut self, f: F)
    where
        F: FnMut(&Transformation) -> bool,
    {
        self.non_recursive_transformations.retain(f);
    }

    /// Get dynamic transformations that depend on IDB collections.
    /// These transformations must be re-evaluated during recursion.
    #[inline]
    pub(crate) fn recursive_transformations(&self) -> &[Transformation] {
        &self.recursive_transformations
    }

    /// Get fingerprints of collections that enter recursion.
    #[inline]
    pub(crate) fn recursion_enter_collections(&self) -> &[u64] {
        &self.recursion_enter_collections
    }

    /// Get fingerprints of accumulative recursive collections within recursion.
    #[inline]
    pub(crate) fn recursion_accumulate_recursive_collections(&self) -> &[u64] {
        &self.recursion_accumulate_recursive_collections
    }

    /// Get fingerprints of iterative recursive collections within recursion.
    #[inline]
    pub(crate) fn recursion_iterative_recursive_collections(&self) -> &[u64] {
        &self.recursion_iterative_recursive_collections
    }

    /// Get fingerprints of collections that leave recursion.
    #[inline]
    pub(crate) fn recursion_leave_collections(&self) -> &[u64] {
        &self.recursion_leave_collections
    }

    /// Output relation fingerprints produced by this stratum.
    #[inline]
    pub(crate) fn output_relations(&self) -> HashSet<u64> {
        self.idb_to_heads_map.keys().cloned().collect()
    }

    /// Get the mapping from each IDB fingerprint to per-rule head fingerprints.
    #[inline]
    pub(crate) fn idb_to_heads_map(&self) -> &HashMap<u64, Vec<u64>> {
        &self.idb_to_heads_map
    }

    /// Get the reverse mapping from per-rule head fingerprint to IDB fingerprint.
    #[inline]
    pub(crate) fn head_to_idb_map(&self) -> &HashMap<u64, u64> {
        &self.head_to_idb_map
    }

    /// Get the mapping from IDB fingerprint to corresponding aggregation.
    /// Returns tuples of `(AggregationOperator, position in output relation,
    /// output arity)`.
    #[inline]
    pub(crate) fn idb_to_aggregation_map(
        &self,
    ) -> &HashMap<u64, (AggregationOperator, usize, usize)> {
        &self.idb_to_aggregation_map
    }

    /// Get the loop condition for this stratum, if it is a loop block.
    #[inline]
    pub(crate) fn loop_condition(&self) -> Option<&LoopCondition> {
        self.loop_condition.as_ref()
    }

    /// Map of atom fingerprint → `"name(arg1, ..., argN)"` label for every
    /// positive/negative atom on any rule's rhs in this stratum. Used by
    /// codegen to annotate operator names with the EDB atom they consume so
    /// the profiler/visualizer can show `[Row -> KV] K:(V0) arc(x, y)` without
    /// any downstream knowledge of atoms.
    #[inline]
    pub(crate) fn atom_fps(&self) -> &HashSet<u64> {
        &self.atom_fps
    }

    /// Check if this stratum is recursive.
    #[inline]
    pub(crate) fn is_recursive(&self) -> bool {
        self.is_recursive
    }

    /// Test-only: per-rule transformations before cross-rule dedup.
    #[cfg(test)]
    pub(crate) fn rule_planners(&self) -> &[RulePlanner] {
        &self.rule_planners
    }
}

impl std::fmt::Display for StratumPlanner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "{}", SECTION_BAR)?;

        let stratum_name = if self.is_recursive {
            "Recursive Stratum"
        } else {
            "Non-Recursive Stratum"
        };
        writeln!(f, "[{stratum_name}] {} rules", self.rule_planners.len())?;
        writeln!(f, "\n{}", SUBSECTION_BAR)?;

        writeln!(f, "Rules:")?;
        for (idx, rule_planner) in self.rule_planners.iter().enumerate() {
            writeln!(f, "  ({:>2}) {:?}", idx, rule_planner.rule())?;
        }
        writeln!(f, "\n{}", SUBSECTION_BAR)?;

        writeln!(
            f,
            "Non-recursive transformations ({}):",
            self.non_recursive_transformations.len()
        )?;
        for (idx, tx) in self.non_recursive_transformations.iter().enumerate() {
            write!(f, "  [N{:>3}] {}", idx, tx)?;
        }
        writeln!(f, "\n{}", SUBSECTION_BAR)?;

        if self.is_recursive {
            writeln!(
                f,
                "Recursive transformations ({}):",
                self.recursive_transformations.len()
            )?;
            for (idx, tx) in self.recursive_transformations.iter().enumerate() {
                write!(f, "  [R{:>3}] {}", idx, tx)?;
            }
        } else {
            writeln!(f, "(Non-recursive stratum: no recursive transformations)")?;
        }

        if !self.idb_to_aggregation_map.is_empty() {
            writeln!(f, "\n{}", SUBSECTION_BAR)?;
            writeln!(f, "IDB to Aggregation Map:")?;
            for (fp, (op, pos, arity)) in &self.idb_to_aggregation_map {
                writeln!(
                    f,
                    "  fp=0x{:016x},\n  op={:?},\n  pos={},\n  arity={}",
                    fp, op, pos, arity
                )?;
            }
        }

        writeln!(f, "{}", SECTION_BAR)
    }
}

// =========================================================================
// Sharing Optimization
// =========================================================================
impl StratumPlanner {
    /// Dedup the per-rule materialized transformations by content
    /// fingerprint (first occurrence wins; order stays topological).
    fn deduplicate_transformations(&mut self) {
        let mut seen = HashSet::new();
        self.transformations = self
            .rule_planners
            .iter()
            .flat_map(|planner| planner.transformations())
            .filter(|tx| seen.insert(tx.output().fingerprint()))
            .cloned()
            .collect();
    }
}

// =========================================================================
// Recursive/Non-Recursive Separation
// =========================================================================
impl StratumPlanner {
    /// Identify non-recursive and recursive transformations for factoring optimization.
    ///
    /// Non-recursive transformations depend only on EDB (base) relations and can be
    /// computed once outside of recursive evaluation loops.
    ///
    /// Recursive transformations depend on IDB (derived) relations and must be
    /// re-evaluated during recursive fixed-point computation.
    ///
    /// # Algorithm
    ///
    /// - Non-recursive strata: All transformations are non-recursive since no recursion occurs.
    /// - Recursive strata: Use left-to-right propagation:
    ///   1. Start with IDB fingerprints (head relations produced in this stratum)
    ///   2. For each transformation from left to right:
    ///      - If any input fingerprint is dynamic, mark transformation as dynamic
    ///      - Add the transformation's output fingerprint to the dynamic set
    ///   3. Remaining transformations are non-recursive
    fn identify_recursive_transformations(&mut self, is_recursive: bool) {
        if !is_recursive {
            // Non-recursive stratum: all transformations are non-recursive
            self.non_recursive_transformations = std::mem::take(&mut self.transformations);
            debug!(
                "Non-recursive stratum: all {} transformations are non-recursive",
                self.non_recursive_transformations.len()
            );
            return;
        }

        // Recursive stratum: propagate dynamic dependencies

        // Step 1: Initialize with output relations fingerprints.
        let mut dynamic_fingerprints: HashSet<u64> =
            self.idb_to_heads_map.keys().copied().collect();

        // Step 2: Left-to-right propagation through transformations
        let mut dynamic_indices = HashSet::new();

        for (i, transformation) in self.transformations.iter().enumerate() {
            // Check if this transformation consumes any dynamic fingerprints
            let consumes_dynamic = if transformation.is_unary() {
                let input_fp = transformation.unary_input().fingerprint();
                dynamic_fingerprints.contains(&input_fp)
            } else {
                let (left, right) = transformation.binary_input();
                dynamic_fingerprints.contains(&left.fingerprint())
                    || dynamic_fingerprints.contains(&right.fingerprint())
            };

            if consumes_dynamic {
                // Mark as dynamic and propagate output fingerprint
                dynamic_indices.insert(i);
                dynamic_fingerprints.insert(transformation.output().fingerprint());
            }
        }

        // Step 3: Separate transformations into non-recursive and recursive vectors
        for (i, transformation) in std::mem::take(&mut self.transformations)
            .into_iter()
            .enumerate()
        {
            if dynamic_indices.contains(&i) {
                self.recursive_transformations.push(transformation);
            } else {
                self.non_recursive_transformations.push(transformation);
            }
        }

        debug!(
            "Recursive stratum: separated {} non-recursive, {} recursive transformations (total: {})",
            self.non_recursive_transformations.len(),
            self.recursive_transformations.len(),
            self.non_recursive_transformations.len() + self.recursive_transformations.len()
        );
    }
}

// =========================================================================
// Metadata Mappings
// =========================================================================
impl StratumPlanner {
    /// Build the fingerprint of collections that enter recursion.
    pub(crate) fn build_recursion_enter_collections(&mut self, available_relations: &HashSet<u64>) {
        // Build sets of input/output fingerprints touched by recursion transformations.
        let mut recursion_input_fps: HashSet<u64> = HashSet::new();
        let mut recursion_output_fps: HashSet<u64> = HashSet::new();
        let mut available_fps = available_relations.clone();

        // Build available output fps from static transformations
        for tx in &self.non_recursive_transformations {
            available_fps.insert(tx.output().fingerprint());
        }

        for tx in &self.recursive_transformations {
            if tx.is_unary() {
                recursion_input_fps.insert(tx.unary_input().fingerprint());
            } else {
                let (left, right) = tx.binary_input();
                recursion_input_fps.insert(left.fingerprint());
                recursion_input_fps.insert(right.fingerprint());
            }
            recursion_output_fps.insert(tx.output().fingerprint());
        }

        // Inputs that never appear as outputs of any recursion transformation are the
        // entering collections for the recursion.
        self.recursion_enter_collections = recursion_input_fps
            .difference(&recursion_output_fps)
            .filter(|fp| available_fps.contains(fp))
            .copied()
            .collect()
    }

    /// Build the mapping from each IDB fingerprint to the rule heads fingerprints
    /// that produce it. Multiple rules may contribute to the same IDB relation.
    ///
    /// Note:
    /// 1. Due to sharing, not every rule has a real IDB -> head mapping, though it may occur
    ///    in the `idb_to_heads_map`.
    fn build_idb_to_heads_map(&mut self, catalogs: &[Catalog]) {
        for (rule_idx, catalog) in catalogs.iter().enumerate() {
            let head_idb_fp = catalog.head_idb_fingerprint();
            let Some(final_tx) = self.rule_planners[rule_idx].transformations().last() else {
                continue;
            };
            let head_fp = final_tx.output().fingerprint();
            // Rules with identical pipelines share one head fp; record it once.
            let heads = self.idb_to_heads_map.entry(head_idb_fp).or_default();
            if !heads.contains(&head_fp) {
                heads.push(head_fp);
            }
            self.head_to_idb_map.insert(head_fp, head_idb_fp);
        }
    }

    /// Build the mapping from each final output collection fingerprint to
    /// its aggregation requirement. Ensures consistent aggregation operator
    /// and position across rules that produce the same relation.
    fn build_idb_to_aggregation_map(
        &mut self,
        catalogs: &[Catalog],
        stratifier: &Stratifier,
    ) -> Result<(), PlanError> {
        // Side map of first-seen head spans used only when constructing
        // the `InconsistentAggregation` diagnostic's `prior_span`.
        let mut prior_spans: HashMap<u64, crate::common::Span> = HashMap::new();

        for catalog in catalogs {
            let head_args = catalog.head_arguments();

            // Rule rewrites (join_modify / projection_modify) rebuild the
            // rule via `FlowLogRule::new`, which leaves a dummy rule span.
            // The head is cloned through unchanged, so prefer its span.
            let current_span = {
                let rule = catalog.rule();
                let head_span = rule.head().span();
                if head_span.is_dummy() {
                    rule.span()
                } else {
                    head_span
                }
            };

            let agg_count = head_args
                .iter()
                .filter(|a| matches!(a, HeadArg::Aggregation(_)))
                .count();
            if agg_count > 1 {
                let head = catalog.rule().head();
                return Err(PlanError::MultipleAggregationsInHead {
                    head_span: current_span,
                    rule_span: catalog.rule().span(),
                    rel: stratifier.display_name(head.head_fingerprint(), head.name()),
                    count: agg_count,
                });
            }

            let Some((pos, op)) = head_args.iter().enumerate().find_map(|(i, arg)| match arg {
                HeadArg::Aggregation(agg) => Some((i, *agg.operator())),
                _ => None,
            }) else {
                continue;
            };

            let arity = head_args.len();
            let head_idb_fp = catalog.head_idb_fingerprint();

            match self.idb_to_aggregation_map.get(&head_idb_fp) {
                Some(&(existing_op, existing_pos, _))
                    if (existing_op, existing_pos) != (op, pos) =>
                {
                    return Err(PlanError::InconsistentAggregation {
                        rule_span: current_span,
                        prior_span: prior_spans
                            .get(&head_idb_fp)
                            .copied()
                            .unwrap_or(crate::common::Span::DUMMY),
                        rel: stratifier.display_name(head_idb_fp, catalog.rule().head().name()),
                        existing_op,
                        existing_pos,
                        found_op: op,
                        found_pos: pos,
                    });
                }
                None => {
                    self.idb_to_aggregation_map
                        .insert(head_idb_fp, (op, pos, arity));
                    prior_spans.insert(head_idb_fp, current_span);
                }
                _ => {} // consistent duplicate — skip
            }
        }
        Ok(())
    }
}