xlog_logic/

compile.rs

//! Compilation pipeline for XLOG programs
//!
//! This module compiles the core XLOG AST into execution plans. The compilation
//! process consists of:
//!
//! 1. **Parsing**: Convert source text to AST (`parser::parse_program`)
//! 2. **Stratification**: Analyze negation/aggregation dependencies (`stratify::stratify`)
//! 3. **Lowering**: Transform AST to Relational IR (`lower::Lowerer::lower_program`)
//!
//! The `Compiler` struct orchestrates these phases and provides a single entry
//! point via the `compile` method. User-defined functions are a source-level
//! normalization step: callers using this low-level compiler must expand them
//! before compilation. Execution-facing `LogicProgram` APIs perform that
//! normalization for parsed and source inputs.

use std::path::{Path, PathBuf};

use xlog_core::{Result, XlogError};
use xlog_ir::{ExecutionPlan, GeneratedQueryRuleProvenance};
use xlog_stats::{StatsManager, StatsSnapshot};

use crate::compiler_config::CompilerConfig;
use crate::list_normalize::normalize_list_builtins_owned;
use crate::lower::Lowerer;
use crate::magic_sets::rewrite_magic_sets_owned;
use crate::meta_normalize::normalize_meta_builtins_owned;
use crate::module::ModuleError;
use crate::optimizer::Optimizer;
use crate::parser::parse_program;
use crate::resolver::ModuleResolver;
use crate::stratify::stratify;
use crate::{BodyLiteral, Program, Query, Rule as AstRule, Term};

/// The XLOG compiler orchestrates the full compilation pipeline.
///
/// This is the core AST-to-RIR compiler. It does not resolve imports or expand
/// user-defined functions; callers must perform those source-level steps first,
/// or use the execution-facing `LogicProgram` compilation APIs.
///
/// # Example
///
/// ```ignore
/// use xlog_logic::compile::Compiler;
///
/// let mut compiler = Compiler::new();
/// let plan = compiler.compile(r#"
///     edge(1, 2).
///     edge(2, 3).
///     reach(X, Y) :- edge(X, Y).
///     reach(X, Z) :- reach(X, Y), edge(Y, Z).
/// "#)?;
/// ```
pub struct Compiler {
    lowerer: Lowerer,
}

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use xlog_core::{RelId, Schema};

impl Default for Compiler {
    fn default() -> Self {
        Self::new()
    }
}

impl Compiler {
    /// Create a new compiler instance.
    pub fn new() -> Self {
        Self {
            lowerer: Lowerer::new(),
        }
    }

    /// Set the maximum active rules for TensorMaskedJoin (16..=128).
    pub fn set_max_active_rules(&mut self, max: usize) {
        self.lowerer.set_max_active_rules(max);
    }

    /// Compile XLOG source code into an execution plan.
    ///
    /// This is the main entry point for core compilation. It chains together:
    /// 1. Parsing (source → AST)
    /// 2. Stratification (analyze dependencies, check for cycles)
    /// 3. Lowering (AST → Relational IR execution plan)
    ///
    /// Function definitions and calls must already be expanded into the core
    /// language. Use an execution-facing `LogicProgram` API when source-level
    /// import resolution and function normalization are required.
    ///
    /// # Arguments
    ///
    /// * `source` - The XLOG source code as a string
    ///
    /// # Returns
    ///
    /// * `Ok(ExecutionPlan)` - The compiled execution plan ready for execution
    /// * `Err(XlogError)` - If any compilation phase fails:
    ///   - `XlogError::Parse` - Syntax errors in the source
    ///   - `XlogError::StratificationCycle` - Unstratifiable negation/aggregation
    ///   - `XlogError::Compilation` - Other semantic errors
    ///
    /// # Example
    ///
    /// ```ignore
    /// let mut compiler = Compiler::new();
    ///
    /// // Compile a simple transitive closure program
    /// let plan = compiler.compile(r#"
    ///     edge(1, 2).
    ///     edge(2, 3).
    ///     reach(X, Y) :- edge(X, Y).
    ///     reach(X, Z) :- reach(X, Y), edge(Y, Z).
    /// "#)?;
    ///
    /// // The plan can now be executed by xlog-runtime
    /// ```
    pub fn compile(&mut self, source: &str) -> Result<ExecutionPlan> {
        self.compile_with_stats_snapshot(source, None)
    }

    /// Compile XLOG source code into an execution plan, optionally seeding the optimizer
    /// with a runtime statistics snapshot.
    ///
    /// This entry point delegates through the composable config-aware API
    /// with `CompilerConfig::default()`, which preserves existing triangle,
    /// 4-cycle, recursive, and selectivity-aware dispatch behavior
    /// bit-identically.
    pub fn compile_with_stats_snapshot(
        &mut self,
        source: &str,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        self.compile_with_config_and_stats_snapshot(
            source,
            &CompilerConfig::default(),
            stats_snapshot,
        )
    }

    /// Composable entry point that accepts a `CompilerConfig`.
    ///
    /// Default-config callers should keep using `compile()` /
    /// `compile_with_stats_snapshot()`. This entry point exists so callers can
    /// enable the variable-ordering cost model per call without an environment
    /// override.
    pub fn compile_with_config_and_stats_snapshot(
        &mut self,
        source: &str,
        config: &CompilerConfig,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        let program = parse_program(source)?;
        self.compile_owned_program_with_config_and_stats_snapshot(program, config, stats_snapshot)
    }

    /// Compile a parsed XLOG program into an execution plan.
    ///
    /// This is useful for callers that want to inspect the AST (facts, queries,
    /// constraints) while compiling without reparsing. The parsed program must
    /// already be normalized to the core language, including function expansion.
    pub fn compile_program(&mut self, program: &Program) -> Result<ExecutionPlan> {
        self.compile_program_with_stats_snapshot(program, None)
    }

    /// Compile a program whose authored constraint identities were prepared at
    /// the outer full-program boundary.
    pub fn compile_prepared_program(&mut self, program: &Program) -> Result<ExecutionPlan> {
        self.compile_prepared_program_with_stats_snapshot(program, None)
    }

    /// Compile a parsed XLOG program into an execution plan, optionally seeding the optimizer.
    ///
    /// Delegates to [`Self::compile_program_with_config_and_stats_snapshot`]
    /// with `CompilerConfig::default()`.
    pub fn compile_program_with_stats_snapshot(
        &mut self,
        program: &Program,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        self.compile_program_with_config_and_stats_snapshot(
            program,
            &CompilerConfig::default(),
            stats_snapshot,
        )
    }

    /// Compile an already prepared program with an optional statistics snapshot.
    pub fn compile_prepared_program_with_stats_snapshot(
        &mut self,
        program: &Program,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        self.compile_prepared_program_with_config_and_stats_snapshot(
            program,
            &CompilerConfig::default(),
            stats_snapshot,
        )
    }

    /// Validate a parsed program through the production lowering contracts without
    /// requiring the dependency graph to be stratifiable.
    ///
    /// This is used by epistemic preparation, where a supported negated modal cycle
    /// must reach well-founded execution rather than fail ordinary stratification.
    /// The normal preprocessing, negation-safety, schema/type, constant, arithmetic,
    /// and projection checks still run unchanged.
    pub(crate) fn validate_program_without_stratification(
        &mut self,
        program: &Program,
    ) -> Result<()> {
        program.validate_prepared_authored_constraint_identity()?;
        let program = run_frontend_passes(program.clone())?;
        self.lowerer.validate_program_without_plan(&program)
    }

    /// Composable program-level entry point.
    ///
    /// `config` is currently consumed only by the promoter when it wires the
    /// variable-ordering cost model. With `CompilerConfig::default()`, the
    /// promoter keeps the default variable order.
    pub fn compile_program_with_config_and_stats_snapshot(
        &mut self,
        program: &Program,
        config: &CompilerConfig,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        self.compile_owned_program_with_config_and_stats_snapshot(
            program.clone(),
            config,
            stats_snapshot,
        )
    }

    /// [`Self::compile_program_with_config_and_stats_snapshot`] for a program
    /// the caller hands over by value (no clone of the AST).
    pub fn compile_owned_program_with_config_and_stats_snapshot(
        &mut self,
        mut program: Program,
        config: &CompilerConfig,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        if program.authored_constraint_source_bound.is_some() {
            program.validate_prepared_authored_constraint_identity()?;
        } else {
            program.prepare_authored_constraint_identity_at_root()?;
        }
        self.compile_prepared_owned_program_with_config_and_stats_snapshot(
            program,
            config,
            stats_snapshot,
        )
    }

    /// Config-aware compilation for an already prepared program.
    pub fn compile_prepared_program_with_config_and_stats_snapshot(
        &mut self,
        program: &Program,
        config: &CompilerConfig,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        self.compile_prepared_owned_program_with_config_and_stats_snapshot(
            program.clone(),
            config,
            stats_snapshot,
        )
    }

    /// [`Self::compile_prepared_program_with_config_and_stats_snapshot`] for a
    /// program handed over by value.
    ///
    /// The frontend passes (desugar → meta → list → magic sets) each take and
    /// return the `Program` by value, so a fact-heavy AST is never cloned
    /// between passes; the `&Program` entry points clone exactly once.
    pub fn compile_prepared_owned_program_with_config_and_stats_snapshot(
        &mut self,
        program: Program,
        config: &CompilerConfig,
        stats_snapshot: Option<&StatsSnapshot>,
    ) -> Result<ExecutionPlan> {
        program.validate_prepared_authored_constraint_identity()?;
        let generated_query_heads = generated_query_heads_for_program(&program)?;
        let program = run_frontend_passes(program)?;

        // Phase 2: Stratify (analyze dependencies, detect cycles)
        let strata = stratify(&program).map_err(map_stratification_to_naf_error)?;

        // Convert strata to the format expected by the lowerer
        let strata_preds: Vec<Vec<String>> = strata.into_iter().map(|s| s.predicates).collect();

        // Phase 3: Lower AST to execution plan
        self.lowerer.set_strata(strata_preds);

        // If we have predicate names for the snapshot, use them to seed lowering-time
        // join ordering with better cardinality estimates.
        let mut cardinality_hints: HashMap<String, u64> = HashMap::new();
        if let Some(snapshot) = stats_snapshot {
            if !snapshot.rel_names.is_empty() {
                let rel_name_by_id: HashMap<RelId, &str> = snapshot
                    .rel_names
                    .iter()
                    .map(|(id, name)| (*id, name.as_str()))
                    .collect();
                for rel in &snapshot.relations {
                    if let Some(name) = rel_name_by_id.get(&rel.rel_id) {
                        cardinality_hints.insert((*name).to_string(), rel.cardinality);
                    }
                }
            }
        }
        self.lowerer.set_cardinality_hints(cardinality_hints);

        let mut plan = self.lowerer.lower_program(&program)?;

        // Phase 4: Optimize (predicate pushdown + cost-aware rewrites)
        //
        // Seed statistics with any known fact cardinalities so cost estimation has
        // at least a baseline for EDB relations.
        let mut mgr = StatsManager::new();
        let mut fact_counts: HashMap<String, u64> = HashMap::new();
        for fact in program.facts() {
            *fact_counts.entry(fact.head.predicate.clone()).or_insert(0) += 1;
        }

        for (pred, rel_id) in self.lowerer.rel_ids() {
            mgr.register_relation(*rel_id);
            let rows = fact_counts.get(pred).copied().unwrap_or(0);
            if rows > 0 {
                mgr.update_cardinality(*rel_id, rows);
                if let Some(schema) = self.lowerer.schemas().get(pred) {
                    mgr.update_byte_size(*rel_id, rows * schema.row_size_bytes() as u64);
                }
            }
        }

        if let Some(snapshot) = stats_snapshot {
            if snapshot.rel_names.is_empty() {
                mgr.merge_snapshot(snapshot);
            } else {
                let rel_name_by_id: HashMap<RelId, &str> = snapshot
                    .rel_names
                    .iter()
                    .map(|(id, name)| (*id, name.as_str()))
                    .collect();

                for rel in &snapshot.relations {
                    let Some(pred) = rel_name_by_id.get(&rel.rel_id) else {
                        continue;
                    };
                    let Some(rel_id) = self.lowerer.rel_ids().get(*pred) else {
                        continue;
                    };

                    let mut remapped = rel.clone();
                    remapped.rel_id = *rel_id;

                    if let Some(schema) = self.lowerer.schemas().get(*pred) {
                        remapped.column_stats.retain(|col| {
                            col.col_idx < schema.arity()
                                && schema.column_type(col.col_idx) == Some(col.dtype)
                        });
                    } else {
                        remapped.column_stats.clear();
                    }

                    mgr.register_relation(*rel_id);
                    if let Some(stats) = mgr.get_relation_stats_mut(*rel_id) {
                        *stats = remapped;
                    }
                }

                for js in &snapshot.join_selectivities {
                    if js.left_keys.len() != js.right_keys.len() {
                        continue;
                    }

                    let Some(left_pred) = rel_name_by_id.get(&js.left_rel) else {
                        continue;
                    };
                    let Some(right_pred) = rel_name_by_id.get(&js.right_rel) else {
                        continue;
                    };
                    let Some(&left_id) = self.lowerer.rel_ids().get(*left_pred) else {
                        continue;
                    };
                    let Some(&right_id) = self.lowerer.rel_ids().get(*right_pred) else {
                        continue;
                    };

                    let Some(left_schema) = self.lowerer.schemas().get(*left_pred) else {
                        continue;
                    };
                    let Some(right_schema) = self.lowerer.schemas().get(*right_pred) else {
                        continue;
                    };
                    if js.left_keys.iter().any(|&k| k >= left_schema.arity())
                        || js.right_keys.iter().any(|&k| k >= right_schema.arity())
                    {
                        continue;
                    }

                    mgr.set_join_selectivity(
                        left_id,
                        right_id,
                        js.left_keys.clone(),
                        js.right_keys.clone(),
                        js.selectivity,
                    );
                }
            }
        }

        // Build schemas by RelId for the optimizer
        let schemas_by_rel_id: HashMap<RelId, Schema> = self
            .lowerer
            .rel_ids()
            .iter()
            .filter_map(|(pred, rel_id)| {
                self.lowerer
                    .schemas()
                    .get(pred)
                    .map(|schema| (*rel_id, schema.clone()))
            })
            .collect();

        let stats_arc = Arc::new(mgr);

        crate::optimizer::helper_split_pass::run(
            &mut plan,
            &schemas_by_rel_id,
            &stats_arc,
            |schema| self.lowerer.create_helper_relation(schema),
        );

        let schemas_by_rel_id: HashMap<RelId, Schema> = self
            .lowerer
            .rel_ids()
            .iter()
            .filter_map(|(pred, rel_id)| {
                self.lowerer
                    .schemas()
                    .get(pred)
                    .map(|schema| (*rel_id, schema.clone()))
            })
            .collect();

        let mut optimizer = Optimizer::new(Arc::clone(&stats_arc));
        optimizer.set_schemas(schemas_by_rel_id);
        for rules in &mut plan.rules_by_scc {
            for rule in rules {
                let body = std::mem::replace(&mut rule.body, xlog_ir::RirNode::Unit);
                rule.body = optimizer.optimize(body);
            }
        }

        // Selectivity-aware reordering pass. Runs BETWEEN the optimizer loop
        // and promote_multiway.
        // Locked compile-pipeline ordering:
        //   lower → helper_split_pass → optimizer → selectivity_pass → promote_multiway
        //
        // Takes `rel_ids` so per-body Scans can be resolved against
        // `StatsManager`. Behavior on empty stats / unseeded relations is
        // no-op (safety floor).
        crate::optimizer::selectivity_pass::run(&mut plan, &stats_arc, self.lowerer.rel_ids());

        // Promote eligible triangle subtrees to RirNode::MultiWayJoin. Runs
        // *after* the optimizer so the optimizer never has to learn the new
        // variant. Fallback identity preserves binary-join semantics on
        // dispatch decline.
        //
        // Pass the lowerer's predicate→RelId map so the promoter can gate
        // recursive-SCC bodies on the count of in-SCC Scans (≤ 1 = promote,
        // ≥ 2 = skip).
        //
        // Also pass `&stats_arc` and the caller-provided `&CompilerConfig`.
        // With `CompilerConfig::default()` (`Disabled`), the promoter never
        // sets `var_order` and default dispatch is bit-identical.
        crate::promote::promote_multiway(&mut plan, self.lowerer.rel_ids(), &stats_arc, config);

        let schemas_by_rel_id: HashMap<RelId, Schema> = self
            .lowerer
            .rel_ids()
            .iter()
            .filter_map(|(pred, rel_id)| {
                self.lowerer
                    .schemas()
                    .get(pred)
                    .map(|schema| (*rel_id, schema.clone()))
            })
            .collect();

        crate::optimizer::helper_split_pass::run_kclique_specs(
            &mut plan,
            &schemas_by_rel_id,
            |schema| self.lowerer.create_helper_relation(schema),
        );

        plan.generated_query_rules = generated_query_heads
            .iter()
            .enumerate()
            .map(|(query_index, head)| {
                let positions = plan
                    .rules_by_scc
                    .iter()
                    .enumerate()
                    .flat_map(|(scc_index, rules)| {
                        rules
                            .iter()
                            .enumerate()
                            .filter_map(move |(rule_index, rule)| {
                                (rule.head == *head).then_some((scc_index, rule_index))
                            })
                    })
                    .collect::<Vec<_>>();
                let [(scc_index, rule_index)] = positions.as_slice() else {
                    return Err(XlogError::Compilation(format!(
                        "generated query head {head} must have exactly one compiled rule, found {}",
                        positions.len()
                    )));
                };
                Ok(GeneratedQueryRuleProvenance {
                    query_index,
                    scc_index: *scc_index,
                    rule_index: *rule_index,
                })
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(plan)
    }

    /// Reset the compiler state for a fresh compilation.
    ///
    /// This creates a new lowerer, clearing any cached schemas or relation IDs
    /// from previous compilations.
    pub fn reset(&mut self) {
        self.lowerer = Lowerer::new();
    }

    /// Get the mapping from predicate names to relation IDs after compilation.
    ///
    /// This mapping is needed to register relations in the executor with
    /// the correct RelIds.
    pub fn rel_ids(&self) -> &HashMap<String, RelId> {
        self.lowerer.rel_ids()
    }

    /// Get the inferred schemas for predicates after compilation.
    ///
    /// These schemas are needed to create GPU buffers with correct column types.
    pub fn schemas(&self) -> &HashMap<String, Schema> {
        self.lowerer.schemas()
    }
}

fn generated_query_heads_for_program(program: &Program) -> Result<Vec<String>> {
    (0..program.queries.len())
        .map(|query_index| {
            let generated_head = format!("__xlog_query_{query_index}");
            let authored_collision = program
                .predicates
                .iter()
                .any(|declaration| declaration.name == generated_head)
                || program
                    .rules
                    .iter()
                    .any(|rule| rule.head.predicate == generated_head)
                || program
                    .learnable_rules
                    .iter()
                    .any(|rule| rule.head.predicate == generated_head);
            if authored_collision {
                return Err(XlogError::Compilation(format!(
                    "authored relation {generated_head} collides with generated query head"
                )));
            }
            Ok(generated_head)
        })
        .collect()
}

/// The source-level frontend passes shared by compilation and validation:
/// desugar queries/constraints → meta normalization → list normalization →
/// magic sets → negation-safety check. All by value, no AST clones.
fn run_frontend_passes(program: Program) -> Result<Program> {
    let program = desugar_queries_and_constraints(program)?;
    let program = normalize_meta_builtins_owned(program)?;
    let program = normalize_list_builtins_owned(program)?;
    let program = rewrite_magic_sets_owned(program)?.program;
    validate_negation_safety(&program)?;
    Ok(program)
}

fn desugar_queries_and_constraints(program: Program) -> Result<Program> {
    let mut out = program;

    // Constraints: `:- body.` becomes `__xlog_constraint_i(1) :- body.`
    for constraint in &out.constraints {
        let authored_index = constraint.require_authored_index()?;
        let pred = format!("__xlog_constraint_{authored_index}");
        out.rules.push(AstRule {
            head: crate::ast::Atom {
                predicate: pred,
                terms: vec![Term::Integer(1)],
            },
            body: constraint.body.clone(),
        });
    }

    // Queries: `?- atom.` becomes `__xlog_query_i(Vars...) :- atom.`
    for (i, Query { atom }) in out.queries.iter().enumerate() {
        let pred = format!("__xlog_query_{}", i);

        let mut head_terms: Vec<Term> = Vec::new();
        let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();

        for term in &atom.terms {
            for name in term.variables() {
                if seen.insert(name) {
                    head_terms.push(Term::Variable(name.to_string()));
                }
            }
        }

        if head_terms.is_empty() {
            head_terms.push(Term::Integer(1));
        }

        out.rules.push(AstRule {
            head: crate::ast::Atom {
                predicate: pred,
                terms: head_terms,
            },
            body: vec![BodyLiteral::Positive(atom.clone())],
        });
    }

    Ok(out)
}

fn validate_negation_safety(program: &Program) -> Result<()> {
    for rule in &program.rules {
        validate_body_naf_safety(&rule.body, &format!("rule {}", rule.head.predicate))?;
    }
    for constraint in &program.constraints {
        validate_body_naf_safety(
            &constraint.body,
            &format!("constraint {}", constraint.require_authored_index()?),
        )?;
    }
    for (idx, learnable) in program.learnable_rules.iter().enumerate() {
        validate_body_naf_safety(&learnable.body, &format!("learnable rule {}", idx))?;
    }
    Ok(())
}

fn validate_body_naf_safety(body: &[BodyLiteral], context: &str) -> Result<()> {
    let mut bound: HashSet<String> = HashSet::new();
    for lit in body {
        match lit {
            BodyLiteral::Positive(atom) => {
                for name in atom.variables() {
                    bound.insert(name.to_string());
                }
            }
            BodyLiteral::Negated(atom) => {
                for name in atom.variables() {
                    if !bound.contains(name) {
                        return Err(naf_error(format!(
                            "unbound variable {} in negated atom {}/{} in {}; bind it before not with a positive atom or deterministic is expression, or use '_' for existential positions",
                            name,
                            atom.predicate,
                            atom.arity(),
                            context
                        )));
                    }
                }
            }
            BodyLiteral::IsExpr(is_expr) => {
                bound.insert(is_expr.target.clone());
            }
            BodyLiteral::Epistemic(_) => {}
            BodyLiteral::Comparison(_) | BodyLiteral::Univ(_) => {}
        }
    }
    Ok(())
}

fn map_stratification_to_naf_error(err: XlogError) -> XlogError {
    match err {
        XlogError::StratificationCycle(cycle) => naf_error(format!(
            "deterministic not atom must be stratified; cycle through negation or aggregation: {}",
            cycle.join(" -> ")
        )),
        other => other,
    }
}

fn naf_error(message: impl Into<String>) -> XlogError {
    XlogError::Compilation(format!("negation safety error: {}", message.into()))
}

/// Convenience function to compile source in one call.
///
/// This creates a short-lived compiler and compiles the source.
/// For multiple compilations, prefer creating a `Compiler` instance directly.
/// Function definitions and calls must already be expanded; execution-facing
/// `LogicProgram` APIs own source normalization and import resolution.
///
/// # Example
///
/// ```ignore
/// use xlog_logic::compile::compile;
///
/// let plan = compile("edge(1, 2). reach(X, Y) :- edge(X, Y).")?;
/// ```
pub fn compile(source: &str) -> Result<ExecutionPlan> {
    let mut compiler = Compiler::new();
    compiler.compile(source)
}

/// Load modules for an entry source file.
///
/// This function:
/// 1. Loads the entry module from the exact supplied path
/// 2. Loads all direct and transitive `.xlog` module dependencies
///
/// # Arguments
///
/// * `entry_file` - Path to the entry source file
/// * `search_paths` - Additional directories to search for modules
///
/// # Returns
///
/// The loaded module resolver with all dependencies resolved, or an error
/// if module resolution fails.
pub fn load_modules(
    entry_file: &Path,
    search_paths: Vec<PathBuf>,
) -> std::result::Result<ModuleResolver, ModuleError> {
    let mut resolver = ModuleResolver::new(search_paths);
    resolver.load_entry_file(entry_file)?;

    Ok(resolver)
}

#[cfg(test)]
mod tests {
    use super::*;
    use xlog_core::ScalarType;
    use xlog_ir::RirNode;
    use xlog_stats::ColumnStats;
    use xlog_stats::RelationStats;
    use xlog_stats::StatsManager;

    #[test]
    fn test_compiler_new() {
        let compiler = Compiler::new();
        // Just verify it can be created
        drop(compiler);
    }

    #[test]
    fn test_compile_fact() {
        let mut compiler = Compiler::new();
        let result = compiler.compile("edge(1, 2).");
        assert!(result.is_ok(), "Failed to compile fact: {:?}", result.err());
    }

    #[test]
    fn test_compile_simple_rule() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            edge(1, 2).
            reach(X, Y) :- edge(X, Y).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile simple rule: {:?}",
            result.err()
        );

        let plan = result.unwrap();
        assert!(!plan.sccs.is_empty(), "Expected at least one SCC");
    }

    #[test]
    fn compiler_records_each_program_query_at_its_compiled_rule_position() {
        let plan = Compiler::new()
            .compile(
                r#"
                    source(a).
                    ?- source(X).
                    ?- source(a).
                "#,
            )
            .expect("queries should compile");

        assert_eq!(plan.generated_query_rules.len(), 2);
        for (query_index, provenance) in plan.generated_query_rules.iter().enumerate() {
            assert_eq!(provenance.query_index, query_index);
            let rule = &plan.rules_by_scc[provenance.scc_index][provenance.rule_index];
            assert_eq!(rule.head, format!("__xlog_query_{query_index}"));
        }
    }

    #[test]
    fn compiler_does_not_infer_query_provenance_from_an_authored_prefix() {
        let mut program = parse_program(
            r#"
                source(a).
                authored(X) :- source(X).
            "#,
        )
        .expect("authored program should parse");
        program
            .rules
            .iter_mut()
            .find(|rule| rule.head.predicate == "authored")
            .expect("authored rule should exist")
            .head
            .predicate = "__xlog_query_authored".to_string();

        let plan = Compiler::new()
            .compile_program(&program)
            .expect("authored prefix should remain an ordinary rule");
        assert!(plan.generated_query_rules.is_empty());
    }

    #[test]
    fn compiler_rejects_an_authored_exact_generated_query_head_collision() {
        let mut program = parse_program(
            r#"
                source(a).
                authored(X) :- source(X).
                ?- source(X).
            "#,
        )
        .expect("collision witness should parse");
        program
            .rules
            .iter_mut()
            .find(|rule| rule.head.predicate == "authored")
            .expect("authored rule should exist")
            .head
            .predicate = "__xlog_query_0".to_string();

        let error = Compiler::new()
            .compile_program(&program)
            .expect_err("authored exact query-head collisions must fail compilation");
        assert!(
            error
                .to_string()
                .contains("authored relation __xlog_query_0 collides with generated query head"),
            "unexpected collision error: {error}"
        );
    }

    #[test]
    fn compiler_reports_bound_is_target_before_arithmetic_type_mismatch() {
        let error = Compiler::new()
            .compile(
                r#"
                val(10).
                bad(Z) :- val(Z), Z is Z + 1.
                ?- bad(Z).
                "#,
            )
            .expect_err("an already-bound is-expression target must be rejected");

        assert_eq!(
            error.to_string(),
            "Compilation error: Variable Z already bound; 'is' requires fresh variable"
        );
    }

    #[test]
    fn compiler_rejects_is_target_bound_by_an_earlier_is_expression() {
        let error = Compiler::new()
            .compile(
                r#"
                val(10).
                bad(Z) :- val(X), Z is X + X, Z is X + X.
                ?- bad(Z).
                "#,
            )
            .expect_err("a prior is-expression target must stay bound");

        assert_eq!(
            error.to_string(),
            "Compilation error: Variable Z already bound; 'is' requires fresh variable"
        );
    }

    #[test]
    fn compiler_accepts_fresh_chained_is_targets() {
        Compiler::new()
            .compile(
                r#"
                val(10).
                good(Z) :- val(X), Y is X + X, Z is Y + Y.
                ?- good(Z).
                "#,
            )
            .expect("distinct source-ordered is-expression targets must remain valid");
    }

    #[test]
    fn compiler_reports_unbound_is_input_before_a_later_duplicate_target() {
        let error = Compiler::new()
            .compile(
                r#"
                seed(1).
                bad(Z) :- seed(X), Z is U + 1, Z is 1 + 1.
                ?- bad(Z).
                "#,
            )
            .expect_err("an unreachable later binding must not preempt the first is error");

        assert_eq!(
            error.to_string(),
            "Compilation error: Variable U used in arithmetic but not bound"
        );
    }

    #[test]
    fn test_compile_transitive_closure() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            edge(1, 2).
            edge(2, 3).
            edge(3, 4).
            reach(X, Y) :- edge(X, Y).
            reach(X, Z) :- reach(X, Y), edge(Y, Z).
        "#,
        );
        assert!(result.is_ok(), "Failed to compile TC: {:?}", result.err());

        let plan = result.unwrap();
        // Should have SCCs for edge and reach
        assert!(!plan.sccs.is_empty());
    }

    #[test]
    fn test_compile_with_negation() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            node(1).
            node(2).
            node(3).
            edge(1, 2).
            isolated(X) :- node(X), not edge(X, _).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile with negation: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_compile_with_comparison() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            value(1).
            value(5).
            value(10).
            value(15).
            small(X) :- value(X), X < 10.
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile with comparison: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_schema_infers_from_rule_body_types() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            path(X, Y) :- reach(X, Y).
            edge(1, 2).
            edge(2, 3).
            reach(X, Y) :- edge(X, Y).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile rule for schema inference: {:?}",
            result.err()
        );

        let schema = compiler
            .schemas()
            .get("reach")
            .expect("missing reach schema");
        assert_eq!(
            schema.column_type(0),
            Some(ScalarType::U32),
            "reach column 0 should match edge column type"
        );
        assert_eq!(
            schema.column_type(1),
            Some(ScalarType::U32),
            "reach column 1 should match edge column type"
        );

        let schema = compiler.schemas().get("path").expect("missing path schema");
        assert_eq!(
            schema.column_type(0),
            Some(ScalarType::U32),
            "path column 0 should inherit the transitive body type"
        );
        assert_eq!(
            schema.column_type(1),
            Some(ScalarType::U32),
            "path column 1 should inherit the transitive body type"
        );
    }

    #[test]
    fn test_schema_propagates_known_columns_when_a_sibling_column_is_unresolved() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            path(X) :- intermediate(X, Z).
            intermediate(X, Z) :- source(X), unresolved(Z).
            unresolved(Z) :- unresolved(Z).
            source(value).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile partial rule-head schema inference: {:?}",
            result.err()
        );

        let schema = compiler.schemas().get("path").expect("missing path schema");
        assert_eq!(
            schema.column_type(0),
            Some(ScalarType::Symbol),
            "path should inherit the known intermediate column independently of its sibling"
        );
    }

    #[test]
    fn test_schema_rejects_conflicting_rule_head_evidence_in_either_order() {
        let symbol_first = r#"
            pred symbols(symbol).
            pred integers(i64).
            symbols(value).
            integers(5000000000).
            mixed(X) :- symbols(X).
            mixed(X) :- integers(X).
        "#;
        let integer_first = r#"
            pred symbols(symbol).
            pred integers(i64).
            symbols(value).
            integers(5000000000).
            mixed(X) :- integers(X).
            mixed(X) :- symbols(X).
        "#;

        for source in [symbol_first, integer_first] {
            let error = Compiler::new()
                .compile(source)
                .expect_err("conflicting inferred rule-head types must be rejected");
            let message = error.to_string();
            assert!(message.contains("mixed"), "{message}");
            assert!(message.contains("column 1"), "{message}");
        }
    }

    #[test]
    fn test_schema_rejects_conflicting_body_evidence_in_either_order() {
        let first_body_order = r#"
            source_number(1).
            source_symbol(value).
            mixed(X) :- source_number(X), source_symbol(X).
            mixed(1) :- source_number(1).
        "#;
        let reversed_body_order = r#"
            source_number(1).
            source_symbol(value).
            mixed(X) :- source_symbol(X), source_number(X).
            mixed(1) :- source_number(1).
        "#;

        for source in [first_body_order, reversed_body_order] {
            let error = Compiler::new()
                .compile(source)
                .expect_err("conflicting body-column types must be rejected");
            let message = error.to_string();
            assert!(message.contains("mixed"), "{message}");
            assert!(message.contains("variable X"), "{message}");
        }
    }

    #[test]
    fn test_schema_infers_arithmetic_binding_result_type() {
        let mut compiler = Compiler::new();
        compiler
            .compile(
                r#"
                computed(X) :- X is cast(1, u64).
            "#,
            )
            .expect("an arithmetic binding should determine its head-column type");

        let schema = compiler
            .schemas()
            .get("computed")
            .expect("missing computed schema");
        assert_eq!(schema.column_type(0), Some(ScalarType::U64));
    }

    #[test]
    fn test_schema_rejects_arithmetic_and_atom_evidence_in_either_rule_order() {
        let arithmetic_first = r#"
            pred integers(i64).
            integers(5000000000).
            mixed(X) :- X is cast(1, u64).
            mixed(X) :- integers(X).
        "#;
        let atom_first = r#"
            pred integers(i64).
            integers(5000000000).
            mixed(X) :- integers(X).
            mixed(X) :- X is cast(1, u64).
        "#;

        for source in [arithmetic_first, atom_first] {
            let error = Compiler::new()
                .compile(source)
                .expect_err("conflicting arithmetic and atom types must be rejected");
            let message = error.to_string();
            assert!(message.contains("mixed"), "{message}");
            assert!(message.contains("column 1"), "{message}");
        }
    }

    #[test]
    fn test_schema_rejects_incompatible_constant_heads_before_runtime() {
        let sources = [
            r#"
                pred p(u32).
                seed().
                p(value) :- seed().
            "#,
            r#"
                p(1).
                p(value).
            "#,
            r#"
                p(value).
                p(1).
            "#,
            r#"
                p(1).
                p(5000000000).
            "#,
            r#"
                p(5000000000).
                p(1).
            "#,
        ];

        for source in sources {
            let error = Compiler::new()
                .compile(source)
                .expect_err("incompatible constant head types must fail during compilation");
            let message = error.to_string();
            assert!(message.contains("p"), "{message}");
            assert!(message.contains("head term at position 0"), "{message}");
        }
    }

    #[test]
    fn test_declared_scalar_schemas_accept_supported_literal_conversions() {
        Compiler::new()
            .compile(
                r#"
                pred signed(i64).
                pred wide(u64).
                pred real(f64).
                pred flag(bool).
                seed().
                signed(1).
                wide(1).
                real(1).
                real(2) :- seed().
                flag(true).
                flag(false) :- seed().
            "#,
            )
            .expect("representable literals should use declared scalar storage types");
    }

    #[test]
    fn test_compile_unstratifiable_fails() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            p :- not q.
            q :- not p.
        "#,
        );
        assert!(result.is_err(), "Should fail with stratification cycle");
    }

    #[test]
    fn test_compile_syntax_error_fails() {
        let mut compiler = Compiler::new();
        let result = compiler.compile("edge(1, 2"); // Missing closing paren and period
        assert!(result.is_err(), "Should fail with syntax error");
    }

    #[test]
    fn test_compile_convenience_function() {
        let result = compile("edge(1, 2).");
        assert!(
            result.is_ok(),
            "Convenience compile failed: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_compiler_reset() {
        let mut compiler = Compiler::new();

        // First compilation
        let result1 = compiler.compile("edge(1, 2).");
        assert!(result1.is_ok());

        // Reset and compile again
        compiler.reset();
        let result2 = compiler.compile("node(1). node(2).");
        assert!(result2.is_ok());
    }

    #[test]
    fn test_compile_with_pred_decl() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            pred edge(u32, u32).
            edge(1, 2).
            edge(2, 3).
            reach(X, Y) :- edge(X, Y).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile with pred decl: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_compile_multi_stratum() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            // Base facts
            edge(1, 2).
            edge(2, 3).
            edge(3, 1).

            // Stratum 0: edge (base)
            // Stratum 1: reach (depends on edge, recursive)
            reach(X, Y) :- edge(X, Y).
            reach(X, Z) :- reach(X, Y), edge(Y, Z).

            // Stratum 2: non_reach (negates reach)
            all_pairs(X, Y) :- edge(X, Z), edge(Y, W).
            non_reach(X, Y) :- all_pairs(X, Y), not reach(X, Y).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile multi-stratum: {:?}",
            result.err()
        );

        let plan = result.unwrap();
        // Should have multiple strata
        assert!(!plan.strata.is_empty(), "Expected multiple strata");
    }

    #[test]
    fn test_compile_aggregation() {
        let mut compiler = Compiler::new();
        let result = compiler.compile(
            r#"
            edge(1, 2).
            edge(1, 3).
            edge(2, 3).
            out_degree(X, count(Y)) :- edge(X, Y).
        "#,
        );
        assert!(
            result.is_ok(),
            "Failed to compile with aggregation: {:?}",
            result.err()
        );

        let plan = result.unwrap();
        let out_degree_rules: Vec<_> = plan
            .rules_by_scc
            .iter()
            .flatten()
            .filter(|r| r.head == "out_degree")
            .collect();
        assert_eq!(out_degree_rules.len(), 1, "Expected one out_degree rule");

        // Aggregation lowering should produce a GroupBy node (wrapped in a Project to match head order).
        let body = &out_degree_rules[0].body;
        match body {
            RirNode::Project { input, .. } => {
                assert!(
                    matches!(input.as_ref(), RirNode::GroupBy { .. }),
                    "Expected Project(GroupBy(..)), got {:?}",
                    input
                );
            }
            other => panic!("Expected Project(GroupBy(..)), got {:?}", other),
        }
    }

    #[test]
    fn test_aggregate_schemas_match_runtime_result_types() {
        let mut compiler = Compiler::new();
        compiler
            .compile(
                r#"
                pred edge(u32, u64).
                pred score(u32, f64).
                pred counted(u32, u64).
                pred summed(u32, u64).
                pred minimum(u32, u64).
                pred maximum(u32, u64).
                pred combined(u32, f64).
                edge(1, 5000000000).
                score(1, 2.5).
                counted(X, count(Y)) :- edge(X, Y).
                summed(X, sum(Y)) :- edge(X, Y).
                minimum(X, min(Y)) :- edge(X, Y).
                maximum(X, max(Y)) :- edge(X, Y).
                combined(X, logsumexp(Y)) :- score(X, Y).
            "#,
            )
            .expect("aggregate declarations should match provider result schemas");

        for predicate in ["counted", "summed", "minimum", "maximum"] {
            assert_eq!(
                compiler
                    .schemas()
                    .get(predicate)
                    .expect("missing aggregate schema")
                    .column_type(1),
                Some(ScalarType::U64),
                "unexpected result type for {predicate}"
            );
        }
        assert_eq!(
            compiler
                .schemas()
                .get("combined")
                .expect("missing log-sum-exp schema")
                .column_type(1),
            Some(ScalarType::F64)
        );
    }

    #[test]
    fn test_undeclared_count_schema_is_u64() {
        let mut compiler = Compiler::new();
        compiler
            .compile(
                r#"
                edge(1, 5000000000).
                degree(X, count(Y)) :- edge(X, Y).
            "#,
            )
            .expect("undeclared count result should compile with its runtime type");
        assert_eq!(
            compiler
                .schemas()
                .get("degree")
                .expect("missing inferred count schema")
                .column_type(1),
            Some(ScalarType::U64)
        );
    }

    #[test]
    fn test_unsupported_aggregate_inputs_are_rejected_before_runtime() {
        let fixtures = [
            r#"
                pred source(u32, i32).
                pred result(u32, u64).
                source(1, 2).
                result(X, sum(Y)) :- source(X, Y).
            "#,
            r#"
                pred source(u32, u32).
                pred result(u32, f64).
                source(1, 2).
                result(X, logsumexp(Y)) :- source(X, Y).
            "#,
        ];

        for source in fixtures {
            let error = Compiler::new()
                .compile(source)
                .expect_err("unsupported aggregate input must fail during compilation");
            let message = error.to_string();
            assert!(message.contains("Unsupported aggregate input"), "{message}");
            assert!(message.contains("execution provider requires"), "{message}");
        }
    }

    #[test]
    fn test_compile_with_stats_snapshot() {
        let mut compiler = Compiler::new();
        let source = r#"
            edge(1, 2).
            edge(2, 3).
            reach(X, Y) :- edge(X, Y).
        "#;

        let _ = compiler.compile(source).expect("Initial compile failed");
        let edge_id = *compiler.rel_ids().get("edge").expect("edge rel_id missing");

        let mut mgr = StatsManager::new();
        mgr.register_relation(edge_id);
        mgr.update_cardinality(edge_id, 42);
        let snapshot = mgr.snapshot();

        let plan = compiler
            .compile_with_stats_snapshot(source, Some(&snapshot))
            .expect("Compile with snapshot failed");
        assert!(!plan.sccs.is_empty());
    }

    #[test]
    fn test_compile_with_named_stats_snapshot_reorders_joins() {
        let mut compiler = Compiler::new();
        let source = r#"
            foo(1).
            edge(1).
            out(X) :- edge(X), foo(X).
        "#;

        // Snapshot uses different RelIds than the compiler will assign for this program.
        // Map: RelId(0) -> edge (small), RelId(1) -> foo (big)
        let mut edge_stats = RelationStats::new(RelId(0));
        edge_stats.update_cardinality(10);
        let mut foo_stats = RelationStats::new(RelId(1));
        foo_stats.update_cardinality(10_000);

        let snapshot = StatsSnapshot {
            relations: vec![edge_stats, foo_stats],
            join_selectivities: Vec::new(),
            rel_names: vec![
                (RelId(0), "edge".to_string()),
                (RelId(1), "foo".to_string()),
            ],
        };

        let plan = compiler
            .compile_with_stats_snapshot(source, Some(&snapshot))
            .expect("Compile with named snapshot failed");

        let foo_id = *compiler.rel_ids().get("foo").expect("foo rel_id missing");
        let edge_id = *compiler.rel_ids().get("edge").expect("edge rel_id missing");

        let out_rule = plan
            .rules_by_scc
            .iter()
            .flatten()
            .find(|r| r.head == "out")
            .expect("out rule missing");

        // Peel projections to reach the join.
        let mut node = &out_rule.body;
        while let RirNode::Project { input, .. } = node {
            node = input;
        }

        match node {
            RirNode::ChainJoin {
                left,
                right,
                fallback,
                ..
            } => {
                // ChainJoin promotion wraps eligible two-atom joins after
                // stats-aware ordering. The chain node and its captured
                // fallback must agree on the build-side choice.
                assert!(matches!(**left, RirNode::Scan { rel } if rel == foo_id));
                assert!(matches!(**right, RirNode::Scan { rel } if rel == edge_id));

                let mut fallback_node = fallback.as_ref();
                while let RirNode::Project { input, .. } = fallback_node {
                    fallback_node = input;
                }
                match fallback_node {
                    RirNode::Join { left, right, .. } => {
                        assert!(matches!(**left, RirNode::Scan { rel } if rel == foo_id));
                        assert!(matches!(**right, RirNode::Scan { rel } if rel == edge_id));
                    }
                    other => panic!("Expected ChainJoin fallback Join node, got {:?}", other),
                }
            }
            RirNode::Join { left, right, .. } => {
                // Prefer building on the smaller relation (right/build side).
                assert!(matches!(**left, RirNode::Scan { rel } if rel == foo_id));
                assert!(matches!(**right, RirNode::Scan { rel } if rel == edge_id));
            }
            other => panic!("Expected Join node, got {:?}", other),
        }
    }

    fn helper_split_source() -> &'static str {
        r#"
            ab(0, 0). bc(0, 0). cd(0, 0). de(0, 0). ef(0, 0). af(0, 0).
            out(A, B, C, D, F) :-
                ab(A, B),
                bc(B, C),
                cd(C, D),
                de(D, E),
                ef(E, F),
                af(A, F).
        "#
    }

    fn helper_split_snapshot(distinct_d: u64) -> StatsSnapshot {
        let mut snapshot_relations = Vec::new();
        for (idx, name) in ["ab", "bc", "cd", "de", "ef", "af"].iter().enumerate() {
            let mut rel_stats = RelationStats::new(RelId(idx as u32));
            rel_stats.update_cardinality(8192);
            if *name == "de" {
                let mut d_col = ColumnStats::new(0, ScalarType::U32);
                d_col.update_distinct(distinct_d);
                rel_stats.add_column(d_col);
            }
            snapshot_relations.push(rel_stats);
        }
        StatsSnapshot {
            relations: snapshot_relations,
            join_selectivities: Vec::new(),
            rel_names: ["ab", "bc", "cd", "de", "ef", "af"]
                .iter()
                .enumerate()
                .map(|(idx, name)| (RelId(idx as u32), (*name).to_string()))
                .collect(),
        }
    }

    #[test]
    fn test_compile_with_named_stats_snapshot_creates_helper_relation() {
        let mut compiler = Compiler::new();
        let snapshot = helper_split_snapshot(1);
        let plan = compiler
            .compile_with_stats_snapshot(helper_split_source(), Some(&snapshot))
            .expect("compile with helper stats");
        let helper = compiler
            .rel_ids()
            .iter()
            .find_map(|(name, rel)| {
                name.starts_with("__kclique_helper_")
                    .then_some((name.clone(), *rel))
            })
            .expect("helper relation allocated");

        let helper_rule_count = plan
            .rules_by_scc
            .iter()
            .flatten()
            .filter(|rule| rule.head == helper.0)
            .count();
        assert_eq!(helper_rule_count, 1);

        let helper_rule = plan
            .rules_by_scc
            .iter()
            .flatten()
            .find(|rule| rule.head == helper.0)
            .expect("helper rule");
        assert!(
            matches!(helper_rule.body, RirNode::ChainJoin { .. }),
            "helper split output should be eligible for ChainJoin promotion"
        );

        let out_rule = plan
            .rules_by_scc
            .iter()
            .flatten()
            .find(|rule| rule.head == "out")
            .expect("out rule");
        assert!(contains_scan(&out_rule.body, helper.1));
    }

    #[test]
    fn test_compile_with_flat_named_stats_keeps_original_rule() {
        let mut compiler = Compiler::new();
        let snapshot = helper_split_snapshot(8192);
        let plan = compiler
            .compile_with_stats_snapshot(helper_split_source(), Some(&snapshot))
            .expect("compile with flat stats");

        assert!(!compiler
            .rel_ids()
            .keys()
            .any(|name| name.starts_with("__kclique_helper_")));
        let out_rules = plan
            .rules_by_scc
            .iter()
            .flatten()
            .filter(|rule| rule.head == "out")
            .count();
        assert_eq!(out_rules, 1);
    }

    fn contains_scan(node: &RirNode, rel: RelId) -> bool {
        match node {
            RirNode::Scan { rel: scan_rel } => *scan_rel == rel,
            RirNode::Join { left, right, .. } | RirNode::ChainJoin { left, right, .. } => {
                contains_scan(left, rel) || contains_scan(right, rel)
            }
            RirNode::Project { input, .. }
            | RirNode::Filter { input, .. }
            | RirNode::Distinct { input, .. }
            | RirNode::GroupBy { input, .. } => contains_scan(input, rel),
            RirNode::Union { inputs } => inputs.iter().any(|input| contains_scan(input, rel)),
            RirNode::Diff { left, right } => contains_scan(left, rel) || contains_scan(right, rel),
            RirNode::Fixpoint {
                base, recursive, ..
            } => contains_scan(base, rel) || contains_scan(recursive, rel),
            RirNode::MultiWayJoin { inputs, .. } => {
                inputs.iter().any(|input| contains_scan(input, rel))
            }
            RirNode::TensorMaskedJoin { rel_index, .. } => {
                rel_index.iter().any(|(input_rel, _)| *input_rel == rel)
            }
            RirNode::Unit => false,
        }
    }
}