polydat_core/iteration/comprehension/

eval_source.rs

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

//! Source evaluation — comprehension_forms.md §10.7.0, §10.7.6,
//! §10.7.8.
//!
//! Makes [`IndexFn`] a contextual query as well as a static AST
//! property: every [`Source`] variant answers
//! `evaluate(ctx) -> EvaluatedSource` carrying its materialized
//! values, observed cardinality, and the index function the
//! emitted values actually satisfy.
//!
//! ## Why this layer exists
//!
//! [`crate::iteration::comprehension::metadata`] computes `IndexFn`
//! at AST-construction time from static source attributes
//! (`cardinality_hint`, declared step, etc.). Two classes of sources
//! cannot claim their real `IndexFn` that way:
//!
//! - **`Source::Generator { expr }`** — the source text resolves
//!   to a list whose shape is only known after evaluation. The
//!   static path declares `Lattice { axis_sizes: [N] }` from
//!   `cardinality_hint` (or `Unbounded` without it).
//! - **`Source::WorkloadParamList { name }`** — the parameter's
//!   list contents are unknown until kernel evaluation.
//!
//! Non-`Lex` strategies (Diagonal / Extrema / Shells / Halton /
//! Sobol / Lhs) need the input's real `IndexFn` shape to
//! check V4 and dispatch their indexed-form algorithms, so V4
//! fires at strategy invocation against the evaluated shape.
//!
//! ## Eval classes
//!
//! Per comprehension_forms.md §10.7.0, sources partition into three
//! eval classes:
//!
//! | Class | Variants | `evaluate(None)` works? |
//! |---|---|---|
//! | [`EvalClass::Static`] | `Literal`, `IntRange`, a `Generator` whose expression references no name | yes |
//! | [`EvalClass::ContextRequired`] | `WorkloadParamList`, a `Generator` whose expression references a name | no — needs `&Context` |
//! | [`EvalClass::Distribution`] | `ContinuousInterval`, `Distribution` (in their "not yet sampled" state) | yes, but `values` is empty — enclosing `Order(_, sampling-strategy, Some(n))` materializes |
//!
//! The class of a generator is decided by the names its expression
//! reads ([`Source::names_read`]), never by a table of
//! generator names: a context-free call evaluates in the empty
//! scope ([`crate::kernel::interp::NoScope`]), and the compile
//! flattens it into a literal of its values
//! (`comprehension::flatten`).
//!
//! [`SourceEval::eval_class`] classifies a source for callers
//! that want to know whether `evaluate(None)` will succeed; the
//! compile-time V4 check in `validate` works from AST metadata
//! and does not consult it. V4 fires again at
//! strategy-invocation time (comprehension_forms.md §10.7.8).
//!
//! ## What this module DOES NOT own
//!
//! - The runtime walker that combines per-clause
//!   `EvaluatedSource`s into the cartesian / zip / union views
//!   strategies actually consume — that lives in
//!   [`crate::iteration::comprehension::runtime`].
//! - The strategy invocation itself — see
//!   [`crate::iteration::comprehension::strategies::Strategy::apply`].
//! - The compile-time V4 fire — see
//!   [`mod@crate::iteration::comprehension::validate`].

use std::sync::Arc;

use crate::ast::Value;
use crate::iteration::comprehension::cardinality::ProductMeasure;
use crate::iteration::comprehension::eval::{SpecError, spec_error};
use crate::iteration::comprehension::metadata::IndexFn;
use crate::iteration::comprehension::source::{LiteralValue, Source};
use crate::kernel::interp::{Layered, Lookup};

/// Result of evaluating one clause's source.
///
/// `values` carries the materialized stream (one [`Value`] per
/// output position). `cardinality` is the count of values
/// (`values.len() as u64`, equivalent to the `IndexFn`'s axis
/// total for discrete sources; `0` for un-sampled continuous
/// sources). `index_fn` is the addressing scheme the emitted
/// values actually satisfy — derived from observed shape for
/// `Generator` / `WorkloadParamList`, declared for static
/// variants.
#[derive(Debug, Clone)]
pub struct EvaluatedSource {
    /// The values, in dispense order.
    pub values: Vec<Value>,
    /// How many values; zero for an unsampled continuous source.
    pub cardinality: u64,
    /// The addressing scheme the values satisfy.
    pub index_fn: IndexFn,
}

/// A name a source read that made it yield nothing
/// (none_semantics.md Rule 1, comprehension_forms.md §5 V3): the name
/// as read, after composition, so a composed `{k_{k}_limits}` read with
/// `k = 3` is `k_3_limits`, and `all(<cursor>)` reads the cursor's
/// extent outputs.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum NoneRead {
    /// Nothing binds the name in the scope the source is evaluated in,
    /// nor does an earlier axis.
    Unbound(String),
    /// The name is bound to None.
    BoundNone(String),
}

impl NoneRead {
    /// The name read.
    pub fn name(&self) -> &str {
        match self {
            NoneRead::Unbound(name) | NoneRead::BoundNone(name) => name,
        }
    }
}

impl std::fmt::Display for NoneRead {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NoneRead::Unbound(name) => write!(f, "`{name}` is not bound"),
            NoneRead::BoundNone(name) => write!(f, "`{name}` is None"),
        }
    }
}

/// The eval-class partition of comprehension_forms.md §10.7.0.
///
/// Tells a caller whether a source can be materialized with
/// `ctx = None`. The compile-time V4 check in `validate` works
/// from AST metadata and does not consult this; V4 otherwise
/// fires at strategy-invocation time.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EvalClass {
    /// Statically evaluable with no kernel / param context.
    /// `evaluate(None)` returns a fully-populated
    /// [`EvaluatedSource`].
    Static,

    /// Requires a kernel context to resolve interpolation
    /// references or workload-param lookups.
    /// `evaluate(None)` returns [`EvalError::NeedsContext`].
    ContextRequired,

    /// Continuous measure / distribution. `evaluate(None)`
    /// succeeds but emits an empty `values` vector; the
    /// `IndexFn` is `Continuous`. The enclosing sampling
    /// `Order(_, strategy, Some(n))` materializes draws.
    Distribution,
}

/// Errors returned by [`SourceEval::evaluate`].
#[derive(Debug, Clone)]
pub enum EvalError {
    /// The source needs a kernel context that wasn't provided.
    NeedsContext,

    /// Evaluation against the supplied context failed. `var`
    /// names the clause; `source` is the spec-text or
    /// description; `message` carries the underlying reason.
    EvalFailed {
        /// The clause's element name.
        var: String,
        /// The source text or description.
        source: String,
        /// The underlying reason.
        message: String,
    },
}

impl std::fmt::Display for EvalError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            EvalError::NeedsContext => f.write_str("source evaluation needs a kernel context"),
            EvalError::EvalFailed {
                var,
                source,
                message,
            } => {
                write!(f, "source '{var} in {source}': {message}")
            }
        }
    }
}

impl std::error::Error for EvalError {}

/// Per-evaluation context for context-required sources.
///
/// Carries the live kernel against which `Source::Generator`
/// spec-text and `Source::WorkloadParamList` lookups resolve.
/// `var_name` lets the source synthesise a useful error
/// message; `prefix` is the prior-axis bindings the evaluator
/// layers in front of `scope` (via `Layered`) so dependent
/// sources see earlier-axis values.
pub struct EvalContext<'a> {
    /// The clause's element name, for messages.
    pub var_name: &'a str,
    /// Where the source's names resolve: the body's scope with the
    /// parent's cascaded wires.
    pub scope: &'a dyn Lookup,
    /// The prior-axis bindings, in axis order.
    pub prefix: &'a [(String, Value)],
}

/// The source-evaluation surface.
///
/// Each [`Source`] variant implements this. The trait is
/// object-safe but typically called through the inherent
/// [`Source`] methods below.
pub trait SourceEval {
    /// Classify this source for the IR planner per
    /// comprehension_forms.md §10.7.0. See [`EvalClass`].
    fn eval_class(&self) -> EvalClass;

    /// Materialize this source.
    ///
    /// Literal / IntRange (`Static`), a context-free Generator
    /// (`Static`, evaluated in the empty scope), and
    /// ContinuousInterval / Distribution (`Distribution`) accept
    /// `ctx = None`. A Generator that references a name and a
    /// WorkloadParamList (`ContextRequired`) require `Some(ctx)` and
    /// return [`EvalError::NeedsContext`] otherwise.
    fn evaluate(&self, ctx: Option<&EvalContext<'_>>) -> Result<EvaluatedSource, EvalError>;
}

impl SourceEval for Source {
    fn eval_class(&self) -> EvalClass {
        match self {
            Source::Literal { .. } | Source::IntRange { .. } => EvalClass::Static,
            Source::ContinuousInterval { .. } | Source::Distribution { .. } => {
                EvalClass::Distribution
            }
            // A generator's class is its expression's: context-free
            // when it reads no name (comprehension_forms.md
            // §10.7.0).
            Source::Generator { .. } if self.names_read().is_empty() => EvalClass::Static,
            Source::Generator { .. } => EvalClass::ContextRequired,
            Source::WorkloadParamList { .. } => EvalClass::ContextRequired,
        }
    }

    fn evaluate(&self, ctx: Option<&EvalContext<'_>>) -> Result<EvaluatedSource, EvalError> {
        evaluate_reading(self, ctx).map(|(evaluated, _)| evaluated)
    }
}

/// [`SourceEval::evaluate`], with the names whose reads made the source
/// yield nothing: a source that reads, after composition, a name nothing
/// binds or a name bound to None yields nothing (none_semantics.md
/// Rule 1, comprehension_forms.md §5 V3), and those names are the second
/// part, empty whenever the source yields.
pub(crate) fn evaluate_reading(
    source: &Source,
    ctx: Option<&EvalContext<'_>>,
) -> Result<(EvaluatedSource, Vec<NoneRead>), EvalError> {
    let evaluated = match source {
        Source::Generator { .. } | Source::WorkloadParamList { .. } => {
            return evaluate_spec_source(source, ctx);
        }
        other => evaluate_static(other),
    };
    Ok((evaluated, Vec::new()))
}

/// Evaluate a `Generator` or `WorkloadParamList` source's spec text in
/// the context's scope, with the prior-axis bindings in front.
fn evaluate_spec_source(
    source: &Source,
    ctx: Option<&EvalContext<'_>>,
) -> Result<(EvaluatedSource, Vec<NoneRead>), EvalError> {
    let spec_text = match source {
        Source::Generator { expr, .. } => expr.clone(),
        Source::WorkloadParamList { name, .. } => format!("{{{name}}}"),
        _ => unreachable!("only spec-text sources"),
    };
    // A context-free generator evaluates in the empty
    // scope; anything that references a name needs the
    // caller's.
    let empty = crate::kernel::interp::NoScope::new();
    let (var_name, scope): (&str, Layered<'_>) = match ctx {
        Some(ctx) => (
            ctx.var_name,
            Layered {
                prefix: ctx.prefix,
                inner: ctx.scope,
            },
        ),
        None if source.eval_class() == EvalClass::Static => (
            "<context-free>",
            Layered {
                prefix: &[],
                inner: &empty,
            },
        ),
        None => return Err(EvalError::NeedsContext),
    };
    match crate::iteration::comprehension::eval::evaluate_spec_internal(&spec_text, &scope) {
        Ok(vals) => {
            let n = vals.len() as u64;
            let index_fn = classify_observed_values(&vals);
            Ok((
                EvaluatedSource {
                    values: vals,
                    cardinality: n,
                    index_fn,
                },
                Vec::new(),
            ))
        }
        // A source over None yields nothing: the evaluation read, after
        // composition, a name nothing binds or one bound to None.
        Err(SpecError::ReadsNone { reads, .. }) => Ok((
            EvaluatedSource {
                values: Vec::new(),
                cardinality: 0,
                index_fn: IndexFn::Lattice {
                    axis_sizes: vec![0],
                },
            },
            reads,
        )),
        Err(SpecError::Failed(message)) => Err(EvalError::EvalFailed {
            var: var_name.to_string(),
            message: spec_error(&spec_text, message).to_string(),
            source: spec_text,
        }),
    }
}

/// The evaluation of a source that reads no name: literals, ranges, and
/// continuous measures.
fn evaluate_static(source: &Source) -> EvaluatedSource {
    match source {
        Source::Literal { values } => {
            let vals: Vec<Value> = values.iter().map(literal_to_value).collect();
            let n = vals.len() as u64;
            EvaluatedSource {
                values: vals,
                cardinality: n,
                // Literal lists carry no shape claim other
                // than length — call them a 1-axis Lattice
                // of that length. Strategies that need
                // arithmetic progression shape (e.g. Halton
                // over a Lattice axis) still get useful
                // behavior because the lookup is by index,
                // not by value.
                index_fn: IndexFn::Lattice {
                    axis_sizes: vec![n],
                },
            }
        }
        Source::IntRange { lo, hi, step } => {
            let step = (*step).max(1);
            let mut vals = Vec::new();
            let mut cur = *lo;
            while cur < *hi {
                vals.push(Value::U64(cur as u64));
                cur += step;
            }
            let n = vals.len() as u64;
            EvaluatedSource {
                values: vals,
                cardinality: n,
                index_fn: IndexFn::Lattice {
                    axis_sizes: vec![n],
                },
            }
        }
        Source::ContinuousInterval { interval, measure } => EvaluatedSource {
            values: Vec::new(),
            cardinality: 0,
            index_fn: IndexFn::Continuous {
                intervals: vec![interval.clone()],
                measure: measure.clone(),
            },
        },
        Source::Distribution {
            distribution,
            support,
            ..
        } => EvaluatedSource {
            values: Vec::new(),
            cardinality: 0,
            // The parameters travel on the AST carrier; the
            // runtime's sampler reads them there
            // (comprehension_forms.md §10.7.6).
            index_fn: IndexFn::Continuous {
                intervals: vec![support.clone()],
                measure: ProductMeasure::Named(*distribution),
            },
        },
        Source::Generator { .. } | Source::WorkloadParamList { .. } => {
            unreachable!("a spec-text source reads names")
        }
    }
}

/// Classify a materialized value list by observed shape.
///
/// The "expand-then-classify" stage of comprehension_forms.md
/// §10.7.6 / §10.7.8: any list of `N` values is a one-axis
/// `Lattice { axis_sizes: [N] }`, whose only shape claim is its
/// length, since a strategy looks values up by position. The shape
/// is read off the evaluated values, never declared from a
/// generator's name.
fn classify_observed_values(vals: &[Value]) -> IndexFn {
    let n = vals.len() as u64;
    IndexFn::Lattice {
        axis_sizes: vec![n],
    }
}

fn literal_to_value(lv: &LiteralValue) -> Value {
    match lv {
        LiteralValue::Int(n) => Value::U64(*n as u64),
        LiteralValue::UInt(n) => Value::U64(*n),
        LiteralValue::Float(f) => Value::F64(*f),
        LiteralValue::String(s) => Value::Str(Arc::from(s.as_str())),
        LiteralValue::Bool(b) => Value::Bool(*b),
        LiteralValue::Json(j) => Value::Json(Arc::new(j.clone())),
    }
}

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

    #[test]
    fn literal_evaluates_without_context() {
        let s = Source::Literal {
            values: vec![
                LiteralValue::Int(1),
                LiteralValue::Int(2),
                LiteralValue::Int(3),
            ],
        };
        assert_eq!(s.eval_class(), EvalClass::Static);
        let ev = s.evaluate(None).unwrap();
        assert_eq!(ev.cardinality, 3);
        assert_eq!(ev.values.len(), 3);
        assert!(matches!(ev.index_fn, IndexFn::Lattice { axis_sizes: ref a } if a == &vec![3]));
    }

    #[test]
    fn int_range_evaluates_without_context() {
        let s = Source::IntRange {
            lo: 0,
            hi: 10,
            step: 2,
        };
        assert_eq!(s.eval_class(), EvalClass::Static);
        let ev = s.evaluate(None).unwrap();
        // 0, 2, 4, 6, 8 = 5 values
        assert_eq!(ev.cardinality, 5);
        assert!(matches!(ev.index_fn, IndexFn::Lattice { axis_sizes: ref a } if a == &vec![5]));
    }

    #[test]
    fn a_context_free_generator_evaluates_without_context() {
        let s = Source::Generator {
            expr: "fib(6)".into(),
            cardinality_hint: None,
        };
        assert_eq!(s.eval_class(), EvalClass::Static);
        let ev = s.evaluate(None).unwrap();
        assert_eq!(ev.cardinality, 6);
    }

    #[test]
    fn generator_without_context_errors() {
        let s = Source::Generator {
            expr: "range(0, {n})".into(),
            cardinality_hint: Some(10),
        };
        assert_eq!(s.eval_class(), EvalClass::ContextRequired);
        match s.evaluate(None) {
            Err(EvalError::NeedsContext) => {}
            other => panic!("expected NeedsContext, got {other:?}"),
        }
    }

    #[test]
    fn workload_param_list_without_context_errors() {
        let s = Source::WorkloadParamList {
            name: "k_values".into(),
            len_hint: Some(5),
        };
        assert_eq!(s.eval_class(), EvalClass::ContextRequired);
        assert!(matches!(s.evaluate(None), Err(EvalError::NeedsContext)));
    }

    #[test]
    fn continuous_interval_yields_continuous_index_fn() {
        let s = Source::ContinuousInterval {
            interval: Interval::closed(0.0, 1.0),
            measure: ProductMeasure::Uniform,
        };
        assert_eq!(s.eval_class(), EvalClass::Distribution);
        let ev = s.evaluate(None).unwrap();
        assert_eq!(ev.cardinality, 0);
        assert!(ev.values.is_empty());
        match ev.index_fn {
            IndexFn::Continuous { intervals, .. } => assert_eq!(intervals.len(), 1),
            other => panic!("expected Continuous, got {other:?}"),
        }
    }

    #[test]
    fn distribution_yields_continuous_index_fn() {
        let s = Source::Distribution {
            distribution: MeasureName::Normal,
            support: Interval {
                lo: f64::NEG_INFINITY,
                hi: f64::INFINITY,
                lo_open: true,
                hi_open: true,
            },
            params: vec![0.0, 1.0],
        };
        assert_eq!(s.eval_class(), EvalClass::Distribution);
        let ev = s.evaluate(None).unwrap();
        assert_eq!(ev.cardinality, 0);
        assert!(matches!(
            ev.index_fn,
            IndexFn::Continuous {
                measure: ProductMeasure::Named(MeasureName::Normal),
                ..
            }
        ));
    }

    #[test]
    fn generator_with_context_evaluates_to_lattice() {
        let canonical = Arc::new(crate::dsl::compile_polydat_interpreter("\n").unwrap());
        let s = Source::Generator {
            expr: "1, 2, 3, 4, 5".into(),
            cardinality_hint: Some(5),
        };
        let ctx = EvalContext {
            var_name: "k",
            scope: &*canonical,
            prefix: &[],
        };
        let ev = s.evaluate(Some(&ctx)).unwrap();
        assert_eq!(ev.cardinality, 5);
        assert!(matches!(ev.index_fn, IndexFn::Lattice { .. }));
    }
}