polydat_grammar/comprehension/

ast.rs

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

//! Operator-tree comprehension AST (comprehension_forms.md §3).
//!
//! Six constructors closed under composition: one source
//! (`clause`), three combinators (`cartesian`, `zip`, `union`),
//! two modifiers (`filter`, `order`). Every comprehension AST is
//! a tree whose nodes are one of these six variants.
//!
//! Closure (comprehension_forms.md §4.1):
//!
//! - C1 — every constructor returns and consumes
//!   `Comprehension`. There is no auxiliary value type at the
//!   AST level.
//! - C2 — well-formedness is decidable in one bottom-up pass
//!   (the `validate` module of the runtime
//!   implements the check).

use serde::{Deserialize, Serialize};

use super::source::Source;
use super::strategy::{StrategyName, ZipMode};

/// The six-variant operator-tree comprehension.
///
/// Closure under composition (comprehension_forms.md §4.1 C1): every variant
/// holds one or more `Comprehension` operands plus
/// constructor-specific scalar parameters (predicate, strategy,
/// truncation, zip mode, source).
///
/// `Box<Comprehension>` appears wherever a variant needs a
/// single child operand; `Vec<Comprehension>` wherever a
/// constructor takes N children (cartesian, zip, union).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum Comprehension {
    /// Leaf source (comprehension_forms.md §3.1). Binds `name` to one value
    /// per dispense, drawn from `source`.
    Clause {
        /// The name bound.
        name: String,
        /// Where the values come from.
        source: Source,
    },

    /// Cross-product combinator (comprehension_forms.md §3.2). Children must
    /// have disjoint name sets (V1).
    Cartesian {
        /// The factors.
        children: Vec<Comprehension>,
    },

    /// Lockstep combinator (comprehension_forms.md §3.3). Children must be
    /// discrete (V7) and have disjoint name sets (V1).
    Zip {
        /// The streams zipped.
        children: Vec<Comprehension>,
        /// The length policy.
        mode: ZipMode,
    },

    /// Concatenation combinator (comprehension_forms.md §3.4). Children must
    /// share an identical tuple shape (V2) and all be discrete
    /// (V9).
    Union {
        /// The streams concatenated, in order.
        children: Vec<Comprehension>,
    },

    /// Selection modifier (comprehension_forms.md §3.5). The
    /// predicate is a Polydat boolean expression (§10.9.1); names
    /// must close over the child's
    /// coordinates plus the parent scope (V3).
    Filter {
        /// The stream filtered.
        child: Box<Comprehension>,
        /// The predicate, a boolean expression over the tuple and the parent scope.
        predicate: String,
    },

    /// Permutation modifier (comprehension_forms.md §3.6). `strategy` must
    /// accept the child's IndexFn (V4); `truncation` limits the
    /// dispensed count; `seed` is the authored seed a seeded
    /// strategy (`Shuffle`, `Lhs`) derives its state from, with a
    /// fixed default when absent.
    Order {
        /// The stream ordered.
        child: Box<Comprehension>,
        /// The strategy applied.
        strategy: StrategyName,
        /// The dispensed count cap, if any.
        truncation: Option<u64>,
        /// The authored seed, read by `Shuffle` and `Lhs`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        seed: Option<u64>,
    },
}

impl Comprehension {
    /// Construct a leaf clause.
    pub fn clause<S: Into<String>>(name: S, source: Source) -> Self {
        Comprehension::Clause {
            name: name.into(),
            source,
        }
    }

    /// Construct a cartesian over the supplied children.
    pub fn cartesian(children: Vec<Comprehension>) -> Self {
        Comprehension::Cartesian { children }
    }

    /// Construct a zip over the supplied children with the
    /// given mode.
    pub fn zip(children: Vec<Comprehension>, mode: ZipMode) -> Self {
        Comprehension::Zip { children, mode }
    }

    /// Construct a union over the supplied children.
    pub fn union(children: Vec<Comprehension>) -> Self {
        Comprehension::Union { children }
    }

    /// Construct a filter wrapping `child` with `predicate`.
    pub fn filter<S: Into<String>>(child: Comprehension, predicate: S) -> Self {
        Comprehension::Filter {
            child: Box::new(child),
            predicate: predicate.into(),
        }
    }

    /// Construct an order node wrapping `child`, with no authored
    /// seed.
    pub fn order(child: Comprehension, strategy: StrategyName, truncation: Option<u64>) -> Self {
        Self::order_seeded(child, strategy, truncation, None)
    }

    /// Construct an order node wrapping `child` with an authored
    /// `seed`, which `Shuffle` and `Lhs` derive their state from.
    pub fn order_seeded(
        child: Comprehension,
        strategy: StrategyName,
        truncation: Option<u64>,
        seed: Option<u64>,
    ) -> Self {
        Comprehension::Order {
            child: Box::new(child),
            strategy,
            truncation,
            seed,
        }
    }

    /// Compute the comprehension's coordinate name set,
    /// recursively. The result preserves declaration order
    /// (comprehension_forms.md §3.2 and §3.4's "in declaration order" tuple
    /// shape rules). Used by V1, V2, V3, and the predicate
    /// analyzer's coord-set input.
    pub fn coordinate_names(&self) -> Vec<String> {
        let mut acc = Vec::new();
        self.collect_coordinate_names(&mut acc);
        acc
    }

    /// Compute `(coordinate_name, source_text)` pairs in
    /// declaration order, deduplicated by name (first
    /// occurrence wins). Source text is the clause text form —
    /// `IntRange { 1, 10, 1 }` → `"1..10"`,
    /// `Literal { [10, 100] }` → `"10, 100"`, etc.
    ///
    /// Gives a consumer a `[(var, spec_expr)]` list for type
    /// detection, the spec-text shape the runtime's probe
    /// pre-evaluation (`pre_evaluate_clause`) takes.
    pub fn coordinate_specs(&self) -> Vec<(String, String)> {
        let mut acc = Vec::new();
        let mut seen = std::collections::HashSet::new();
        self.collect_coordinate_specs(&mut acc, &mut seen);
        acc
    }

    /// Grammar-based extraction of the free names referenced by
    /// every source spec in this comprehension subtree —
    /// workload params, outer iter-vars, and wires that a
    /// `Generator` spec (`concat(foo)`, bare `eh_values`)
    /// consumes. Each spec is parsed with the canonical Polydat
    /// expression grammar (`crate::refs::referenced_names`)
    /// rather than byte-scanned, so a bare source reference is
    /// recognised exactly as the kernel compiler would resolve
    /// it. `WorkloadParamList { name }` contributes `name`
    /// directly; literals / ranges / intervals contribute
    /// nothing. Used by the workload validator's
    /// declared-but-unreferenced check.
    pub fn referenced_source_names(&self) -> std::collections::BTreeSet<String> {
        let mut out = std::collections::BTreeSet::new();
        self.walk_sources(&mut |source| out.extend(source.referenced_names()));
        out
    }

    /// Visit every leaf [`Source`] in this comprehension subtree.
    fn walk_sources(&self, visit: &mut impl FnMut(&super::source::Source)) {
        match self {
            Comprehension::Clause { source, .. } => visit(source),
            Comprehension::Cartesian { children }
            | Comprehension::Zip { children, .. }
            | Comprehension::Union { children } => {
                for c in children {
                    c.walk_sources(visit);
                }
            }
            Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
                child.walk_sources(visit);
            }
        }
    }

    fn collect_coordinate_specs(
        &self,
        acc: &mut Vec<(String, String)>,
        seen: &mut std::collections::HashSet<String>,
    ) {
        match self {
            Comprehension::Clause { name, source } => {
                if seen.insert(name.clone()) {
                    let spec_text = source.to_text().unwrap_or_else(|| "<source>".to_string());
                    acc.push((name.clone(), spec_text));
                }
            }
            Comprehension::Cartesian { children } | Comprehension::Zip { children, .. } => {
                for c in children {
                    c.collect_coordinate_specs(acc, seen);
                }
            }
            Comprehension::Union { children } => {
                for c in children {
                    c.collect_coordinate_specs(acc, seen);
                }
            }
            Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
                child.collect_coordinate_specs(acc, seen);
            }
        }
    }

    fn collect_coordinate_names(&self, acc: &mut Vec<String>) {
        match self {
            Comprehension::Clause { name, .. } => {
                if !acc.contains(name) {
                    acc.push(name.clone());
                }
            }
            Comprehension::Cartesian { children } | Comprehension::Zip { children, .. } => {
                for c in children {
                    c.collect_coordinate_names(acc);
                }
            }
            Comprehension::Union { children } => {
                // V2 requires identical shape; take the first
                // child's coordinates as canonical.
                if let Some(first) = children.first() {
                    first.collect_coordinate_names(acc);
                }
            }
            Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
                child.collect_coordinate_names(acc);
            }
        }
    }

    /// `true` if this node is a leaf clause.
    pub fn is_clause(&self) -> bool {
        matches!(self, Comprehension::Clause { .. })
    }

    /// `true` if this node is one of the three combinators.
    pub fn is_combinator(&self) -> bool {
        matches!(
            self,
            Comprehension::Cartesian { .. }
                | Comprehension::Zip { .. }
                | Comprehension::Union { .. }
        )
    }

    /// `true` if this node is a modifier (`filter` or `order`).
    pub fn is_modifier(&self) -> bool {
        matches!(
            self,
            Comprehension::Filter { .. } | Comprehension::Order { .. }
        )
    }

    /// Iterate this node's direct operand children. Returns
    /// an empty iterator for leaf clauses.
    pub fn children(&self) -> Box<dyn Iterator<Item = &Comprehension> + '_> {
        match self {
            Comprehension::Clause { .. } => Box::new(std::iter::empty()),
            Comprehension::Cartesian { children }
            | Comprehension::Zip { children, .. }
            | Comprehension::Union { children } => Box::new(children.iter()),
            Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
                Box::new(std::iter::once(child.as_ref()))
            }
        }
    }

    /// Count of nodes in the AST (this node + all descendants).
    /// Used by the optimizer's well-founded measure for
    /// termination (comprehension_forms.md §10.6, property 3).
    pub fn node_count(&self) -> usize {
        1 + self.children().map(|c| c.node_count()).sum::<usize>()
    }

    /// Maximum depth of the AST. Constant for flat composition,
    /// O(log N) for balanced trees. Bounds the operator stack
    /// (comprehension_forms.md §9.3).
    pub fn depth(&self) -> usize {
        1 + self.children().map(|c| c.depth()).max().unwrap_or(0)
    }
}

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

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

    #[test]
    fn clause_coordinates() {
        let c = lit_int_clause("k", &[1, 2, 3]);
        assert_eq!(c.coordinate_names(), vec!["k"]);
        assert!(c.is_clause());
        assert!(!c.is_combinator());
        assert!(!c.is_modifier());
    }

    #[test]
    fn continuous_interval_spec_text_round_trips_as_float() {
        // A continuous `[1.0, 5.0)` must reconstruct to a spec the
        // source parser re-classifies as CONTINUOUS, not an integer
        // range. `{}` on `1.0f64` prints "1", so the reconstruction
        // must use a float-preserving format ("1.0..5.0"), else a
        // downstream type-probe re-parses "1..5" and types the
        // iter-var `U64` — silently corrupting a float optimize axis.
        use crate::comprehension::cardinality::{Interval, ProductMeasure};
        let c = Comprehension::clause(
            "ef",
            Source::ContinuousInterval {
                interval: Interval {
                    lo: 1.0,
                    hi: 5.0,
                    lo_open: false,
                    hi_open: true,
                },
                measure: ProductMeasure::Uniform,
            },
        );
        let (var, spec_text) = c.coordinate_specs().into_iter().next().unwrap();
        assert_eq!(var, "ef");
        // Re-parsing the reconstructed text must yield a continuous
        // interval again — the round-trip the kernel-type probe relies on.
        let reparsed = crate::comprehension::spec::parse_source(&spec_text).unwrap();
        assert!(
            matches!(reparsed, Source::ContinuousInterval { .. }),
            "reconstructed '{spec_text}' re-parsed to {reparsed:?}, expected ContinuousInterval"
        );
    }

    #[test]
    fn referenced_source_names_grammar_based() {
        // `eh in eh_values` — a bare source reference parses to
        // a Generator whose free name is the workload param.
        let bare = Comprehension::clause(
            "eh",
            Source::Generator {
                expr: "eh_values".into(),
                cardinality_hint: None,
            },
        );
        let got: Vec<String> = bare.referenced_source_names().into_iter().collect();
        assert_eq!(got, vec!["eh_values"]);

        // `(nbo) in (concat(nbo_v_values))` — the source is a
        // function call; the callee `concat` is NOT a reference
        // but its argument IS.
        let call = Comprehension::clause(
            "nbo",
            Source::Generator {
                expr: "concat(nbo_v_values)".into(),
                cardinality_hint: None,
            },
        );
        let got: Vec<String> = call.referenced_source_names().into_iter().collect();
        assert_eq!(got, vec!["nbo_v_values"]);

        // `{profiles}` — an explicit WorkloadParamList contributes
        // its name directly.
        let wpl = Comprehension::clause(
            "p",
            Source::WorkloadParamList {
                name: "profiles".into(),
                len_hint: None,
            },
        );
        let got: Vec<String> = wpl.referenced_source_names().into_iter().collect();
        assert_eq!(got, vec!["profiles"]);

        // Literal sources contribute nothing.
        let lit = lit_int_clause("k", &[1, 2, 3]);
        assert!(lit.referenced_source_names().is_empty());

        // Cartesian unions the per-clause references.
        let cart = Comprehension::cartesian(vec![bare, call]);
        let got: Vec<String> = cart.referenced_source_names().into_iter().collect();
        assert_eq!(got, vec!["eh_values", "nbo_v_values"]);
    }

    #[test]
    fn cartesian_coordinates_in_declaration_order() {
        let c = Comprehension::cartesian(vec![
            lit_int_clause("k", &[1, 2]),
            lit_int_clause("limit", &[10, 20, 30]),
        ]);
        assert_eq!(c.coordinate_names(), vec!["k", "limit"]);
        assert!(c.is_combinator());
    }

    #[test]
    fn zip_coordinates() {
        let c = Comprehension::zip(
            vec![
                lit_int_clause("x", &[1, 2, 3]),
                lit_int_clause("y", &[10, 20, 30]),
            ],
            ZipMode::Strict,
        );
        assert_eq!(c.coordinate_names(), vec!["x", "y"]);
    }

    #[test]
    fn union_takes_first_childs_shape() {
        let a = Comprehension::cartesian(vec![
            lit_int_clause("k", &[10]),
            lit_int_clause("limit", &[10, 20]),
        ]);
        let b = Comprehension::cartesian(vec![
            lit_int_clause("k", &[100]),
            lit_int_clause("limit", &[100, 200]),
        ]);
        let u = Comprehension::union(vec![a, b]);
        assert_eq!(u.coordinate_names(), vec!["k", "limit"]);
    }

    #[test]
    fn filter_and_order_pass_through_coordinates() {
        let inner = Comprehension::cartesian(vec![
            lit_int_clause("k", &[1, 2]),
            lit_int_clause("limit", &[10]),
        ]);
        let filtered = Comprehension::filter(inner.clone(), "{k} > 0");
        assert_eq!(filtered.coordinate_names(), vec!["k", "limit"]);
        assert!(filtered.is_modifier());

        let ordered = Comprehension::order(inner, StrategyName::Lex, Some(5));
        assert_eq!(ordered.coordinate_names(), vec!["k", "limit"]);
        assert!(ordered.is_modifier());
    }

    #[test]
    fn node_count_and_depth() {
        let inner = Comprehension::cartesian(vec![
            lit_int_clause("k", &[1, 2]),
            lit_int_clause("limit", &[10]),
        ]);
        // inner: 1 (cartesian) + 2 (clauses) = 3 nodes; depth 2
        assert_eq!(inner.node_count(), 3);
        assert_eq!(inner.depth(), 2);

        let filtered = Comprehension::filter(inner, "{k} > 0");
        // filtered: 1 (filter) + 3 = 4 nodes; depth 3
        assert_eq!(filtered.node_count(), 4);
        assert_eq!(filtered.depth(), 3);
    }

    #[test]
    fn round_trip_serde() {
        let c = Comprehension::order(
            Comprehension::filter(
                Comprehension::cartesian(vec![
                    lit_int_clause("k", &[1, 2, 3]),
                    lit_int_clause("limit", &[10, 20]),
                ]),
                "{k} * {limit} > 5",
            ),
            StrategyName::Halton,
            Some(10),
        );
        let json = serde_json::to_string(&c).unwrap();
        let back: Comprehension = serde_json::from_str(&json).unwrap();
        assert_eq!(c, back);
    }

    /// An order's authored seed rides through serde, and an order
    /// written without one reads back as unseeded.
    #[test]
    fn an_orders_seed_round_trips_and_defaults_to_none() {
        let c = Comprehension::order_seeded(
            Comprehension::clause(
                "k",
                Source::IntRange {
                    lo: 1,
                    hi: 4,
                    step: 1,
                },
            ),
            StrategyName::Shuffle,
            Some(2),
            Some(42),
        );
        let json = serde_json::to_string(&c).unwrap();
        assert!(json.contains("\"seed\":42"), "{json}");
        let back: Comprehension = serde_json::from_str(&json).unwrap();
        assert_eq!(back, c);
        let unseeded = json.replace(",\"seed\":42", "");
        let back: Comprehension = serde_json::from_str(&unseeded).unwrap();
        assert!(
            matches!(back, Comprehension::Order { seed: None, .. }),
            "{back:?}"
        );
    }
}