flowlog-build 0.3.4

//! Common utilities for rule planning.
//!
//! This module provides shared functionality for rule planning operations including:
//! - Semijoin operations (positive and anti-semijoin)
//! - Comparison predicate pushdown
//! - Projection and unused argument removal
//! - Producer-consumer relationship management

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

use super::RulePlanner;
use crate::catalog::{
    ArithmeticPos, AtomArgumentSignature, AtomSignature, Catalog, JoinPredicates, KvPredicates,
};
use crate::planner::{KeyValueLayout, PlanError, TransformationInfo};

// =========================================================================
// Semijoin & Comparison Operations
// =========================================================================
impl RulePlanner {
    /// Attempts to apply semijoin or comparison/fn_call pushdown optimizations.
    ///
    /// This method tries the following optimizations in order:
    /// 1. Comparison pushdown: When a comparison predicate can be pushed to atoms
    /// 2. FnCall pushdown: When a boolean UDF predicate can be pushed to atoms
    /// 3. Positive semijoin: When a positive atom has positive supersets
    /// 4. Anti-semijoin: When a negative atom has positive supersets  
    ///
    /// The reason why we try comparison/fn_call pushdown first is that it can avoid fuse a comparison
    /// with a neg join producer, which is undefined in my understanding.
    ///
    /// WARNING: this method should be called as a total instead of calling individual
    /// semijoin or comparison methods, because they may interfere with each other in fuse logic,
    /// as the reason showed above.
    ///
    /// Returns `true` if any optimization was applied, `false` otherwise.
    pub(super) fn apply_semijoin(&mut self, catalog: &mut Catalog) -> Result<bool, PlanError> {
        // (1) Comparison predicate pushdown
        // When a comparison can be evaluated against specific atoms
        if let Some((lhs_comp_idx, rhs_pos_indices)) = catalog
            .comparison_supersets()
            .iter()
            .enumerate()
            .find(|(_, v)| !v.is_empty())
            .map(|(idx, indices)| (idx, indices.clone()))
        {
            trace!(
                "Comparison pushdown:\n  Comparison: {}\n  RHS atoms: {:?}",
                catalog.comparison_predicate(lhs_comp_idx),
                rhs_pos_indices
                    .iter()
                    .map(|&i| (
                        &catalog.rule().rhs()[catalog.positive_atom_rhs_id(i)],
                        catalog.positive_atom_rhs_id(i)
                    ))
                    .collect::<Vec<_>>()
            );
            return self.apply_comparison_pushdown(catalog, lhs_comp_idx, &rhs_pos_indices);
        }

        // (2) FnCall predicate pushdown
        // When a boolean UDF predicate can be evaluated against specific atoms
        if let Some((lhs_fn_call_idx, rhs_pos_indices)) = catalog
            .fn_call_supersets()
            .iter()
            .enumerate()
            .find(|(_, v)| !v.is_empty())
            .map(|(idx, indices)| (idx, indices.clone()))
        {
            trace!(
                "FnCall pushdown:\n  FnCall: {}\n  RHS atoms: {:?}",
                catalog.fn_call_predicate(lhs_fn_call_idx),
                rhs_pos_indices
                    .iter()
                    .map(|&i| (
                        &catalog.rule().rhs()[catalog.positive_atom_rhs_id(i)],
                        catalog.positive_atom_rhs_id(i)
                    ))
                    .collect::<Vec<_>>()
            );
            return self.apply_fn_call_pushdown(catalog, lhs_fn_call_idx, &rhs_pos_indices);
        }

        // (3) Positive semijoin optimization
        // When a positive atom has positive supersets, we can join them.
        // Note we need premap for both LHS and RHS atoms if they are original EDBs (row format).
        if let Some((lhs_pos_idx, rhs_pos_indices)) = catalog
            .positive_supersets()
            .iter()
            .enumerate()
            .find(|(_, v)| !v.is_empty())
            .map(|(idx, indices)| (idx, indices.clone()))
        {
            self.apply_positive_semijoin_premap(catalog, lhs_pos_idx, &rhs_pos_indices)?;

            trace!(
                "Positive semijoin:\n  LHS atom: ({}, {})\n  RHS atoms: {:?}",
                catalog.rule().rhs()[catalog.positive_atom_rhs_id(lhs_pos_idx)],
                catalog.positive_atom_rhs_id(lhs_pos_idx),
                rhs_pos_indices
                    .iter()
                    .map(|&i| (
                        &catalog.rule().rhs()[catalog.positive_atom_rhs_id(i)],
                        catalog.positive_atom_rhs_id(i)
                    ))
                    .collect::<Vec<_>>()
            );
            return self.apply_positive_semijoin(catalog, lhs_pos_idx, &rhs_pos_indices);
        }

        // (4) Anti-semijoin optimization
        // When a negative atom has positive supersets, we can anti-join them.
        // Note we need premap for both LHS and RHS atoms if they are original EDBs (row format).
        if let Some((lhs_neg_idx, rhs_pos_indices)) = catalog
            .negative_supersets()
            .iter()
            .enumerate()
            .find(|(_, v)| !v.is_empty())
            .map(|(idx, indices)| (idx, indices.clone()))
        {
            self.apply_anti_semijoin_premap(catalog, lhs_neg_idx, &rhs_pos_indices)?;

            trace!(
                "Anti-semijoin:\n  LHS negative atom: ({}, !{})\n  RHS atoms: {:?}",
                catalog.rule().rhs()[catalog.negative_atom_rhs_id(lhs_neg_idx)],
                catalog.negative_atom_rhs_id(lhs_neg_idx),
                rhs_pos_indices
                    .iter()
                    .map(|&i| (
                        &catalog.rule().rhs()[catalog.positive_atom_rhs_id(i)],
                        catalog.positive_atom_rhs_id(i)
                    ))
                    .collect::<Vec<_>>()
            );
            return self.apply_anti_semijoin(catalog, lhs_neg_idx, &rhs_pos_indices);
        }

        Ok(false)
    }

    /// Apply premap for positive semijoin optimization if needed..
    fn apply_positive_semijoin_premap(
        &mut self,
        catalog: &mut Catalog,
        lhs_pos_idx: usize,
        rhs_pos_indices: &[usize],
    ) -> Result<(), PlanError> {
        // Even for empty key key-value layout, we still need premap for EDB relations.
        // ((), (value)) is not the same as (value) in differential dataflow.

        // Process LHS atom for positive semijoin.
        if catalog
            .original_atom_fingerprints()
            .contains(&catalog.positive_atom_fingerprint(lhs_pos_idx))
        {
            self.create_edb_premap_transformations(catalog, lhs_pos_idx, true)?;
        }

        // Process each RHS atom for positive semijoin.
        for &rhs_idx in rhs_pos_indices {
            if catalog
                .original_atom_fingerprints()
                .contains(&catalog.positive_atom_fingerprint(rhs_idx))
            {
                self.create_edb_premap_transformations(catalog, rhs_idx, true)?;
            }
        }
        Ok(())
    }

    /// Applies positive semijoin optimization.
    ///
    /// Positive semijoin: Joins the LHS atom with each RHS atom and keeps only the RHS.
    fn apply_positive_semijoin(
        &mut self,
        catalog: &mut Catalog,
        lhs_pos_idx: usize,
        rhs_pos_indices: &[usize],
    ) -> Result<bool, PlanError> {
        // Extract LHS atom information
        let lhs_pos_args = catalog
            .positive_atom_argument_signature(lhs_pos_idx)
            .clone();
        let lhs_pos_fp = catalog.positive_atom_fingerprint(lhs_pos_idx);
        let left_atom_signature = AtomSignature::new(true, lhs_pos_idx);
        // Build join keys from LHS arguments - these become the join condition
        let lhs_keys: Vec<ArithmeticPos> = lhs_pos_args
            .iter()
            .map(|&s| ArithmeticPos::from_var_signature(s))
            .collect();
        let lhs_key_names: Vec<String> = lhs_pos_args
            .iter()
            .map(|sig| catalog.signature_to_argument_str(sig).clone())
            .collect();
        trace!(
            "Semijoin keys: {:?}",
            lhs_keys
                .iter()
                .map(|pos| (
                    pos,
                    catalog.signature_to_argument_str(pos.init().as_var_signature().unwrap())
                ))
                .collect::<Vec<_>>()
        );

        // Initialize collections for new atoms created by the semijoin
        let mut new_names = Vec::new();
        let mut new_fps = Vec::new();
        let mut new_arg_lists = Vec::new();
        let mut right_sigs = Vec::new();

        // Process each RHS atom for semijoin
        for &rhs_idx in rhs_pos_indices {
            let current_transformation_index = self.transformation_infos.len();

            // Register LHS atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                lhs_pos_fp,
                current_transformation_index,
            )?;

            // Extract RHS atom information
            let rhs_args = catalog.positive_atom_argument_signature(rhs_idx).clone();
            let rhs_fp = catalog.positive_atom_fingerprint(rhs_idx);

            // Register RHS atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                rhs_fp,
                current_transformation_index,
            )?;

            // Build RHS atom argument signatures; keys must mirror the LHS key ordering
            let (new_rhs_args, rhs_keys, rhs_vals) = Self::reorder_rhs_arguments(
                &rhs_args,
                &lhs_key_names,
                catalog,
                "positive semijoin",
            )?;
            trace!(
                "Semijoin RHS values: {:?}",
                rhs_vals
                    .iter()
                    .map(|arg| (
                        arg,
                        catalog.signature_to_argument_str(arg.init().as_var_signature().unwrap())
                    ))
                    .collect::<Vec<_>>()
            );

            // Store join result argument list and signature for catalog update
            new_arg_lists.push(new_rhs_args.clone());
            right_sigs.push(AtomSignature::new(true, rhs_idx));

            // Create the join transformation
            let lhs_name = catalog.positive_atom_name(lhs_pos_idx)?.to_string();
            let rhs_name = catalog.positive_atom_name(rhs_idx)?.to_string();
            let new_name = Self::semijoin_name(&lhs_name, &rhs_name, &lhs_key_names);
            let tx = TransformationInfo::join_to_kv(
                lhs_pos_fp,
                lhs_name,
                rhs_fp,
                rhs_name,
                new_name.clone(),
                KeyValueLayout::new(lhs_keys.clone(), Vec::new()), // LHS: keys only
                KeyValueLayout::new(rhs_keys, rhs_vals.clone()),   // RHS: keys aligned + values
                KeyValueLayout::new(lhs_keys.clone(), rhs_vals),
                JoinPredicates::default(), // no additional comparisons and fn call predicates
            );

            let new_fp = tx.output_info_fp();

            // Register this transformation as a producer
            self.insert_producer(new_fp, current_transformation_index);

            trace!("Positive semijoin transformation:\n{}", tx);

            new_names.push(new_name);
            new_fps.push(new_fp);

            // Store the transformation info
            self.transformation_infos.push(tx);
        }

        // Update catalog with the new joined atoms
        catalog.join_modify(
            left_atom_signature,
            right_sigs,
            new_arg_lists,
            new_names,
            new_fps,
        )?;
        Ok(true)
    }

    /// Apply premap for anti-semijoin optimization.
    fn apply_anti_semijoin_premap(
        &mut self,
        catalog: &mut Catalog,
        lhs_neg_idx: usize,
        rhs_pos_indices: &[usize],
    ) -> Result<(), PlanError> {
        // Even for empty key key-value layout, we still need premap for EDB relations.
        // ((), (value)) is not the same as (value) in differential dataflow.

        // Process LHS atom for anti-semijoin.
        if catalog
            .original_atom_fingerprints()
            .contains(&catalog.negative_atom_fingerprint(lhs_neg_idx))
        {
            self.create_edb_premap_transformations(catalog, lhs_neg_idx, false)?;
        }

        // Process each RHS atom for anti-semijoin.
        for &rhs_idx in rhs_pos_indices {
            if catalog
                .original_atom_fingerprints()
                .contains(&catalog.positive_atom_fingerprint(rhs_idx))
            {
                self.create_edb_premap_transformations(catalog, rhs_idx, true)?;
            }
        }
        Ok(())
    }

    /// Applies anti-semijoin optimization.
    ///
    /// Anti-semijoin: Keeps LHS rows whose join keys do NOT exist in any RHS atom.
    fn apply_anti_semijoin(
        &mut self,
        catalog: &mut Catalog,
        lhs_neg_idx: usize,
        rhs_pos_indices: &[usize],
    ) -> Result<bool, PlanError> {
        // Extract LHS negative atom information
        let lhs_neg_args = catalog
            .negative_atom_argument_signature(lhs_neg_idx)
            .clone();
        let lhs_neg_fp = catalog.negative_atom_fingerprint(lhs_neg_idx);
        let left_atom_signature = AtomSignature::new(false, lhs_neg_idx);
        // Build join keys from LHS arguments - these become the join condition
        let lhs_keys: Vec<ArithmeticPos> = lhs_neg_args
            .iter()
            .map(|&s| ArithmeticPos::from_var_signature(s))
            .collect();
        let lhs_key_names: Vec<String> = lhs_neg_args
            .iter()
            .map(|sig| catalog.signature_to_argument_str(sig).clone())
            .collect();
        trace!(
            "Semijoin keys: {:?}",
            lhs_keys
                .iter()
                .map(|pos| (
                    pos,
                    catalog.signature_to_argument_str(pos.init().as_var_signature().unwrap())
                ))
                .collect::<Vec<_>>()
        );

        // Initialize collections for new atoms created by the anti-semijoin
        let mut new_names = Vec::new();
        let mut new_fps = Vec::new();
        let mut new_arg_lists = Vec::new();
        let mut right_sigs = Vec::new();

        // Process each RHS atom for anti-semijoin
        for &rhs_idx in rhs_pos_indices {
            let current_transformation_index = self.transformation_infos.len();

            // Register LHS negative atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                lhs_neg_fp,
                current_transformation_index,
            )?;

            // Extract RHS atom information
            let rhs_args = catalog.positive_atom_argument_signature(rhs_idx).clone();
            let rhs_fp = catalog.positive_atom_fingerprint(rhs_idx);

            // Register RHS atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                rhs_fp,
                current_transformation_index,
            )?;

            let (new_rhs_args, rhs_keys, rhs_vals) =
                Self::reorder_rhs_arguments(&rhs_args, &lhs_key_names, catalog, "anti-semijoin")?;
            trace!(
                "Semijoin RHS values: {:?}",
                rhs_vals
                    .iter()
                    .map(|arg| (
                        arg,
                        catalog.signature_to_argument_str(arg.init().as_var_signature().unwrap())
                    ))
                    .collect::<Vec<_>>()
            );

            // Store join result argument list and signature for catalog update
            new_arg_lists.push(new_rhs_args.clone());
            right_sigs.push(AtomSignature::new(true, rhs_idx));

            // Create the anti-join transformation
            let lhs_name = catalog.negative_atom_name(lhs_neg_idx)?.to_string();
            let rhs_name = catalog.positive_atom_name(rhs_idx)?.to_string();
            let new_name = Self::antijoin_name(&lhs_name, &rhs_name, &lhs_key_names);
            let tx = TransformationInfo::anti_join_to_kv(
                lhs_neg_fp,
                lhs_name,
                rhs_fp,
                rhs_name,
                new_name.clone(),
                KeyValueLayout::new(lhs_keys.clone(), Vec::new()), // LHS: keys only
                KeyValueLayout::new(rhs_keys.clone(), rhs_vals.clone()), // RHS: values only
                KeyValueLayout::new(rhs_keys, rhs_vals),
            );

            let new_fp = tx.output_info_fp();

            // Update producer transformation index
            self.insert_producer(new_fp, current_transformation_index);

            trace!("Anti semijoin transformation:\n{}", tx);

            new_names.push(new_name);
            new_fps.push(new_fp);

            // Store the transformation info
            self.transformation_infos.push(tx);
        }

        // Update catalog with the new anti-joined atoms
        catalog.join_modify(
            left_atom_signature,
            right_sigs,
            new_arg_lists,
            new_names,
            new_fps,
        )?;
        Ok(true)
    }

    /// Pushes down a comparison predicate onto RHS atoms that fully cover its variables.
    ///
    /// Comparison pushdown moves evaluation of comparison predicates closer to data sources,
    /// reducing the volume of data processed in subsequent operations. This is particularly
    /// effective for selective predicates that can eliminate many tuples early.
    fn apply_comparison_pushdown(
        &mut self,
        catalog: &mut Catalog,
        lhs_comp_idx: usize,
        rhs_pos_indices: &[usize],
    ) -> Result<bool, PlanError> {
        // Initialize collections for new atoms created by the comparison pushdown
        let mut new_names = Vec::new();
        let mut new_fps = Vec::new();
        let mut right_sigs = Vec::new();

        // Process each RHS atom for comparison pushdown
        for &rhs_idx in rhs_pos_indices {
            let current_transformation_index = self.transformation_infos.len();

            // Extract RHS atom information
            let rhs_args = catalog.positive_atom_argument_signature(rhs_idx).clone();
            let rhs_fp = catalog.positive_atom_fingerprint(rhs_idx);

            // Register RHS atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                rhs_fp,
                current_transformation_index,
            )?;

            // Store signature for catalog update
            right_sigs.push(AtomSignature::new(true, rhs_idx));

            let in_vals = rhs_args
                .iter()
                .map(|&sig| ArithmeticPos::from_var_signature(sig))
                .collect::<Vec<_>>();

            let input_name = catalog.positive_atom_name(rhs_idx)?.to_string();
            let cond = catalog.comparison_predicate(lhs_comp_idx).to_string();
            let new_name = Self::filter_name(&input_name, &cond);
            let tx = TransformationInfo::kv_to_kv(
                rhs_fp,
                input_name,
                new_name.clone(),
                catalog.original_atom_fingerprints().contains(&rhs_fp),
                KeyValueLayout::new(vec![], in_vals.clone()),
                KeyValueLayout::new(vec![], in_vals),
                KvPredicates {
                    compare_exprs: vec![
                        catalog.resolve_comparison_predicates(rhs_idx, lhs_comp_idx)?,
                    ],
                    ..Default::default() // no const-eq, no var-eq and fn call predicates
                },
            );

            let new_fp = tx.output_info_fp();

            // Register this transformation as a producer
            self.insert_producer(new_fp, current_transformation_index);

            trace!("Comparison transformation:\n{}", tx);

            new_names.push(new_name);
            new_fps.push(new_fp);

            // Store the transformation info
            self.transformation_infos.push(tx);
        }

        catalog.comparison_modify(lhs_comp_idx, right_sigs, new_names, new_fps)?;
        Ok(true)
    }

    /// Pushes down a boolean UDF predicate onto RHS atoms that fully cover its variables.
    ///
    /// FnCall pushdown moves evaluation of UDF predicates closer to data sources,
    /// reducing the volume of data processed in subsequent operations.
    fn apply_fn_call_pushdown(
        &mut self,
        catalog: &mut Catalog,
        lhs_fn_call_idx: usize,
        rhs_pos_indices: &[usize],
    ) -> Result<bool, PlanError> {
        // Initialize collections for new atoms created by the fn_call pushdown
        let mut new_names = Vec::new();
        let mut new_fps = Vec::new();
        let mut right_sigs = Vec::new();

        // Process each RHS atom for fn_call pushdown
        for &rhs_idx in rhs_pos_indices {
            let current_transformation_index = self.transformation_infos.len();

            // Extract RHS atom information
            let rhs_args = catalog.positive_atom_argument_signature(rhs_idx).clone();
            let rhs_fp = catalog.positive_atom_fingerprint(rhs_idx);

            // Register RHS atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                rhs_fp,
                current_transformation_index,
            )?;

            // Store signature for catalog update
            right_sigs.push(AtomSignature::new(true, rhs_idx));

            let in_vals = rhs_args
                .iter()
                .map(|&sig| ArithmeticPos::from_var_signature(sig))
                .collect::<Vec<_>>();

            let input_name = catalog.positive_atom_name(rhs_idx)?.to_string();
            let cond = catalog.fn_call_predicate(lhs_fn_call_idx).to_string();
            let new_name = Self::filter_name(&input_name, &cond);
            let tx = TransformationInfo::kv_to_kv(
                rhs_fp,
                input_name,
                new_name.clone(),
                catalog.original_atom_fingerprints().contains(&rhs_fp),
                KeyValueLayout::new(vec![], in_vals.clone()),
                KeyValueLayout::new(vec![], in_vals),
                KvPredicates {
                    fn_call_preds: vec![
                        catalog.resolve_fn_call_predicates(rhs_idx, lhs_fn_call_idx)?,
                    ],
                    ..Default::default() // no const-eq, var-eq, or comparisons
                },
            );

            let new_fp = tx.output_info_fp();

            // Register this transformation as a producer
            self.insert_producer(new_fp, current_transformation_index);

            trace!("FnCall transformation:\n{}", tx);

            new_names.push(new_name);
            new_fps.push(new_fp);

            // Store the transformation info
            self.transformation_infos.push(tx);
        }

        catalog.fn_call_modify(lhs_fn_call_idx, right_sigs, new_names, new_fps)?;
        Ok(true)
    }
}

// =========================================================================
// Projection & Unused Argument Removal
// =========================================================================
impl RulePlanner {
    /// Removes arguments across atoms that are provably unused for outputs.
    pub(super) fn remove_unused_arguments(
        &mut self,
        catalog: &mut Catalog,
    ) -> Result<bool, PlanError> {
        // Get groups of unused arguments per atom from catalog analysis
        let groups = catalog.unused_arguments_per_atom();
        if groups.is_empty() {
            return Ok(false);
        }

        let mut applied = false;

        // Process each atom that has unused arguments
        for (atom_signature, to_delete) in groups.clone() {
            let current_transformation_index = self.transformation_infos.len();

            trace!(
                "Unused-arg removal:\n  Atom: {}, {}\n  To delete: {:?}",
                catalog.rule().rhs()[catalog.rhs_index_from_signature(atom_signature)],
                atom_signature,
                to_delete
                    .iter()
                    .map(|sig| (catalog.signature_to_argument_str(sig), sig))
                    .collect::<Vec<_>>()
            );

            // Resolve atom information from signature
            let (args, atom_fp, _atom_id, input_name) = catalog.resolve_atom(&atom_signature)?;
            let input_name = input_name.to_string();

            // Register atom as consumer of this transformation
            self.insert_consumer(
                catalog.original_atom_fingerprints(),
                atom_fp,
                current_transformation_index,
            )?;

            // Create set of positions to drop for efficient filtering
            let drop_set: HashSet<ArithmeticPos> = to_delete
                .iter()
                .map(|&sig| ArithmeticPos::from_var_signature(sig))
                .collect();

            // Get current values for this atom
            let in_vals = args
                .iter()
                .map(|&sig| ArithmeticPos::from_var_signature(sig))
                .collect::<Vec<_>>();

            // Filter out unused arguments from both keys and values
            let out_vals: Vec<ArithmeticPos> = in_vals
                .iter()
                .filter(|&sig| !drop_set.contains(sig))
                .cloned()
                .collect();

            trace!("Output KV layout: keys=[], values={:?}", out_vals);

            // Create projection transformation that removes unused arguments
            let kept_attrs = Self::attrs_from_positions(&out_vals, catalog);
            let new_name = Self::proj_name(&input_name, &kept_attrs);
            let tx = TransformationInfo::kv_to_kv(
                atom_fp,
                input_name,
                new_name.clone(),
                catalog.original_atom_fingerprints().contains(&atom_fp),
                KeyValueLayout::new(vec![], in_vals),
                KeyValueLayout::new(vec![], out_vals),
                KvPredicates::default(), // no const-eq, var-eq, comparisons and fn_call predicates
            );

            let new_fp = tx.output_info_fp();

            // Register this transformation as a producer
            self.insert_producer(new_fp, current_transformation_index);

            trace!("Unused transformation:\n{}", tx);

            // Store the transformation info
            self.transformation_infos.push(tx);

            // Modify the catalog to reflect the projected atom
            catalog.projection_modify(atom_signature, to_delete, new_name, new_fp)?;

            applied = true;
        }

        Ok(applied)
    }
}

// =========================================================================
// Premap for EDB relations
// =========================================================================
impl RulePlanner {
    /// Creates premap transformations for EDB relations that are not in row format.
    /// We do not care about real key-value layout output here, as these premap is just
    /// an indicator we need to map the EDB from read-in row format.
    /// The key-value layout adjustment is handled in later fuse phase if needed.
    pub(super) fn create_edb_premap_transformations(
        &mut self,
        catalog: &mut Catalog,
        atom_idx: usize,
        is_positive: bool,
    ) -> Result<(), PlanError> {
        // Both positive and negative atoms can be pre-mapped.
        let (edb_fp, edb_args) = if is_positive {
            (
                catalog.positive_atom_fingerprint(atom_idx),
                catalog.positive_atom_argument_signature(atom_idx),
            )
        } else {
            (
                catalog.negative_atom_fingerprint(atom_idx),
                catalog.negative_atom_argument_signature(atom_idx),
            )
        };
        let edb_atom_signature = AtomSignature::new(is_positive, atom_idx);
        let current_transformation_index = self.transformation_infos.len();

        let edb_layout = KeyValueLayout::new(
            Vec::new(),
            edb_args
                .iter()
                .map(|&sig| ArithmeticPos::from_var_signature(sig))
                .collect(),
        );

        // EDB premap is a layout-only pass-through: keep the atom's current
        // hierarchical name as both input and output name.
        let edb_name = if is_positive {
            catalog.positive_atom_name(atom_idx)?.to_string()
        } else {
            catalog.negative_atom_name(atom_idx)?.to_string()
        };
        let tx = TransformationInfo::kv_to_kv(
            edb_fp,
            edb_name.clone(),
            edb_name.clone(),
            true,
            edb_layout.clone(),
            edb_layout,
            KvPredicates::default(), // no const-eq, var-eq, comparisons and fn_call predicates
        );

        let new_name = edb_name;
        let new_fp = tx.output_info_fp();

        // Store the transformation info
        self.transformation_infos.push(tx);

        // Register this transformation as consumer of EDB atom
        self.insert_consumer(
            catalog.original_atom_fingerprints(),
            edb_fp,
            current_transformation_index,
        )?;

        // Register this transformation as a producer
        self.insert_producer(new_fp, current_transformation_index);

        // Update catalog to reflect the premap atom
        catalog.map_modify(edb_atom_signature, new_name, new_fp)?;
        Ok(())
    }
}

// =========================================================================
// Private Utilities
// =========================================================================
impl RulePlanner {
    /// Hierarchical name wrappers. These describe the logical operation
    /// applied to the input(s), so each transformation's `output_name`
    /// records the full construction path from EDBs.
    #[inline]
    pub(super) fn proj_name(input_name: &str, attrs: &[String]) -> String {
        format!("π[{}]({})", attrs.join(","), input_name)
    }

    #[inline]
    pub(super) fn filter_name(input_name: &str, cond: &str) -> String {
        format!("σ[{}]({})", cond, input_name)
    }

    /// Extract named-attribute names from `ArithmeticPos` entries.
    /// Non-binding signatures (placeholders, const-eq, var-eq) and
    /// non-variable arithmetic expressions are skipped — they carry no
    /// user-visible name and only add noise to the rendered operator label.
    pub(super) fn attrs_from_positions(
        positions: &[ArithmeticPos],
        catalog: &Catalog,
    ) -> Vec<String> {
        let filters = catalog.filters();
        positions
            .iter()
            .filter_map(|pos| pos.init().as_var_signature())
            .filter(|sig| !filters.is_const_or_var_eq_or_placeholder(sig))
            .map(|sig| catalog.signature_to_argument_str(sig).clone())
            .collect()
    }

    #[inline]
    pub(super) fn join_name(left: &str, right: &str, keys: &[String]) -> String {
        format!("({} ⋈[{}] {})", left, keys.join(","), right)
    }

    #[inline]
    pub(super) fn semijoin_name(left: &str, right: &str, keys: &[String]) -> String {
        format!("({} ⋉[{}] {})", left, keys.join(","), right)
    }

    #[inline]
    pub(super) fn antijoin_name(left: &str, right: &str, keys: &[String]) -> String {
        format!("({} ▷[{}] {})", left, keys.join(","), right)
    }

    /// Orders RHS arguments so join keys come first in the same order as the LHS keys.
    #[allow(clippy::type_complexity)]
    fn reorder_rhs_arguments(
        rhs_args: &[AtomArgumentSignature],
        lhs_key_names: &[String],
        catalog: &Catalog,
        context: &str,
    ) -> Result<
        (
            Vec<AtomArgumentSignature>,
            Vec<ArithmeticPos>,
            Vec<ArithmeticPos>,
        ),
        PlanError,
    > {
        let mut remaining: Vec<(String, AtomArgumentSignature)> = rhs_args
            .iter()
            .map(|&sig| (catalog.signature_to_argument_str(&sig).clone(), sig))
            .collect();
        let mut ordered = Vec::with_capacity(rhs_args.len());

        for key_name in lhs_key_names {
            let position = remaining
                .iter()
                .position(|(name, _)| name == key_name)
                .ok_or_else(|| {
                    PlanError::internal(format!(
                        "reorder_rhs_arguments: RHS missing key {key_name} for {context}"
                    ))
                })?;
            let (_, sig) = remaining.remove(position);
            ordered.push(sig);
        }

        for (_, sig) in remaining {
            ordered.push(sig);
        }

        let key_count = lhs_key_names.len();
        let rhs_keys = ordered[..key_count]
            .iter()
            .map(|&sig| ArithmeticPos::from_var_signature(sig))
            .collect();
        let rhs_vals = ordered[key_count..]
            .iter()
            .map(|&sig| ArithmeticPos::from_var_signature(sig))
            .collect();

        Ok((ordered, rhs_keys, rhs_vals))
    }

    /// Partitions atom arguments into join keys and remaining values.
    pub(super) fn partition_shared_keys(
        catalog: &Catalog,
        lhs_sigs: &[AtomArgumentSignature],
        rhs_sigs: &[AtomArgumentSignature],
    ) -> (
        Vec<ArithmeticPos>,
        Vec<ArithmeticPos>,
        Vec<ArithmeticPos>,
        Vec<ArithmeticPos>,
    ) {
        // Build mapping from argument names to RHS signatures for efficient lookup
        let mut rhs_name_to_sig = HashMap::new();
        for sig in rhs_sigs {
            let name = catalog.signature_to_argument_str(sig);
            rhs_name_to_sig.insert(name.clone(), *sig);
        }

        // Partition LHS arguments into join keys and remaining values
        let mut left_keys = Vec::new();
        let mut left_remains = Vec::new();
        let mut matched_names = Vec::new(); // Keep order for right_keys

        for sig in lhs_sigs {
            let name = catalog.signature_to_argument_str(sig);
            if rhs_name_to_sig.contains_key(name) {
                // This variable appears in both atoms - it's a join key
                left_keys.push(ArithmeticPos::from_var_signature(*sig));
                matched_names.push(name.clone());
            } else {
                // This variable only appears in LHS - it's a payload value
                left_remains.push(ArithmeticPos::from_var_signature(*sig))
            }
        }

        // Build right_keys in the same order as left_keys
        let right_keys: Vec<ArithmeticPos> = matched_names
            .iter()
            .map(|name| {
                let sig = rhs_name_to_sig[name];
                ArithmeticPos::from_var_signature(sig)
            })
            .collect();

        // Collect RHS arguments that don't participate in join (RHS payload)
        let mut right_remains = Vec::new();
        for sig in rhs_sigs {
            let name = catalog.signature_to_argument_str(sig);
            if !matched_names.contains(name) {
                right_remains.push(ArithmeticPos::from_var_signature(*sig));
            }
        }

        (left_keys, left_remains, right_keys, right_remains)
    }
}

// =========================================================================
// Producer-Consumer Relationship Management
// =========================================================================
impl RulePlanner {
    /// Registers a producer transformation for data with a given fingerprint.
    ///
    /// Multiple transformations can produce the same data (i.e., multiple producers per output).
    /// We will deduplicate in the later fusion logic.
    pub(super) fn insert_producer(&mut self, producer_fp: u64, producer_idx: usize) {
        self.producer_consumer
            .entry(producer_fp)
            .and_modify(|(producers, _)| producers.push(producer_idx))
            .or_insert((vec![producer_idx], vec![]));
    }

    /// Registers a consumer transformation for data with a given fingerprint.
    ///
    /// Multiple transformations can consume the same data (i.e., multiple consumers per producer).
    /// This function tracks these consumption relationships, but ignores original atom fingerprints
    /// since those represent input data rather than transformation outputs.
    ///
    /// # Arguments
    /// * `original_atom_fp` - Set of fingerprints for original input atoms
    /// * `producer_fp` - Fingerprint of the data being consumed
    /// * `consumer_idx` - Index of the transformation that consumes this data
    ///
    /// Returns an internal `PlanError` if a consumer registers for a non-original
    /// fingerprint that has no producer yet (planner bug: inconsistent graph).
    pub(super) fn insert_consumer(
        &mut self,
        original_atom_fp: &HashSet<u64>,
        producer_fp: u64,
        consumer_idx: usize,
    ) -> Result<(), PlanError> {
        // If this is an original atom fingerprint (read-in atom),
        // it's not produced by any transformation, so just return
        if original_atom_fp.contains(&producer_fp) {
            return Ok(());
        }

        match self.producer_consumer.get_mut(&producer_fp) {
            Some((_, consumers)) => {
                consumers.push(consumer_idx);
                Ok(())
            }
            None => Err(PlanError::internal(format!(
                "insert_consumer: no producer for transformation fingerprint {producer_fp:#018x}"
            ))),
        }
    }

    /// Retrieves the indices of producer transformations for a given fingerprint.
    #[inline]
    pub(super) fn producer_indices(&self, fp: u64) -> Result<Vec<usize>, PlanError> {
        self.producer_consumer
            .get(&fp)
            .map(|(producers, _)| producers.clone())
            .ok_or_else(|| {
                PlanError::internal(format!(
                    "producer_indices: no producer for transformation fingerprint {fp:#018x}"
                ))
            })
    }

    /// Retrieves the indices of consumer transformations for a given fingerprint.
    #[inline]
    pub(super) fn consumer_indices(&self, fp: u64) -> Result<Vec<usize>, PlanError> {
        self.producer_consumer
            .get(&fp)
            .map(|(_, consumers)| consumers.clone())
            .ok_or_else(|| {
                PlanError::internal(format!(
                    "consumer_indices: no consumers for transformation fingerprint {fp:#018x}"
                ))
            })
    }
}

/// Shared setup helper for planner phase tests.
///
/// Parses a program source and builds a fresh `(RulePlanner, Catalog)` for
/// its first rule. Bypasses the typechecker, so literal constants stay
/// polymorphic (`ConstType::Int(_)`) — matching the behavior of
/// `tests/catalog_errors.rs`.
#[cfg(test)]
pub(super) fn test_setup(src: &str) -> (RulePlanner, Catalog) {
    use crate::common::SourceMap;
    use crate::parser::Program;
    use std::io::Write;

    let mut tmp = tempfile::NamedTempFile::new().expect("tempfile");
    tmp.write_all(src.as_bytes()).expect("write");
    let mut sm = SourceMap::new();
    let program =
        Program::parse(&tmp.path().to_string_lossy(), false, &mut sm).expect("parse failed");
    let rule = program.rules()[0].clone();
    let catalog = Catalog::from_rule(&rule).expect("catalog build failed");
    let planner = RulePlanner::new(rule);
    (planner, catalog)
}