polydat_grammar/comprehension/

source.rs

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

//! Clause source values (comprehension_forms.md §3.1).
//!
//! A `clause(name, source)` binds a name to the values
//! produced by its source. Sources split into two families:
//!
//! - **Discrete stream producers** — literal lists, integer
//!   ranges, generator functions, workload-param references.
//!   Cardinality is `Bounded`, `BoundedAtMost`, or `Unbounded`.
//! - **Continuous measures** — real intervals with an
//!   integrable measure (uniform on bounded intervals; named
//!   probability distributions like Normal / Exponential).
//!   Cardinality is `Continuous`; V8 requires an enclosing
//!   sampling `order(_, strategy, Some(n))` before dispense.
//!
//! Sources are stream producers — they do not pre-materialize
//! into `Vec<Value>`. This is the load-bearing model property
//! of comprehension_forms.md §3.1 and §6.2.

use serde::{Deserialize, Serialize};

use super::cardinality::{CardinalityClass, Interval, MeasureName, ProductMeasure};

/// A clause's source of values.
///
/// Discrete variants produce a stream of `Value` via the
/// runtime evaluator; continuous variants describe a measure
/// that a downstream sampling strategy will draw from.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Source {
    /// Literal comma list (e.g., `[1, 2, 4, 8]`). Stream
    /// producer over the list contents.
    Literal {
        /// The values, in order.
        values: Vec<LiteralValue>,
    },

    /// Integer half-open range `lo..hi` with optional step.
    /// Default step is 1.
    IntRange {
        /// The first value.
        lo: i64,
        /// One past the last.
        hi: i64,
        /// The step between values, 1 by default.
        step: i64,
    },

    /// Generator function call expressed as a Polydat source string.
    /// Its eval class follows its free names (comprehension_forms.md §10.7.0): an
    /// expression that references no name is context-free and is
    /// evaluated at compile, so a clause over it becomes a literal of
    /// its values; one that references a name (an outer coordinate,
    /// a parameter, a wire) is evaluated at traversal.
    Generator {
        /// The generator call, as Polydat source.
        expr: String,
        /// How many values it yields, when known: the count the
        /// compile established by evaluating a context-free call
        /// whose values are not literal-representable (a partition
        /// list, for one). `None` for a call evaluated at traversal,
        /// whose cardinality is `Unbounded` until then.
        cardinality_hint: Option<u64>,
    },

    /// Reference to a workload-level parameter that resolves to
    /// a list of values. Cardinality is the parameter's
    /// declared list length.
    WorkloadParamList {
        /// The parameter's name.
        name: String,
        /// The list's length, when known.
        len_hint: Option<u64>,
    },

    /// Real interval (continuous source). Combined with a
    /// `measure` to form a `Continuous` cardinality.
    /// Integrability is checked at parse via V8.
    ContinuousInterval {
        /// The interval.
        interval: Interval,
        /// The measure drawn from.
        measure: ProductMeasure,
    },

    /// Named continuous distribution. The distribution carries
    /// its own support; the `support` field records the
    /// effective interval for V8's check.
    Distribution {
        /// The distribution.
        distribution: MeasureName,
        /// Its effective support.
        support: Interval,
        /// Its parameters, in the order of
        /// [`MeasureName::parameter_names`]; empty for the standard
        /// parameters ([`MeasureName::default_params`]).
        params: Vec<f64>,
    },
}

/// A literal value carried in a `Source::Literal`. Subset of
/// the polydat `Value` type — the kinds clauses can directly
/// bind. Extension to richer value types lives in the source
/// evaluator, not the AST.
///
/// Serialized untagged because the variants are primitives;
/// the JSON/YAML representation is just the bare value
/// (`1` / `"x"` / `true` / `1.5`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum LiteralValue {
    /// An integer.
    Int(i64),
    /// An integer above `i64::MAX`, which only a `u64` holds. An
    /// integer that fits `i64` is `Int`, so every integer has one
    /// form ([`LiteralValue::unsigned`]). After `Int`, so an untagged
    /// read takes an integer as `Int` when it fits.
    UInt(u64),
    /// A float.
    Float(f64),
    /// A string.
    String(String),
    /// A boolean.
    Bool(bool),
    /// A JSON value carrying its own kind: an item of a JSON list a
    /// generator supplied at run time, bound where the element is
    /// declared `json`. Last, so an untagged read tries the scalar
    /// forms first.
    Json(serde_json::Value),
}

impl LiteralValue {
    /// The literal of an unsigned integer: `Int` when it fits `i64`,
    /// `UInt` above.
    pub fn unsigned(n: u64) -> Self {
        match i64::try_from(n) {
            Ok(n) => LiteralValue::Int(n),
            Err(_) => LiteralValue::UInt(n),
        }
    }
}

impl Source {
    /// The names this source references (comprehension_forms.md §10.7.0): a generator
    /// expression's parsed free identifiers (`concat(foo)`) and its
    /// `{name}` interpolation placeholders, and a workload parameter
    /// list's own name. Literals, ranges, and intervals reference
    /// nothing. A source with no references is context-free.
    pub fn referenced_names(&self) -> std::collections::BTreeSet<String> {
        let mut out = std::collections::BTreeSet::new();
        match self {
            Source::WorkloadParamList { name, .. } => {
                out.insert(name.clone());
            }
            Source::Generator { expr, .. } => {
                out.extend(crate::refs::referenced_names(expr));
                crate::refs::collect_string_interpolation_refs(expr, &mut out);
            }
            Source::Literal { .. }
            | Source::IntRange { .. }
            | Source::ContinuousInterval { .. }
            | Source::Distribution { .. } => {}
        }
        out
    }

    /// The names evaluating this source reads from the prior axes and
    /// the scope it is evaluated in (comprehension_forms.md §5 V3,
    /// §10.9.1). This differs from [`Self::referenced_names`] in two
    /// forms. A composed name (`{k_{k}_limits}`) reads its leaves
    /// (`k`) here; the name they compose to is known only once the
    /// leaves are bound, and is read when the source is evaluated. The
    /// cursor form `all(<cursor>)` reads the cursor's extent outputs
    /// ([`cursor_extent_names`]), not a value named after the cursor.
    pub fn names_read(&self) -> std::collections::BTreeSet<String> {
        let mut out = std::collections::BTreeSet::new();
        match self {
            Source::WorkloadParamList { name, .. } => {
                crate::refs::collect_string_interpolation_refs(&format!("{{{name}}}"), &mut out);
            }
            Source::Generator { expr, .. } => match all_cursor_argument(expr) {
                Some(cursor) => out.extend(cursor_extent_names(cursor)),
                None => out.extend(self.referenced_names()),
            },
            Source::Literal { .. }
            | Source::IntRange { .. }
            | Source::ContinuousInterval { .. }
            | Source::Distribution { .. } => {}
        }
        out
    }

    /// Declare this source's cardinality class for use by
    /// `clause` metadata propagation.
    pub fn cardinality(&self) -> CardinalityClass {
        match self {
            Source::Literal { values } => CardinalityClass::Bounded(values.len() as u64),
            Source::IntRange { lo, hi, step } => {
                let step = (*step).max(1).unsigned_abs();
                if hi <= lo {
                    CardinalityClass::Bounded(0)
                } else {
                    let span = (hi - lo) as u64;
                    let n = span.div_ceil(step);
                    CardinalityClass::Bounded(n)
                }
            }
            Source::Generator {
                cardinality_hint, ..
            } => match cardinality_hint {
                Some(n) => CardinalityClass::Bounded(*n),
                None => CardinalityClass::Unbounded,
            },
            Source::WorkloadParamList { len_hint, .. } => match len_hint {
                Some(n) => CardinalityClass::Bounded(*n),
                None => CardinalityClass::Unbounded,
            },
            Source::ContinuousInterval { interval, measure } => CardinalityClass::Continuous {
                intervals: vec![interval.clone()],
                measure: measure.clone(),
            },
            Source::Distribution { support, .. } => CardinalityClass::Continuous {
                intervals: vec![support.clone()],
                measure: ProductMeasure::Named(*self.distribution_name()),
            },
        }
    }

    /// `true` if this source is continuous (Continuous /
    /// Distribution variants). Used by V7 (zip must be all
    /// discrete) and V9 (union must be all discrete) without
    /// a full cardinality computation.
    pub fn is_continuous(&self) -> bool {
        matches!(
            self,
            Source::ContinuousInterval { .. } | Source::Distribution { .. }
        )
    }

    /// `true` if this source is discrete (every variant except
    /// the continuous ones).
    pub fn is_discrete(&self) -> bool {
        !self.is_continuous()
    }

    fn distribution_name(&self) -> &MeasureName {
        match self {
            Source::Distribution { distribution, .. } => distribution,
            _ => panic!("distribution_name called on non-Distribution source"),
        }
    }
}

// ── iteration interior and string-comprehension striping
//    (comprehension_forms.md §3.1.2, §3.1.3) ──

/// The string-comprehension separator rule (comprehension_forms.md
/// §3.1.3), in one place
/// so the parse-time (`source_parser`) and runtime (`eval`)
/// striping can never drift: split on runs of comma / semicolon /
/// ASCII whitespace, trim, drop empties. Every other character
/// (`:` `.` `-` `/` …) stays in the token. Returns the raw token
/// substrings; callers type them (Value or LiteralValue).
pub fn split_string_comprehension(s: &str) -> Vec<&str> {
    s.split(|c: char| c == ',' || c == ';' || c.is_ascii_whitespace())
        .map(str::trim)
        .filter(|t| !t.is_empty())
        .collect()
}

// ── the cursor form `all(<cursor>)` (comprehension_forms.md §10.9.1) ──

/// The cursor `text` enumerates when it is the source form
/// `all(<cursor>)`, whitespace aside; `None` for any other text.
pub fn all_cursor_argument(text: &str) -> Option<&str> {
    let cursor = text.trim().strip_prefix("all(")?.strip_suffix(')')?.trim();
    let mut chars = cursor.chars();
    let starts = chars
        .next()
        .is_some_and(|c| c.is_ascii_alphabetic() || c == '_');
    (starts && chars.all(|c| c.is_ascii_alphanumeric() || c == '_')).then_some(cursor)
}

/// The auxiliary outputs a `cursor <name> = ...` declaration compiles
/// to, holding its extent: `[start, end]`. `all(<cursor>)` reads these.
pub fn cursor_extent_names(cursor: &str) -> [String; 2] {
    [
        format!("__cursor_extent_{cursor}_start"),
        format!("__cursor_extent_{cursor}_end"),
    ]
}

/// The cursor whose extent the output `name` holds, when `name` is one
/// of [`cursor_extent_names`].
pub fn cursor_of_extent_name(name: &str) -> Option<&str> {
    let rest = name.strip_prefix("__cursor_extent_")?;
    rest.strip_suffix("_start")
        .or_else(|| rest.strip_suffix("_end"))
        .filter(|cursor| !cursor.is_empty())
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A composed name reads its leaves, and the cursor form reads the
    /// cursor's extent; every other source reads what it references.
    #[test]
    fn names_read_are_the_leaves_of_a_composition_and_a_cursor_extent() {
        let names = |s: Source| s.names_read().into_iter().collect::<Vec<_>>();
        assert_eq!(
            names(Source::WorkloadParamList {
                name: "k_{k}_limits".into(),
                len_hint: None,
            }),
            ["k"]
        );
        assert_eq!(
            names(Source::WorkloadParamList {
                name: "k_values".into(),
                len_hint: None,
            }),
            ["k_values"]
        );
        assert_eq!(
            names(Source::Generator {
                expr: " all( row ) ".into(),
                cardinality_hint: None,
            }),
            ["__cursor_extent_row_end", "__cursor_extent_row_start"]
        );
        assert_eq!(
            names(Source::Generator {
                expr: "pow2({n})".into(),
                cardinality_hint: None,
            }),
            ["n"]
        );
        assert_eq!(all_cursor_argument("all(1)"), None);
        assert_eq!(all_cursor_argument("all(a, b)"), None);
        for extent in cursor_extent_names("row") {
            assert_eq!(cursor_of_extent_name(&extent), Some("row"));
        }
        assert_eq!(cursor_of_extent_name("row"), None);
    }

    #[test]
    fn literal_cardinality_is_list_length() {
        let s = Source::Literal {
            values: vec![
                LiteralValue::Int(1),
                LiteralValue::Int(2),
                LiteralValue::Int(3),
            ],
        };
        assert!(matches!(s.cardinality(), CardinalityClass::Bounded(3)));
    }

    /// An integer has one literal form, `Int` up to `i64::MAX` and
    /// `UInt` above, and its serialized form reads back as that form.
    #[test]
    fn an_unsigned_literal_keeps_its_value_through_serde() {
        assert_eq!(
            LiteralValue::unsigned(i64::MAX as u64),
            LiteralValue::Int(i64::MAX)
        );
        assert_eq!(
            LiteralValue::unsigned(u64::MAX),
            LiteralValue::UInt(u64::MAX)
        );
        for v in [
            LiteralValue::Int(-3),
            LiteralValue::Int(i64::MAX),
            LiteralValue::UInt(1 << 63),
            LiteralValue::UInt(u64::MAX),
        ] {
            let json = serde_json::to_string(&v).unwrap();
            let back: LiteralValue = serde_json::from_str(&json).unwrap();
            assert_eq!(back, v, "{json}");
        }
    }

    #[test]
    fn int_range_step_1() {
        let s = Source::IntRange {
            lo: 1,
            hi: 10,
            step: 1,
        };
        assert!(matches!(s.cardinality(), CardinalityClass::Bounded(9)));
    }

    #[test]
    fn int_range_with_step() {
        let s = Source::IntRange {
            lo: 0,
            hi: 10,
            step: 2,
        };
        // 0,2,4,6,8 = 5 values
        assert!(matches!(s.cardinality(), CardinalityClass::Bounded(5)));
    }

    #[test]
    fn int_range_empty() {
        let s = Source::IntRange {
            lo: 5,
            hi: 5,
            step: 1,
        };
        assert!(matches!(s.cardinality(), CardinalityClass::Bounded(0)));
    }

    #[test]
    fn generator_without_hint_is_unbounded() {
        let s = Source::Generator {
            expr: "live_query()".into(),
            cardinality_hint: None,
        };
        assert!(matches!(s.cardinality(), CardinalityClass::Unbounded));
    }

    #[test]
    fn generator_with_hint_is_bounded() {
        let s = Source::Generator {
            expr: "first_100()".into(),
            cardinality_hint: Some(100),
        };
        assert!(matches!(s.cardinality(), CardinalityClass::Bounded(100)));
    }

    #[test]
    fn continuous_interval_produces_continuous_class() {
        let s = Source::ContinuousInterval {
            interval: Interval::closed(0.0, 1.0),
            measure: ProductMeasure::Uniform,
        };
        match s.cardinality() {
            CardinalityClass::Continuous { intervals, measure } => {
                assert_eq!(intervals.len(), 1);
                assert!(matches!(measure, ProductMeasure::Uniform));
            }
            other => panic!("expected Continuous, got {other:?}"),
        }
        assert!(s.is_continuous());
        assert!(!s.is_discrete());
    }

    #[test]
    fn distribution_source_classification() {
        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!(s.is_continuous());
        match s.cardinality() {
            CardinalityClass::Continuous {
                measure: ProductMeasure::Named(MeasureName::Normal),
                ..
            } => {}
            other => panic!("expected Continuous with Named(Normal), got {other:?}"),
        }
    }
}