polydat-core 0.6.0

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

//! `CompiledComprehension` — the entry point for the three
//! consumption surfaces.
//!
//! Holds an `Arc<Program>` (immutable IR per comprehension_forms.md
//! §9.1). Each factory method on this handle (`coordinate_stream`,
//! `scoped_kernel_stream`, `scope_once`) returns a fresh
//! streamer with its own dispense state but shares the
//! `Arc<Program>`, so siblings never recompile (§9.5.2's
//! independence contract).

use std::sync::Arc;

use crate::iteration::comprehension::ast::Comprehension;
use crate::iteration::comprehension::eval_source::{EvalClass, SourceEval};
use crate::iteration::comprehension::flatten::flatten_static_sources;
use crate::iteration::comprehension::ir::{Program, compile as compile_to_ir};
use crate::iteration::comprehension::optimize::optimize;
use crate::iteration::comprehension::predicate::recognizers::extract_coord_refs;
use crate::iteration::comprehension::validate::{
    Mode, Surface, ValidationError, ValidationReport, ValidationWarning, unresolved_names, validate,
};

use crate::kernel::interp::NoScope;

use super::coord_stream::CoordinateStream;
use super::instance::{KernelScope, ScopedKernelInstance};
use super::scope_once::scope_once_with;
use super::scoped_stream::ScopedKernelStream;

/// A comprehension that has been compiled to immutable IR
/// and is ready to dispense. The single source of truth for
/// the underlying program across the consumption surfaces.
///
/// Construction is via [`from_ast`](Self::from_ast) (compiles once) or
/// [`from_program`](Self::from_program) (when the IR was compiled elsewhere).
/// Cloning a `CompiledComprehension` is cheap — just an
/// `Arc::clone` on the program.
#[derive(Debug, Clone)]
pub struct CompiledComprehension {
    program: Arc<Program>,
}

impl CompiledComprehension {
    /// Compile an AST: validation (§5, V1–V9) first, then the §10
    /// optimizer, then the AST → IR pass, once. Validation and
    /// optimization are stages of this compile
    /// (comprehension_forms.md §5, §9.4, §10.6): a tree that
    /// violates a V-axiom is refused here, and the §9.3 resource
    /// bounds hold for the program this returns.
    pub fn from_ast(ast: &Comprehension) -> Result<Self, ValidationError> {
        Self::from_ast_with(ast, Mode::Permissive).map(|(compiled, _)| compiled)
    }

    /// [`from_ast`](Self::from_ast) under a validation mode
    /// (comprehension_forms.md §5.8), with the validator's report: in
    /// `Permissive` mode a degenerate composition is a warning in the
    /// report and the comprehension compiles; in `Strict` mode it is
    /// the error. Context-free sources are flattened first
    /// (§10.7.0): a generator that references no name is evaluated
    /// here, in the empty scope, and validated as the literal of its
    /// values.
    ///
    /// The comprehension is compiled with no enclosing scope, so a name
    /// it reads and does not bind is resolved nowhere (V3): in
    /// `Permissive` mode it is a `ValidationWarning::UnresolvedNames` in
    /// the report and reads None, and in `Strict` mode it is
    /// `ValidationError::V3UnresolvedNames`. [`from_ast_in`](Self::from_ast_in)
    /// compiles one bound in a scope.
    pub fn from_ast_with(
        ast: &Comprehension,
        mode: Mode,
    ) -> Result<(Self, ValidationReport), ValidationError> {
        Self::from_ast_in(ast, mode, &|_| false)
    }

    /// [`from_ast_with`](Self::from_ast_with) for a comprehension bound in
    /// an enclosing scope, which has the names `in_scope` answers true
    /// for, as a producer wire's comprehension is (comprehension_forms.md
    /// §5 V3, §9.5.2).
    ///
    /// These surfaces supply no name. A name the enclosing scope does not
    /// have either is resolved nowhere (V3): outside `Strict` mode it is
    /// a `ValidationWarning::UnresolvedNames` in the report and reads
    /// None, so a source reading one yields nothing and a predicate that
    /// reads one for a tuple keeps it not; in `Strict` mode it is
    /// `ValidationError::V3UnresolvedNames`. A name only that scope has
    /// is one a `for` traversal captures when it opens and a stream cannot
    /// see: a source reading one is `ValidationError::ContextRequired`, as
    /// is a source reading an earlier axis, which only a traversal
    /// evaluates, and a predicate reading one is
    /// `ValidationError::PredicateContextRequired`.
    pub fn from_ast_in(
        ast: &Comprehension,
        mode: Mode,
        in_scope: &dyn Fn(&str) -> bool,
    ) -> Result<(Self, ValidationReport), ValidationError> {
        let ast = flatten_static_sources(ast, &NoScope::new());
        // Resolved nowhere: neither bound nor a name a traversal of it in
        // the enclosing scope would capture.
        let unresolved = unresolved_names(&ast, Surface::Traversal(in_scope));
        if mode == Mode::Strict && !unresolved.is_empty() {
            return Err(ValidationError::V3UnresolvedNames { reads: unresolved });
        }
        if let Some((name, references)) = first_context_required(&ast, in_scope) {
            return Err(ValidationError::ContextRequired { name, references });
        }
        if let Some((predicate, references)) = first_unbound_predicate(&ast, in_scope) {
            return Err(ValidationError::PredicateContextRequired {
                predicate,
                references,
            });
        }
        if let Some((name, message)) = first_failed_static(&ast) {
            return Err(ValidationError::SourceFailed { name, message });
        }
        let mut report = validate(&ast, mode)?;
        if !unresolved.is_empty() {
            report
                .warnings
                .insert(0, ValidationWarning::UnresolvedNames { reads: unresolved });
        }
        Ok((
            Self {
                program: Arc::new(compile_to_ir(&optimize(ast))),
            },
            report,
        ))
    }

    /// Wrap an already-compiled program (tests use this for
    /// hand-built IR).
    pub fn from_program(program: Arc<Program>) -> Self {
        Self { program }
    }

    /// Access the underlying compiled program (immutable per
    /// comprehension_forms.md §9.1).
    pub fn program(&self) -> &Program {
        &self.program
    }

    /// Clone the `Arc<Program>` for sharing with other
    /// handles. Used internally by the streamer factories.
    pub(crate) fn program_arc(&self) -> Arc<Program> {
        Arc::clone(&self.program)
    }

    /// **First-order surface** (comprehension_forms.md §9.5).
    ///
    /// Return a fresh [`CoordinateStream`]. Each call
    /// allocates new per-streamer state; siblings share the
    /// underlying IR but dispense independently per §9.5.2's
    /// independence contract.
    pub fn coordinate_stream(&self) -> CoordinateStream {
        CoordinateStream::new(self.program_arc())
    }

    /// **Second-order surface** (comprehension_forms.md §9.5).
    ///
    /// Return a fresh [`ScopedKernelStream`] wrapping the
    /// supplied parent kernel. Each `advance()` pulls one
    /// coord tuple from the underlying IR and applies
    /// `parent.scope(&coords)` to produce a
    /// [`ScopedKernelInstance`].
    ///
    /// Independence: pulling from this stream does NOT
    /// advance any [`CoordinateStream`] obtained from the
    /// same `CompiledComprehension`.
    pub fn scoped_kernel_stream<K: KernelScope>(&self, parent: K) -> ScopedKernelStream<K> {
        ScopedKernelStream::new(self.program_arc(), parent)
    }

    /// **One-shot surface** (comprehension_forms.md §9.5.3).
    ///
    /// Apply `parent.scope(coords)` directly, without
    /// constructing any streamer. Pure function — no
    /// cursor consulted, no dispense state advanced. Used
    /// for replay, debugging, and point queries where a
    /// specific coord tuple is already known.
    pub fn scope_once<K: KernelScope>(
        &self,
        parent: &K,
        coords: &crate::iteration::comprehension::strategies::Tuple,
    ) -> ScopedKernelInstance<K::Scoped> {
        scope_once_with(parent, coords)
    }
}

/// The first clause of `ast` whose source needs a scope
/// (comprehension_forms.md §10.7.0), with the names it references:
/// the scope-less surfaces refuse such a comprehension by name. A source
/// that reads a name resolved nowhere, bound neither by an earlier axis
/// nor in the enclosing scope (`in_scope`), reads None and yields
/// nothing on every surface, so it needs no scope.
fn first_context_required(
    ast: &Comprehension,
    in_scope: &dyn Fn(&str) -> bool,
) -> Option<(String, Vec<String>)> {
    fn walk(
        c: &Comprehension,
        in_scope: &dyn Fn(&str) -> bool,
        before: &mut Vec<String>,
    ) -> Option<(String, Vec<String>)> {
        match c {
            Comprehension::Clause { name, source } => {
                let references = source.referenced_names();
                let reads_none = references
                    .iter()
                    .any(|n| !before.contains(n) && !in_scope(n));
                (source.eval_class() == EvalClass::ContextRequired && !reads_none)
                    .then(|| (name.clone(), references.into_iter().collect()))
            }
            Comprehension::Cartesian { children } => {
                let depth = before.len();
                let mut found = None;
                for child in children {
                    found = walk(child, in_scope, before);
                    if found.is_some() {
                        break;
                    }
                    before.extend(child.coordinate_names());
                }
                before.truncate(depth);
                found
            }
            Comprehension::Zip { children, .. } | Comprehension::Union { children } => children
                .iter()
                .find_map(|child| walk(child, in_scope, before)),
            Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
                walk(child, in_scope, before)
            }
        }
    }
    walk(ast, in_scope, &mut Vec::new())
}

/// The first filter of `ast` whose predicate names what its tuples do
/// not bind and the enclosing scope has (`in_scope`), with those names:
/// the scope-less surfaces evaluate a predicate in the empty scope, so
/// they refuse it by name. A name the scope does not have either reads
/// None on every surface.
fn first_unbound_predicate(
    ast: &Comprehension,
    in_scope: &dyn Fn(&str) -> bool,
) -> Option<(String, Vec<String>)> {
    match ast {
        Comprehension::Clause { .. } => None,
        Comprehension::Cartesian { children }
        | Comprehension::Zip { children, .. }
        | Comprehension::Union { children } => children
            .iter()
            .find_map(|child| first_unbound_predicate(child, in_scope)),
        Comprehension::Filter { child, predicate } => {
            let bound = child.coordinate_names();
            let unbound: Vec<String> = extract_coord_refs(predicate)
                .into_iter()
                .filter(|name| !bound.contains(name) && in_scope(name))
                .collect();
            if unbound.is_empty() {
                first_unbound_predicate(child, in_scope)
            } else {
                Some((predicate.clone(), unbound))
            }
        }
        Comprehension::Order { child, .. } => first_unbound_predicate(child, in_scope),
    }
}

/// The first context-free generator the flatten could not evaluate,
/// with the evaluator's message. After [`flatten_static_sources`] in
/// the empty scope, a context-free generator that is still a call with
/// no cardinality hint is one whose evaluation failed: an empty result
/// carries a hint of 0 and a non-literal result its count. Nothing
/// later evaluates it on the scope-less surfaces, so it is evaluated
/// once more here for its message, and refused.
fn first_failed_static(ast: &Comprehension) -> Option<(String, String)> {
    use crate::iteration::comprehension::eval_source::EvalContext;
    use crate::iteration::comprehension::source::Source;
    match ast {
        Comprehension::Clause { name, source } => match source {
            Source::Generator {
                cardinality_hint: None,
                ..
            } if source.eval_class() == EvalClass::Static => {
                let scope = NoScope::new();
                let ctx = EvalContext {
                    var_name: name,
                    scope: &scope,
                    prefix: &[],
                };
                source
                    .evaluate(Some(&ctx))
                    .err()
                    .map(|e| (name.clone(), e.to_string()))
            }
            _ => None,
        },
        Comprehension::Cartesian { children }
        | Comprehension::Zip { children, .. }
        | Comprehension::Union { children } => children.iter().find_map(first_failed_static),
        Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
            first_failed_static(child)
        }
    }
}

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

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

    #[test]
    fn from_ast_compiles_once() {
        let ast = clause("k", &[1, 2, 3]);
        let compiled = CompiledComprehension::from_ast(&ast).unwrap();
        assert!(!compiled.program().is_empty());
    }

    /// Optimization is mandatory before IR compilation
    /// (comprehension_forms.md §9.4, §10.6): `from_ast` compiles the
    /// optimized tree, so a program it returns is the program of the
    /// optimized AST, and where a rule fires it is not the program of
    /// the raw one.
    #[test]
    fn from_ast_compiles_the_optimized_tree() {
        // A nested cartesian: R0b (A2) flattens it, so the raw and
        // optimized trees compile to different programs.
        let inner = Comprehension::cartesian(vec![clause("a", &[1, 2]), clause("b", &[3])]);
        let ast = Comprehension::cartesian(vec![inner, clause("c", &[4])]);
        let compiled = CompiledComprehension::from_ast(&ast).unwrap();
        assert_eq!(*compiled.program(), compile_to_ir(&optimize(ast.clone())));
        assert_ne!(*compiled.program(), compile_to_ir(&ast));
    }

    /// Validation is a stage of the compile (comprehension_forms.md
    /// §5): a tree that violates a V-axiom is refused by `from_ast`
    /// with the axiom's error, not compiled.
    #[test]
    fn from_ast_refuses_a_tree_that_violates_a_v_axiom() {
        // V1: a cartesian whose children bind the same name.
        let ast = Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("k", &[3, 4])]);
        let err = CompiledComprehension::from_ast(&ast).unwrap_err();
        assert!(
            matches!(err, ValidationError::V1DuplicateName { ref name, .. } if name == "k"),
            "{err}"
        );
        assert!(err.to_string().starts_with("V1:"), "{err}");
    }

    /// Validation modes (comprehension_forms.md §5.8): a degenerate
    /// composition is a warning in the permissive report and the error
    /// of a strict compile.
    #[test]
    fn from_ast_with_reports_or_refuses_a_degenerate_composition() {
        use crate::iteration::comprehension::strategy::StrategyName;
        use crate::iteration::comprehension::validate::ValidationWarning;
        let ast = Comprehension::order(clause("k", &[1, 2, 3]), StrategyName::Extrema, Some(1));
        let (_, report) = CompiledComprehension::from_ast_with(&ast, Mode::Permissive).unwrap();
        assert!(matches!(
            report.warnings.as_slice(),
            [ValidationWarning::DegenerateGeometric { .. }]
        ));
        let err = CompiledComprehension::from_ast_with(&ast, Mode::Strict).unwrap_err();
        assert!(matches!(err, ValidationError::StrictWarning(_)), "{err}");
        assert!(err.to_string().starts_with("strict mode:"), "{err}");
    }

    #[test]
    fn cloning_compiled_shares_arc() {
        let ast = clause("k", &[1, 2, 3]);
        let a = CompiledComprehension::from_ast(&ast).unwrap();
        let b = a.clone();
        // Same Arc — strong_count goes up.
        let count = Arc::strong_count(&a.program);
        assert!(count >= 2, "expected shared Arc, count = {count}");
        drop(b);
    }

    #[test]
    fn two_coordinate_streams_share_program() {
        let ast = clause("k", &[1, 2, 3]);
        let compiled = CompiledComprehension::from_ast(&ast).unwrap();
        let _s1 = compiled.coordinate_stream();
        let _s2 = compiled.coordinate_stream();
        // Both streams hold an Arc; count is at least 3 (compiled +
        // two streamers, possibly more if internal clones happen).
        let count = Arc::strong_count(&compiled.program);
        assert!(
            count >= 3,
            "expected shared program across streamers, count = {count}"
        );
    }

    /// A context-required source has no coordinate stream
    /// (comprehension_forms.md §9.5.2, §10.7.0): compiled in a scope
    /// that has the names it reads, the compile refuses it by name, with
    /// the names it needs, instead of dispensing nothing. With no scope
    /// those names resolve nowhere (V3): a permissive compile warns and
    /// the source, reading None, yields nothing; a strict one refuses it.
    #[test]
    fn from_ast_refuses_a_context_required_source_by_name() {
        let ast = Comprehension::cartesian(vec![
            clause("k", &[1, 2]),
            Comprehension::clause(
                "j",
                Source::Generator {
                    expr: "pow2({n})".into(),
                    cardinality_hint: None,
                },
            ),
        ]);
        let err =
            CompiledComprehension::from_ast_in(&ast, Mode::Permissive, &|n| n == "n").unwrap_err();
        assert!(
            matches!(
                err,
                ValidationError::ContextRequired { ref name, ref references }
                    if name == "j" && references == &["n".to_string()]
            ),
            "{err}"
        );
        assert!(err.to_string().contains("traverse it with `for`"), "{err}");
        let (compiled, report) = CompiledComprehension::from_ast_with(&ast, Mode::Permissive)
            .unwrap_or_else(|e| panic!("{e}"));
        assert!(
            matches!(report.warnings.as_slice(),
                [ValidationWarning::UnresolvedNames { reads }] if reads.len() == 1),
            "{:?}",
            report.warnings
        );
        assert_eq!(compiled.coordinate_stream().count(), 0);
        let err = CompiledComprehension::from_ast_with(&ast, Mode::Strict).unwrap_err();
        assert!(
            matches!(err, ValidationError::V3UnresolvedNames { ref reads } if reads.len() == 1),
            "{err}"
        );
    }

    /// A predicate naming what its tuples do not bind has no scope to
    /// resolve in on a coordinate stream: compiled in a scope that has
    /// the name, the compile refuses it by name. With no scope the name
    /// resolves nowhere (V3): a permissive compile warns and the name
    /// reads None, so only a tuple the predicate decides before reading
    /// it is kept; a strict one refuses it.
    #[test]
    fn from_ast_refuses_a_predicate_that_needs_a_scope() {
        let ast = Comprehension::filter(clause("k", &[1, 2, 3]), "{k} == 1 || {k} > {limit}");
        let err = CompiledComprehension::from_ast_in(&ast, Mode::Permissive, &|n| n == "limit")
            .unwrap_err();
        assert!(
            matches!(
                err,
                ValidationError::PredicateContextRequired { ref references, .. }
                    if references == &["limit".to_string()]
            ),
            "{err}"
        );
        assert!(err.to_string().contains("traverse it with `for`"), "{err}");
        let compiled = CompiledComprehension::from_ast(&ast).unwrap_or_else(|e| panic!("{e}"));
        let kept: Vec<_> = compiled
            .coordinate_stream()
            .collect::<Result<_, _>>()
            .unwrap();
        assert_eq!(kept.len(), 1);
        assert_eq!(
            kept[0].bindings[0].1,
            crate::iteration::comprehension::strategies::TupleValue::I64(1)
        );
        let err = CompiledComprehension::from_ast_with(&ast, Mode::Strict).unwrap_err();
        assert!(
            matches!(err, ValidationError::V3UnresolvedNames { .. }),
            "{err}"
        );
        let bound = Comprehension::filter(clause("k", &[1, 2, 3]), "{k} > 1");
        assert!(CompiledComprehension::from_ast(&bound).is_ok());
    }
}