polydat_core/iteration/

source.rs

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

//! Data sources: typed sequences that drive workload iteration.
//!
//! A **source** is a data provider with identity. It knows what it yields
//! (schema), how much it has (extent), where the consumer is (cursor),
//! and how to partition across concurrent fibers.
//!
//! Sources replace the `cycles` counter as the workload iteration driver.
//! A Polydat program declares one with the `cursor` keyword
//! (`cursor q = range(0, N) [over <partition>]`). A host pulls from
//! sources to drive dispatch; when a source is exhausted, the activation
//! is done.
//!
//! ## Source Types
//!
//! - **Range**: `range(0, N)` — a finite sequence of ordinals, served by
//!   [`RangeSourceFactory`]. Replaces `cycles: N`.
//! - **Extending**: `until_elapsed(base, min_ms[, delta])` and the other
//!   `until_*` constructors, compiled to a `CursorKind::Extending*` and
//!   served by [`ExtendingRangeSourceFactory`].
//! - **Host-supplied**: any [`DataSourceFactory`] a host implements, such
//!   as a dataset reader whose items carry vectors or metadata; this
//!   crate ships only the range factories.
//!
//! ## Crate Sovereignty
//!
//! All source API surface lives here in `polydat-core`, re-exported by
//! `polydat` at `polydat::iteration::source`. Host runtimes and adapters
//! consume these types but don't define them.

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::ast::{PortType, Value};

/// Whether a source item can be reconstructed from its ordinal without
/// consulting or advancing mutable source state.
///
/// The default is deliberately conservative. A forward-only cursor is not
/// sufficient evidence that an earlier item is safe to render again.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum SourceReplayStability {
    /// Rendering is destructive, stateful, externally mutable, or otherwise
    /// not proven stable for repeated calls at one ordinal.
    #[default]
    Consumptive,
    /// `render_item(ordinal)` is pure, total, and byte-identical for the
    /// lifetime of the advertised source generation.
    StableByOrdinal,
}

/// A compact description of values which can be generated without rendering
/// an opaque source item first.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum SourceValueForm {
    /// The source may be replayable, but the item must still be rendered.
    #[default]
    Opaque,
    /// The yielded scalar value is the ordinal itself. A typed consumer may
    /// apply its declared narrowing or wrapping conversion while vectorizing.
    Ordinal,
}

/// Runtime source capability used by offset-stamped batch plans.
///
/// `generation` must change whenever rendering the same ordinal may produce a
/// different value. Activation identity remains executor-owned because a
/// phase rewind can intentionally replay one unchanged source generation as a
/// new logical evaluation.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct SourceReplayContract {
    /// Whether rendering an ordinal again gives the same value.
    pub stability: SourceReplayStability,
    /// The generation of the source's values; changes when an ordinal may render differently.
    pub generation: u64,
    /// What a value is: an ordinal, or opaque.
    pub value_form: SourceValueForm,
}

impl SourceReplayContract {
    /// The contract of a source that is consumed as it advances: no replay.
    pub const fn consumptive() -> Self {
        Self {
            stability: SourceReplayStability::Consumptive,
            generation: 0,
            value_form: SourceValueForm::Opaque,
        }
    }

    /// The contract of a source stable by ordinal in `generation`.
    pub const fn stable_ordinal(generation: u64) -> Self {
        Self {
            stability: SourceReplayStability::StableByOrdinal,
            generation,
            value_form: SourceValueForm::Ordinal,
        }
    }

    /// Whether any owned ordinal renders again without advancing the source.
    pub const fn is_stable_by_ordinal(self) -> bool {
        matches!(self.stability, SourceReplayStability::StableByOrdinal)
    }

    /// True when the source ordinal is simultaneously a replay key, scalar
    /// value, packet number, and low-order lane clock.
    pub const fn is_perfect_ordinal(self) -> bool {
        self.is_stable_by_ordinal() && matches!(self.value_form, SourceValueForm::Ordinal)
    }
}

/// A single item yielded by a source.
#[derive(Clone, Debug)]
pub struct SourceItem {
    /// Position in the source sequence.
    pub ordinal: u64,
    /// Named field values. Empty for range sources (ordinal IS the data).
    /// For dataset sources: `[("vector", Value::Json(...)), ("metadata", Value::U64(...))]`.
    pub fields: Vec<(String, Value)>,
}

impl SourceItem {
    /// Create a range item (ordinal only, no fields).
    pub fn ordinal(ordinal: u64) -> Self {
        Self {
            ordinal,
            fields: Vec::new(),
        }
    }

    /// Create an item with ordinal and named fields.
    pub fn with_fields(ordinal: u64, fields: Vec<(String, Value)>) -> Self {
        Self { ordinal, fields }
    }

    /// Get a field value by name.
    pub fn field(&self, name: &str) -> Option<&Value> {
        self.fields.iter().find(|(n, _)| n == name).map(|(_, v)| v)
    }
}

/// Schema describing what a source yields.
#[derive(Clone, Debug)]
pub struct SourceSchema {
    /// Source name as declared in the Polydat graph.
    pub name: String,
    /// Field names and types available for projection (e.g., `ordinal: U64`, `vector: Json`).
    pub projections: Vec<(String, PortType)>,
    /// Known extent, if finite. `None` for infinite sources.
    ///
    /// `None` may also mean the extent is computable but only at runtime
    /// (e.g., the cursor's `range(...)` bounds depend on iteration-variable
    /// externs). In that case `extent_outputs` carries the kernel output
    /// names whose values yield `[start, end)` once externs are bound.
    pub extent: Option<u64>,
    /// Optional aux output names `(start, end)` that the runtime can pull
    /// from the kernel after externs are populated to compute extent. Set
    /// when `range(...)` bounds are non-literal (e.g., wire-bound dataset
    /// function calls). The compiler also writes back to `extent` if both
    /// values fold to constants at compile time.
    pub extent_outputs: Option<(String, String)>,
    /// Optional cursor limit clamp, taken from
    /// `CompileOptions::cursor_limit` (a host's `--limit`-style option).
    /// Applied after runtime extent evaluation.
    pub extent_limit: Option<u64>,
    /// What kind of cursor this is. Default `Range`, the bounded
    /// `range(start, end)` cursor; the `Extending*` variants declare a
    /// runtime-extending cursor under a wall-clock, count, or
    /// predicate policy. A host picks [`ExtendingRangeSourceFactory`]
    /// for those, and `cursor_over_partitions` treats them as
    /// open-extent when it resolves an `over` clause.
    pub cursor_kind: CursorKind,
    /// SRD 71: name of the kernel output that carries the
    /// partition-narrowing source (the `over <expr>` clause on
    /// the cursor declaration). `cursor_over_partitions` pulls it and
    /// resolves it through `cursor_partition::resolve_over` into the
    /// full partition list: a spec string or a `PartitionSpec`
    /// resolves against the cursor's extent, a `Partition` or a
    /// `PartitionList` is re-projected onto it, and `Value::None`
    /// yields no partitions. A host, such as activation setup,
    /// selects one partition and writes it with `set_cursor`.
    ///
    /// `None` means the cursor was declared without an `over`
    /// clause; the cursor uses its full declared extent.
    pub partition_output: Option<String>,
    /// The partitions the `over` clause denotes, resolved by the
    /// compiler when the clause is a literal spec and the extent is
    /// known at build (engines.md §3.5). A host reads them
    /// without evaluating anything; `cursor_over_partitions` returns
    /// them without a pull. `None` when the clause or the extent is
    /// only known at run time, or the cursor has no `over` clause.
    pub partitions: Option<Vec<crate::iteration::cursor_partition::Partition>>,
}

/// Cursor-construction discriminator. Set by the Polydat compiler
/// when it recognises a cursor's constructor expression
/// (`range(...)`, `until_elapsed(...)`, etc.); read by the
/// executor at phase setup to instantiate the matching
/// data-source factory + policy.
///
/// Each `Extending*` variant carries the Polydat output names the
/// runtime should pull at phase setup to resolve the policy's
/// parameters (same pattern as `extent_outputs` for `range`).
/// `delta_output` is optional — when `None`, the cursor's base
/// is used as the extension delta.
#[derive(Clone, Debug, Default)]
pub enum CursorKind {
    /// Bounded `range(start, end)` cursor. Default — keeps
    /// every existing call site working.
    #[default]
    Range,
    /// `until_elapsed(base, min_ms[, delta])`.
    /// Extends while elapsed_ms < min_ms.
    ExtendingTimed {
        /// The output holding the minimum elapsed milliseconds.
        min_ms_output: String,
        /// The output holding the extension step, if one was given.
        delta_output: Option<String>,
    },
    /// `until_passes(base, min_passes[, delta])`.
    /// Extends while completed passes < min_passes.
    ExtendingPasses {
        /// The output holding the minimum number of passes.
        min_passes_output: String,
        /// The output holding the extension step, if one was given.
        delta_output: Option<String>,
    },
    /// `until_count(base, min_count[, delta])`.
    /// Extends while raw consumed count < min_count.
    ExtendingCount {
        /// The output holding the minimum consumed count.
        min_count_output: String,
        /// The output holding the extension step, if one was given.
        delta_output: Option<String>,
    },
    /// `until_elapsed_and_passes(base, min_ms, min_passes[, delta])`.
    /// AND: stops when EITHER target is reached. Extends while
    /// BOTH conditions are still below target.
    ExtendingElapsedAndPasses {
        /// The output holding the minimum elapsed milliseconds.
        min_ms_output: String,
        /// The output holding the minimum number of passes.
        min_passes_output: String,
        /// The output holding the extension step, if one was given.
        delta_output: Option<String>,
    },
    /// `until_elapsed_or_passes(base, min_ms, min_passes[, delta])`.
    /// OR: stops only when BOTH targets are reached. Extends
    /// while EITHER condition is still below target.
    ExtendingElapsedOrPasses {
        /// The output holding the minimum elapsed milliseconds.
        min_ms_output: String,
        /// The output holding the minimum number of passes.
        min_passes_output: String,
        /// The output holding the extension step, if one was given.
        delta_output: Option<String>,
    },
}

/// Consumption API for data sources. One instance per fiber.
///
/// The interaction model has two phases:
///
/// 1. **Reserve** — `reserve(stride)` atomically claims a range of
///    ordinals via CAS on the shared cursor. This touches shared
///    state but is instantaneous (one atomic op). Returns `None`
///    when the source is globally exhausted.
///
/// 2. **Render** — the fiber uses the reserved range with its own
///    Polydat instance to produce field values. No shared state, no
///    contention between fibers. For range sources, rendering is
///    trivial (ordinal IS the data). For dataset sources, rendering
///    reads vectors/metadata from mmap'd storage.
///
/// The `next_chunk` convenience method combines both phases. Use
/// `reserve` directly when the rendering is handled by the
/// executor's Polydat fiber.
pub trait DataSource: Send {
    /// Atomically reserve up to `stride` ordinals from the source.
    ///
    /// Returns the half-open range `[start..end)` of reserved
    /// ordinals, or `None` if the source is exhausted. The range
    /// may be shorter than `stride` at the tail of the source.
    ///
    /// This is the only method that touches shared state (the
    /// global cursor). It must be lock-free — a single CAS or
    /// fetch_add.
    fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>>;

    /// Pull the next item. `None` = source exhausted.
    fn next(&mut self) -> Option<SourceItem> {
        let range = self.reserve(1)?;
        Some(self.render_item(range.start))
    }

    /// Pull up to `limit` items. Combines reserve + render.
    /// Returns fewer than `limit` only when the source is globally
    /// exhausted. Empty vec = exhausted.
    fn next_chunk(&mut self, limit: usize) -> Vec<SourceItem> {
        let range = match self.reserve(limit) {
            Some(r) => r,
            None => return Vec::new(),
        };
        (range.start..range.end)
            .map(|ordinal| self.render_item(ordinal))
            .collect()
    }

    /// Produce a source item for a previously reserved ordinal.
    ///
    /// This is the fiber-local rendering step — no shared state.
    /// For range sources: returns `SourceItem::ordinal(ordinal)`.
    /// For dataset sources: reads vector/metadata from storage.
    fn render_item(&self, ordinal: u64) -> SourceItem;

    /// Known extent, if finite.
    fn extent(&self) -> Option<u64>;

    /// Items consumed so far (for progress reporting).
    fn consumed(&self) -> u64;

    /// The schema of items this source yields.
    fn schema(&self) -> &SourceSchema;

    /// Replay/addressability capability for offset-stamped batch execution.
    /// External source implementations inherit the safe consumptive default
    /// until they explicitly prove the stronger contract.
    fn replay_contract(&self) -> SourceReplayContract {
        SourceReplayContract::consumptive()
    }
}

/// Factory that creates per-fiber `DataSource` readers.
///
/// Holds shared state (atomic cursor, partition pool) that's
/// distributed across readers. Each fiber gets its own reader.
///
/// ## Dispatch model
///
/// The **stride** is the stanza length — the number of source items
/// a fiber acquires as an atomic unit. One stanza of ops processes
/// one stride of source items. Strides are inseparable: a fiber
/// that acquires a stride processes all items before acquiring the
/// next.
///
/// The default implementation (`RangeSourceFactory`) uses a shared
/// atomic cursor — all fibers pull strides from the same counter,
/// producing natural monotonic striping. This is correct for range
/// sources where items are independent ordinals.
///
/// For dataset sources with locality benefits (mmap prefetch),
/// factories can implement partitioned allocation: each fiber gets
/// a pre-assigned range of strides, and when exhausted, steals
/// strides from a shared pool. The stride is the minimum unit of
/// work stealing — a fiber never steals partial stanzas.
pub trait DataSourceFactory: Send + Sync {
    /// Create a new reader for a fiber.
    fn create_reader(&self) -> Box<dyn DataSource>;

    /// Schema for all readers from this factory.
    fn schema(&self) -> &SourceSchema;

    /// Global items consumed across all readers (for progress reporting).
    fn global_consumed(&self) -> u64;

    /// Known extent, if finite. Same as schema().extent but avoids clone.
    fn global_extent(&self) -> Option<u64> {
        self.schema().extent
    }

    /// Replay/addressability capability shared by readers from this factory.
    fn replay_contract(&self) -> SourceReplayContract {
        SourceReplayContract::consumptive()
    }

    /// Rewind the factory's shared cursor to the start so a
    /// subsequent `create_reader()` produces a fresh stream
    /// covering the same ordinal range: what a host that re-runs a
    /// source between poll rounds calls after each round exhausts it.
    ///
    /// Default impl: returns `false` to signal the factory
    /// doesn't support rewinding. Factories that DO support it
    /// ([`RangeSourceFactory`], [`ExtendingRangeSourceFactory`])
    /// override and reset their internal cursor / extent
    /// state. A host that needs to rewind should reject a factory
    /// that returns `false` with a clear diagnostic rather than
    /// complete silently after the first round.
    fn rewind_for_poll(&self) -> bool {
        false
    }
}

// =========================================================================
// RangeSource: finite sequence of ordinals
// =========================================================================

/// Factory for range sources. Shared atomic cursor distributes
/// ordinals across fibers.
pub struct RangeSourceFactory {
    cursor: Arc<AtomicU64>,
    end: u64,
    schema: SourceSchema,
}

impl RangeSourceFactory {
    /// Create a range source from `[start, end)`.
    pub fn new(start: u64, end: u64) -> Self {
        Self {
            cursor: Arc::new(AtomicU64::new(start)),
            end,
            schema: SourceSchema {
                name: "_range".into(),
                projections: vec![("ordinal".into(), PortType::U64)],
                extent: Some(end.saturating_sub(start)),
                extent_outputs: None,
                extent_limit: None,
                cursor_kind: CursorKind::Range,
                partition_output: None,
                partitions: None,
            },
        }
    }

    /// Create a range source with a named schema.
    pub fn named(name: &str, start: u64, end: u64) -> Self {
        let mut factory = Self::new(start, end);
        factory.schema.name = name.to_string();
        factory
    }
}

impl DataSourceFactory for RangeSourceFactory {
    fn create_reader(&self) -> Box<dyn DataSource> {
        Box::new(RangeSource {
            cursor: self.cursor.clone(),
            end: self.end,
            consumed: 0,
            schema: self.schema.clone(),
        })
    }

    fn schema(&self) -> &SourceSchema {
        &self.schema
    }

    fn global_consumed(&self) -> u64 {
        let pos = self.cursor.load(Ordering::Relaxed);
        let start = self.end.saturating_sub(self.schema.extent.unwrap_or(0));
        pos.saturating_sub(start)
            .min(self.schema.extent.unwrap_or(u64::MAX))
    }

    fn replay_contract(&self) -> SourceReplayContract {
        SourceReplayContract::stable_ordinal(0)
    }

    fn rewind_for_poll(&self) -> bool {
        // Reset the shared atomic cursor back to the
        // start-of-range. SRD-75 phase-poll uses this between
        // iterations: each iteration covers ordinals
        // `[start, end)` afresh, so the workload's ops see the
        // same cycle space per iteration. With concurrency=1
        // (mandated by phase-poll), there's a single fiber
        // and no race on the reset.
        let start = self.end.saturating_sub(self.schema.extent.unwrap_or(0));
        self.cursor.store(start, Ordering::Relaxed);
        true
    }
}

/// Per-fiber range reader. Pulls ordinals from a shared atomic cursor.
struct RangeSource {
    cursor: Arc<AtomicU64>,
    end: u64,
    consumed: u64,
    schema: SourceSchema,
}

impl DataSource for RangeSource {
    fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>> {
        let base = self.cursor.fetch_add(stride as u64, Ordering::Relaxed);
        if base >= self.end {
            return None;
        }
        let actual_end = (base + stride as u64).min(self.end);
        let count = actual_end - base;
        self.consumed += count;
        Some(base..actual_end)
    }

    fn render_item(&self, ordinal: u64) -> SourceItem {
        SourceItem::ordinal(ordinal)
    }

    fn extent(&self) -> Option<u64> {
        self.schema.extent
    }

    fn consumed(&self) -> u64 {
        self.consumed
    }

    fn schema(&self) -> &SourceSchema {
        &self.schema
    }

    fn replay_contract(&self) -> SourceReplayContract {
        SourceReplayContract::stable_ordinal(0)
    }
}

// =========================================================================
// ExtendingRangeSource: runtime-growable extent
// =========================================================================

/// Snapshot of cursor state passed to an [`ExtensionPolicy`]
/// at each end-reach decision point. Policies are pure
/// predicates over this context — they don't carry their own
/// clocks or counters.
#[derive(Clone, Copy, Debug)]
pub struct ExtensionContext {
    /// Wall-clock milliseconds since the source factory was
    /// constructed (typically phase start).
    pub elapsed_ms: u64,
    /// Global ordinals consumed so far — `cursor.load() - start`.
    /// Pass count is `consumed / base`.
    pub consumed: u64,
    /// The `base` chunk size declared by the cursor. Used by
    /// policies that reason in passes rather than raw counts.
    pub base: u64,
}

impl ExtensionContext {
    /// Convenience: integer pass count (consumed / base).
    /// Returns 0 when `base == 0` (degenerate cursor).
    pub fn passes(&self) -> u64 {
        self.consumed.checked_div(self.base).unwrap_or(0)
    }
}

/// Policy that decides whether to extend an
/// `ExtendingRangeSource` when its current end is reached.
///
/// Implementations are pure predicates over the
/// [`ExtensionContext`] — no internal state, no side effects.
/// The source provides elapsed time and consumed counts; the
/// policy returns `Some(delta)` to grow the extent or `None`
/// to terminate. Cheap-and-possibly-duplicated call contract
/// (concurrent fibers may invoke under racing end-reach; only
/// one CAS-wins extension takes effect per round).
pub trait ExtensionPolicy: Send + Sync {
    /// Decide how to extend (if at all) given the current
    /// cursor context.
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64>;
}

/// Factory for ExtendingRangeSource. Differs from
/// [`RangeSourceFactory`] in that `end` is an atomic that the
/// extension policy may grow over the lifetime of the phase.
/// `global_extent()` returns the CURRENT end so phase-status
/// displays reflect any growth honestly.
pub struct ExtendingRangeSourceFactory {
    cursor: Arc<AtomicU64>,
    end: Arc<AtomicU64>,
    start: u64,
    /// Per-pass chunk size — also the default extension delta
    /// when the policy reports "continue". Exposed to the
    /// policy via `ExtensionContext::base` so pass-count
    /// predicates work.
    base: u64,
    /// SRD 71 — hard upper bound on growth. When a cursor is
    /// narrowed by a partition (`until_elapsed(...) over p`),
    /// the partition's end ordinal caps the extension: the
    /// policy keeps making its time / pass / count decisions,
    /// but the source terminates the moment the partition is
    /// exhausted, whichever comes first. `None` = unbounded
    /// (the policy alone decides).
    max_end: Option<u64>,
    /// Wall-clock baseline. Captured at factory construction
    /// so per-phase factories yield per-phase elapsed numbers
    /// without external clock plumbing.
    started: std::time::Instant,
    policy: Arc<dyn ExtensionPolicy>,
    schema: SourceSchema,
}

impl ExtendingRangeSourceFactory {
    /// Build with an initial extent `[start, start + initial_extent)`.
    /// The extension policy is consulted only when the cursor
    /// reaches the end — the first stride consumed produces
    /// indices in the initial range.
    pub fn new(
        name: &str,
        start: u64,
        initial_extent: u64,
        policy: Arc<dyn ExtensionPolicy>,
    ) -> Self {
        let end = start.saturating_add(initial_extent);
        Self {
            cursor: Arc::new(AtomicU64::new(start)),
            end: Arc::new(AtomicU64::new(end)),
            start,
            base: initial_extent,
            max_end: None,
            started: std::time::Instant::now(),
            policy,
            schema: SourceSchema {
                name: name.to_string(),
                projections: vec![("ordinal".into(), PortType::U64)],
                extent: Some(initial_extent),
                extent_outputs: None,
                cursor_kind: CursorKind::Range,
                extent_limit: None,
                partition_output: None,
                partitions: None,
            },
        }
    }

    /// Cap growth at `max_end` (absolute ordinal, exclusive) —
    /// the partition-narrowing bound from SRD 71. The initial
    /// extent is clamped too, so a base chunk larger than the
    /// partition never reserves past it.
    pub fn bounded(mut self, max_end: u64) -> Self {
        self.max_end = Some(max_end);
        let clamped = self.end.load(Ordering::Acquire).min(max_end);
        self.end.store(clamped, Ordering::Release);
        self
    }
}

impl DataSourceFactory for ExtendingRangeSourceFactory {
    fn create_reader(&self) -> Box<dyn DataSource> {
        Box::new(ExtendingRangeSource {
            cursor: self.cursor.clone(),
            end: self.end.clone(),
            policy: self.policy.clone(),
            start: self.start,
            base: self.base,
            max_end: self.max_end,
            started: self.started,
            consumed: 0,
            schema: self.schema.clone(),
        })
    }

    fn schema(&self) -> &SourceSchema {
        &self.schema
    }

    fn global_consumed(&self) -> u64 {
        let pos = self.cursor.load(Ordering::Relaxed);
        pos.saturating_sub(self.start)
    }

    fn global_extent(&self) -> Option<u64> {
        // Live extent — readers / status displays see growth.
        Some(self.end.load(Ordering::Acquire).saturating_sub(self.start))
    }

    fn rewind_for_poll(&self) -> bool {
        // A new round starts over: the cursor at the start and the end
        // at the base chunk, which the policy grows again as it did
        // the first time. A partition bound (`max_end`) still caps it.
        self.cursor.store(self.start, Ordering::Release);
        let mut end = self.start.saturating_add(self.base);
        if let Some(cap) = self.max_end {
            end = end.min(cap);
        }
        self.end.store(end, Ordering::Release);
        true
    }

    fn replay_contract(&self) -> SourceReplayContract {
        SourceReplayContract::stable_ordinal(0)
    }
}

/// Per-fiber reader for an extending range source. Reservation
/// uses CAS so concurrent fibers race cleanly without lock
/// contention; on end-reach the policy is consulted (possibly
/// duplicated under concurrent end-reach), and the end atomic
/// is grown via a single CAS (only one extension takes effect
/// per round; losers retry and see the new end).
struct ExtendingRangeSource {
    cursor: Arc<AtomicU64>,
    end: Arc<AtomicU64>,
    policy: Arc<dyn ExtensionPolicy>,
    start: u64,
    base: u64,
    /// SRD 71 partition cap — see
    /// [`ExtendingRangeSourceFactory::bounded`].
    max_end: Option<u64>,
    started: std::time::Instant,
    consumed: u64,
    schema: SourceSchema,
}

impl DataSource for ExtendingRangeSource {
    fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>> {
        loop {
            let cur = self.cursor.load(Ordering::Acquire);
            let end = self.end.load(Ordering::Acquire);
            if cur < end {
                // Try to claim [cur, min(cur+stride, end)).
                let target = (cur.saturating_add(stride as u64)).min(end);
                match self
                    .cursor
                    .compare_exchange(cur, target, Ordering::AcqRel, Ordering::Acquire)
                {
                    Ok(_) => {
                        let count = target - cur;
                        self.consumed += count;
                        return Some(cur..target);
                    }
                    Err(_) => continue, // raced; retry
                }
            }
            // SRD 71: a partition-bound cursor terminates the
            // moment the partition is exhausted, whether or not
            // the policy's time / pass / count target was
            // reached — no policy consultation past the cap.
            if let Some(max) = self.max_end
                && end >= max
            {
                return None;
            }
            // Cursor has caught up to end. Consult policy with
            // a snapshot of the current state. Each fiber that
            // races to this point gets its own context read;
            // duplicate consultations are idempotent for pure
            // predicates.
            let ctx = ExtensionContext {
                elapsed_ms: self.started.elapsed().as_millis() as u64,
                consumed: end.saturating_sub(self.start),
                base: self.base,
            };
            match self.policy.next_extension(&ctx) {
                Some(delta) if delta > 0 => {
                    // CAS the end forward, clamped at the
                    // partition cap when one is set. If someone
                    // else extended ahead of us, that's fine —
                    // our duplicate policy consultation is
                    // harmless and the new end is at least as
                    // far as ours would have been.
                    let mut new_end = end.saturating_add(delta);
                    if let Some(max) = self.max_end {
                        new_end = new_end.min(max);
                    }
                    if new_end == end {
                        return None;
                    }
                    let _ = self.end.compare_exchange(
                        end,
                        new_end,
                        Ordering::AcqRel,
                        Ordering::Acquire,
                    );
                    continue;
                }
                _ => return None,
            }
        }
    }

    fn render_item(&self, ordinal: u64) -> SourceItem {
        SourceItem::ordinal(ordinal)
    }

    fn extent(&self) -> Option<u64> {
        // Live extent — same convention as the factory's
        // `global_extent`: subscribers see the current ceiling
        // even after it grows.
        Some(self.end.load(Ordering::Acquire).saturating_sub(self.start))
    }

    fn consumed(&self) -> u64 {
        self.consumed
    }

    fn schema(&self) -> &SourceSchema {
        &self.schema
    }

    fn replay_contract(&self) -> SourceReplayContract {
        SourceReplayContract::stable_ordinal(0)
    }
}

// =========================================================================
// Cursors: provenance-driven cursor targeting
// =========================================================================

/// A cursor target: a DataSource reader paired with its Polydat input index.
struct CursorTarget {
    /// The DataSource reader that provides values for this cursor.
    reader: Box<dyn DataSource>,
    /// The Polydat input index where the cursor's ordinal is injected.
    input_index: usize,
    /// Source name (for diagnostics).
    source_name: String,
    /// Stable process-independent identity used in ordinal packet stamps.
    stream_id: u64,
}

/// Ownership token for one range reserved from a perfect ordinal source.
///
/// The token contains no reader borrow or rendered payload because the source
/// contract proves that the scalar value is the ordinal itself. Moving this
/// value transfers responsibility for every ordinal in `range`; the batch
/// executor must drain all of them or process the remainder through its scalar
/// recovery path. The shared source cursor is never rewound.
#[derive(Debug)]
pub struct OrdinalBatchLease {
    source_name: String,
    stream_id: u64,
    source_generation: u64,
    input_index: usize,
    sequence: u64,
    range: std::ops::Range<u64>,
}

impl OrdinalBatchLease {
    /// The cursor the lease reserves from.
    pub fn source_name(&self) -> &str {
        &self.source_name
    }

    /// The stream the ordinals belong to.
    pub const fn stream_id(&self) -> u64 {
        self.stream_id
    }

    /// The generation of the source's values.
    pub const fn source_generation(&self) -> u64 {
        self.source_generation
    }

    /// The kernel input the cursor's ordinal is written to.
    pub const fn input_index(&self) -> usize {
        self.input_index
    }

    /// The lease's position in reservation order.
    pub const fn sequence(&self) -> u64 {
        self.sequence
    }

    /// The ordinals reserved.
    pub fn range(&self) -> std::ops::Range<u64> {
        self.range.clone()
    }

    /// How many ordinals the lease holds.
    pub const fn len_u64(&self) -> u64 {
        self.range.end - self.range.start
    }

    /// Whether the lease holds no ordinal.
    pub const fn is_empty(&self) -> bool {
        self.range.start == self.range.end
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Why a cursor batch could not be reserved.
pub enum CursorBatchError {
    /// No ordinals were asked for.
    ZeroDemand,
    /// The reservation names several targets; a batch drives one.
    RequiresSingleTarget {
        /// The targets named.
        targets: usize,
    },
    /// The source is not a perfect ordinal stream.
    SourceNotPerfectOrdinal {
        /// The source's name.
        source: String,
    },
}

impl std::fmt::Display for CursorBatchError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ZeroDemand => f.write_str("ordinal batch demand must be greater than zero"),
            Self::RequiresSingleTarget { targets } => write!(
                f,
                "Tier-1 ordinal batching requires exactly one cursor target, found {targets}"
            ),
            Self::SourceNotPerfectOrdinal { source } => write!(
                f,
                "source '{source}' does not declare the perfect ordinal replay contract"
            ),
        }
    }
}

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

/// Provenance-driven advancer that targets only the cursor nodes
/// relevant to a specific set of output fields.
///
/// Built by a host from the output names it will read, tracing Polydat
/// provenance from those fields back to root cursor nodes. Only those cursors
/// advance — unused cursors are left untouched.
pub struct Cursors {
    targets: Vec<CursorTarget>,
    /// Last items read from each target (for injecting into Polydat state).
    last_items: Vec<Option<SourceItem>>,
    /// Total advances performed.
    advances: u64,
    /// Reservation order for owned ordinal leases from this cursor set.
    next_batch_sequence: u64,
}

impl Cursors {
    /// Build an advancer from a Polydat program, a set of output field names,
    /// and a map of source name → DataSourceFactory.
    ///
    /// Traces provenance: for each field, finds the output node, gets its
    /// input provenance bitmask, and identifies which inputs are cursor
    /// ordinals (matching `{source}__ordinal` pattern). Creates a reader
    /// for each targeted source.
    pub fn for_fields(
        program: &crate::kernel::PolydatProgram,
        field_names: &[&str],
        source_factories: &std::collections::HashMap<String, Arc<dyn DataSourceFactory>>,
    ) -> Self {
        // Collect the union of input provenance for all referenced
        // fields — exact at any input count (multi-word ProvMask;
        // the former one-word form over-matched cursor inputs >= 63).
        let mut combined_provenance = crate::kernel::ProvMask::empty();
        for name in field_names {
            if let Some((node_idx, _)) = program.resolve_output(name)
                && let Some(prov) = program.input_provenance_for(node_idx)
            {
                combined_provenance.union_with(prov);
            }
        }

        // Find cursor inputs in the provenance
        let input_names = program.input_names();
        let mut targets = Vec::new();
        let mut seen_sources = std::collections::HashSet::new();

        for (idx, input_name) in input_names.iter().enumerate() {
            if !combined_provenance.contains(idx) {
                continue;
            }

            // Check if this input is a source projection ({source}__ordinal)
            if let Some(source_name) = input_name.strip_suffix("__ordinal") {
                if seen_sources.contains(source_name) {
                    continue;
                }
                seen_sources.insert(source_name.to_string());

                if let Some(factory) = source_factories.get(source_name) {
                    let stream_id =
                        xxhash_rust::xxh3::xxh3_64(format!("{source_name}@{idx}").as_bytes());
                    targets.push(CursorTarget {
                        reader: factory.create_reader(),
                        input_index: idx,
                        source_name: source_name.to_string(),
                        stream_id,
                    });
                }
            }
        }

        let target_count = targets.len();
        Cursors {
            targets,
            last_items: vec![None; target_count],
            advances: 0,
            next_batch_sequence: 0,
        }
    }

    /// Reserve one owned batch from a single perfect ordinal cursor.
    ///
    /// This is the source-side entry point for Tier-1 SIMD execution. It is
    /// intentionally unavailable for multiple cursor targets or sources that
    /// merely happen to be monotonic: only the explicit perfect-ordinal
    /// contract permits payload-free replay from the returned range.
    pub fn reserve_ordinal_batch(
        &mut self,
        demand: usize,
    ) -> Result<Option<OrdinalBatchLease>, CursorBatchError> {
        if demand == 0 {
            return Err(CursorBatchError::ZeroDemand);
        }
        if self.targets.len() != 1 {
            return Err(CursorBatchError::RequiresSingleTarget {
                targets: self.targets.len(),
            });
        }

        let target = &mut self.targets[0];
        let contract = target.reader.replay_contract();
        if !contract.is_perfect_ordinal() {
            return Err(CursorBatchError::SourceNotPerfectOrdinal {
                source: target.source_name.clone(),
            });
        }
        let Some(range) = target.reader.reserve(demand) else {
            return Ok(None);
        };
        let claimed = range.end - range.start;
        let lease = OrdinalBatchLease {
            source_name: target.source_name.clone(),
            stream_id: target.stream_id,
            source_generation: contract.generation,
            input_index: target.input_index,
            sequence: self.next_batch_sequence,
            range,
        };
        self.next_batch_sequence = self.next_batch_sequence.wrapping_add(1);
        self.advances = self.advances.saturating_add(claimed);
        Ok(Some(lease))
    }

    /// Advance all targeted cursors. Returns `false` if any targeted
    /// cursor is exhausted (no more data).
    ///
    /// After advancing, the new ordinals and field projections are
    /// available via `inject_into_state()`.
    pub fn advance(&mut self) -> bool {
        for (i, target) in self.targets.iter_mut().enumerate() {
            match target.reader.next() {
                Some(item) => {
                    self.last_items[i] = Some(item);
                }
                None => return false, // this cursor exhausted
            }
        }
        self.advances += 1;
        true
    }

    /// Inject the current cursor values into a Polydat state.
    ///
    /// Sets each cursor's ordinal at its input index. Field projections
    /// are not written here; a host reads them from `last_items()`.
    pub fn inject_into_state(&self, state: &mut crate::kernel::PolydatState) {
        for (i, target) in self.targets.iter().enumerate() {
            if let Some(ref item) = self.last_items[i] {
                state.set_input(target.input_index, crate::ast::Value::U64(item.ordinal));
                // Field projections (e.g. base__vector) are the host's to
                // write from `last_items()`; only the ordinal lands here.
            }
        }
    }

    /// The last items read from all targets, for a host to project fields from.
    pub fn last_items(&self) -> &[Option<SourceItem>] {
        &self.last_items
    }

    /// Known extent of the driving cursor (smallest among targeted
    /// cursors with known extent). Used for progress reporting.
    pub fn extent(&self) -> Option<u64> {
        self.targets.iter().filter_map(|t| t.reader.extent()).min()
    }

    /// Total advances performed so far.
    pub fn consumed(&self) -> u64 {
        self.advances
    }

    /// Number of targeted cursors.
    pub fn target_count(&self) -> usize {
        self.targets.len()
    }

    /// Whether this advancer has any targets.
    pub fn is_empty(&self) -> bool {
        self.targets.is_empty()
    }
}

// =========================================================================
// Built-in ExtensionPolicy implementations
// =========================================================================
//
// Each policy is a pure predicate over `ExtensionContext` and
// returns `Some(delta)` when the cursor should continue or
// `None` when it should stop. The delta is independent of the
// stop condition — workloads can extend in chunks unrelated to
// the cursor's base size (e.g. base=10000 but delta=1000 to
// check the condition more often).

/// Extend while `ctx.elapsed_ms < min_ms`, projecting the
/// remaining work from the observed rate and committing it in
/// one batch.
///
/// Per call: estimate `repeats = floor((consumed * remaining_ms
/// / elapsed_ms) / base)` — the number of base-sized passes
/// needed to fill the time remaining at the current rate. Apply
/// a 5% under-bias (`repeats * 95 / 100`) so successive
/// extensions converge from below the target rather than
/// overshooting once. Commit `repeats * base` cycles.
///
/// The under-bias gives a geometric convergence: each call
/// covers ~95% of the time remaining, so 3–4 calls typically
/// land within 0.01% of the time budget. The first call (no
/// rate signal yet — `elapsed_ms == 0` or `consumed == 0`)
/// falls back to one base chunk via the `delta` field.
pub struct UntilElapsedPolicy {
    /// The elapsed milliseconds to reach.
    pub min_ms: u64,
    /// The extension step for the first call, before a rate is known.
    pub delta: u64,
}

impl ExtensionPolicy for UntilElapsedPolicy {
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
        if ctx.elapsed_ms >= self.min_ms {
            return None;
        }
        // No rate signal yet — first end-reach with elapsed time
        // below the clock resolution, or a degenerate empty
        // cursor. Step by the declared `delta` (which the
        // executor sets to `base` when the workload didn't
        // supply an explicit `delta` arg) so the next call has
        // a rate measurement to project from.
        if ctx.elapsed_ms == 0 || ctx.consumed == 0 {
            return Some(self.delta.max(1));
        }
        let remaining_ms = self.min_ms - ctx.elapsed_ms;
        // u128 saturating math: consumed * remaining_ms can
        // overflow u64 for long-running phases with high
        // throughput (e.g. 10^9 ops over 10^4 ms).
        let est_remaining =
            (ctx.consumed as u128).saturating_mul(remaining_ms as u128) / (ctx.elapsed_ms as u128);
        let biased = est_remaining.saturating_mul(95) / 100;
        let base = ctx.base.max(1) as u128;
        let repeats = biased / base;
        if repeats == 0 {
            return None;
        }
        let delta = repeats.saturating_mul(base);
        Some(u64::try_from(delta).unwrap_or(u64::MAX))
    }
}

/// Extend by `delta` while `ctx.passes() < min_passes`. Passes
/// are whole multiples of `ctx.base`.
pub struct UntilPassesPolicy {
    /// The passes to reach.
    pub min_passes: u64,
    /// The extension step.
    pub delta: u64,
}

impl ExtensionPolicy for UntilPassesPolicy {
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
        if ctx.passes() < self.min_passes {
            Some(self.delta)
        } else {
            None
        }
    }
}

/// Extend by `delta` while `ctx.consumed < min_count`.
pub struct UntilCountPolicy {
    /// The consumed count to reach.
    pub min_count: u64,
    /// The extension step.
    pub delta: u64,
}

impl ExtensionPolicy for UntilCountPolicy {
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
        if ctx.consumed < self.min_count {
            Some(self.delta)
        } else {
            None
        }
    }
}

/// Continue (extend) while ALL child policies say continue.
/// Stops the moment any child policy returns `None`. The
/// delta is the minimum of the children's deltas — conservative
/// step size keeps any single condition from over-shooting its
/// stop point.
pub struct AndPolicy {
    /// The policies that must all ask to continue.
    pub policies: Vec<Arc<dyn ExtensionPolicy>>,
}

impl ExtensionPolicy for AndPolicy {
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
        let mut min_delta = u64::MAX;
        for p in &self.policies {
            let d = p.next_extension(ctx)?;
            min_delta = min_delta.min(d);
        }
        if min_delta == u64::MAX || min_delta == 0 {
            None
        } else {
            Some(min_delta)
        }
    }
}

/// Continue (extend) while ANY child policy says continue.
/// Stops only when every child policy returns `None`. The
/// delta is the maximum of the policies that said continue —
/// matches the most aggressive child still pushing forward.
pub struct OrPolicy {
    /// The policies of which any may ask to continue.
    pub policies: Vec<Arc<dyn ExtensionPolicy>>,
}

impl ExtensionPolicy for OrPolicy {
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
        let mut max_delta: Option<u64> = None;
        for p in &self.policies {
            if let Some(d) = p.next_extension(ctx) {
                max_delta = Some(max_delta.map(|m| m.max(d)).unwrap_or(d));
            }
        }
        max_delta
    }
}

/// Back-compat alias for the original time-only policy.
/// Constructs an [`UntilElapsedPolicy`] with delta = base.
pub struct TimeElapsedPolicy {
    inner: UntilElapsedPolicy,
}

impl TimeElapsedPolicy {
    /// A policy extending in steps of `base` until `min_ms` have elapsed.
    pub fn new(base: u64, min_ms: u64) -> Self {
        Self {
            inner: UntilElapsedPolicy {
                min_ms,
                delta: base,
            },
        }
    }
}

impl ExtensionPolicy for TimeElapsedPolicy {
    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
        self.inner.next_extension(ctx)
    }
}

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

    #[test]
    fn range_sources_declare_the_perfect_ordinal_replay_contract() {
        let factory = RangeSourceFactory::new(10, 20);
        assert!(factory.replay_contract().is_perfect_ordinal());
        let reader = factory.create_reader();
        assert!(reader.replay_contract().is_perfect_ordinal());
        let item = reader.render_item(13);
        assert_eq!(item.ordinal, 13);
        assert!(item.fields.is_empty());
    }

    /// Policy that always extends by `base` — stands in for a
    /// time/pass policy whose target is far away, so the
    /// partition cap is the only thing that can stop growth.
    struct AlwaysExtend;
    impl ExtensionPolicy for AlwaysExtend {
        fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
            Some(ctx.base.max(1))
        }
    }

    #[test]
    fn extending_source_bounded_terminates_at_partition_end() {
        // SRD 71: `until_*(base, ...) over p` — base-sized chunks
        // walk within the partition; the cap stops growth even
        // though the policy would keep extending. Partition
        // [100, 125) with base 10 → 10 + 10 + 5, then exhausted.
        let factory =
            ExtendingRangeSourceFactory::new("q", 100, 10, Arc::new(AlwaysExtend)).bounded(125);
        let mut reader = factory.create_reader();
        let mut total = 0u64;
        let mut last_end = 100;
        while let Some(r) = reader.reserve(7) {
            assert!(r.end <= 125, "reservation past the partition cap: {r:?}");
            assert_eq!(r.start, last_end, "contiguous reservations");
            last_end = r.end;
            total += r.end - r.start;
        }
        assert_eq!(total, 25, "exactly the partition's cardinality");
        assert_eq!(last_end, 125);
    }

    #[test]
    fn extending_source_bounded_clamps_oversized_base() {
        // A base chunk larger than the partition never reserves
        // past it.
        let factory =
            ExtendingRangeSourceFactory::new("q", 0, 1000, Arc::new(AlwaysExtend)).bounded(30);
        let mut reader = factory.create_reader();
        let r = reader.reserve(usize::MAX).unwrap();
        assert_eq!(r, 0..30);
        assert!(reader.reserve(1).is_none());
    }

    #[test]
    fn extending_source_unbounded_keeps_policy_semantics() {
        // Without a cap the policy alone decides — three
        // extensions of a terminating policy.
        struct NTimes(std::sync::atomic::AtomicU64);
        impl ExtensionPolicy for NTimes {
            fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
                if self.0.fetch_add(1, Ordering::Relaxed) < 3 {
                    Some(ctx.base)
                } else {
                    None
                }
            }
        }
        let factory = ExtendingRangeSourceFactory::new(
            "q",
            0,
            10,
            Arc::new(NTimes(std::sync::atomic::AtomicU64::new(0))),
        );
        let mut reader = factory.create_reader();
        let mut total = 0u64;
        while let Some(r) = reader.reserve(64) {
            total += r.end - r.start;
        }
        assert_eq!(total, 40, "initial 10 + three 10-ordinal extensions");
    }

    #[test]
    fn range_source_yields_ordinals() {
        let factory = RangeSourceFactory::new(0, 5);
        let mut reader = factory.create_reader();
        assert_eq!(reader.extent(), Some(5));

        for i in 0..5 {
            let item = reader.next().unwrap();
            assert_eq!(item.ordinal, i);
            assert!(item.fields.is_empty());
        }
        assert!(reader.next().is_none());
        assert_eq!(reader.consumed(), 5);
    }

    #[test]
    fn range_source_chunk() {
        let factory = RangeSourceFactory::new(0, 10);
        let mut reader = factory.create_reader();

        let chunk = reader.next_chunk(3);
        assert_eq!(chunk.len(), 3);
        assert_eq!(chunk[0].ordinal, 0);
        assert_eq!(chunk[2].ordinal, 2);

        let chunk = reader.next_chunk(100);
        assert_eq!(chunk.len(), 7); // only 7 remaining
        assert_eq!(chunk[0].ordinal, 3);
        assert_eq!(chunk[6].ordinal, 9);

        let chunk = reader.next_chunk(1);
        assert!(chunk.is_empty()); // exhausted
    }

    #[test]
    fn range_source_concurrent_readers() {
        let factory = RangeSourceFactory::new(0, 100);
        let mut r1 = factory.create_reader();
        let mut r2 = factory.create_reader();

        // Each reader gets unique ordinals from the shared cursor
        let a = r1.next().unwrap().ordinal;
        let b = r2.next().unwrap().ordinal;
        assert_ne!(a, b);

        // Drain both readers
        let mut total = 2;
        while r1.next().is_some() {
            total += 1;
        }
        while r2.next().is_some() {
            total += 1;
        }
        assert_eq!(total, 100);
    }

    #[test]
    fn source_item_field_access() {
        let item = SourceItem::with_fields(
            42,
            vec![
                ("name".into(), Value::Str("test".into())),
                ("score".into(), Value::F64(0.95)),
            ],
        );
        assert_eq!(item.ordinal, 42);
        assert_eq!(item.field("name"), Some(&Value::Str("test".into())));
        assert_eq!(item.field("score"), Some(&Value::F64(0.95)));
        assert_eq!(item.field("missing"), None);
    }

    #[test]
    fn range_source_named() {
        let factory = RangeSourceFactory::named("users", 0, 1000);
        assert_eq!(factory.schema().name, "users");
        assert_eq!(factory.schema().extent, Some(1000));
    }

    // ── ExtendingRangeSource ─────────────────────────────────

    /// Trivial extension policy for tests: lets the caller
    /// decide how many extensions remain. Each call to
    /// `next_extension` decrements the count.
    struct FixedExtensions {
        delta: u64,
        remaining: std::sync::atomic::AtomicU64,
    }

    impl FixedExtensions {
        fn new(delta: u64, times: u64) -> Self {
            Self {
                delta,
                remaining: std::sync::atomic::AtomicU64::new(times),
            }
        }
    }

    impl ExtensionPolicy for FixedExtensions {
        fn next_extension(&self, _ctx: &ExtensionContext) -> Option<u64> {
            let prev = self.remaining.fetch_sub(1, Ordering::Relaxed);
            if prev == 0 || prev > i64::MAX as u64 {
                self.remaining.store(0, Ordering::Relaxed);
                None
            } else {
                Some(self.delta)
            }
        }
    }

    #[test]
    fn extending_source_consumes_initial_extent_then_extends() {
        // Initial extent = 5; policy extends once by 5 more.
        // Total expected: 10 ordinals consumed.
        let policy = Arc::new(FixedExtensions::new(5, 1));
        let factory = ExtendingRangeSourceFactory::new("ext", 0, 5, policy);
        let mut reader = factory.create_reader();
        let mut got: Vec<u64> = Vec::new();
        while let Some(item) = reader.next() {
            got.push(item.ordinal);
            if got.len() > 50 {
                panic!("runaway extension");
            }
        }
        assert_eq!(got, (0..10).collect::<Vec<u64>>());
        assert_eq!(reader.consumed(), 10);
    }

    #[test]
    fn extending_source_zero_extensions_behaves_like_fixed_range() {
        let policy = Arc::new(FixedExtensions::new(0, 0));
        let factory = ExtendingRangeSourceFactory::new("ext", 0, 3, policy);
        let mut reader = factory.create_reader();
        let mut got: Vec<u64> = Vec::new();
        while let Some(item) = reader.next() {
            got.push(item.ordinal);
        }
        assert_eq!(got, vec![0, 1, 2]);
    }

    #[test]
    fn extending_source_global_extent_grows_after_extension() {
        // global_extent must reflect the CURRENT end so phase
        // status displays show growth honestly.
        let policy = Arc::new(FixedExtensions::new(10, 2));
        let factory = ExtendingRangeSourceFactory::new("ext", 0, 5, policy);
        assert_eq!(factory.global_extent(), Some(5));
        let mut reader = factory.create_reader();
        // Drain the first 5 to force an extension.
        for _ in 0..5 {
            reader.next().unwrap();
        }
        // Trigger the extension by attempting one more pull.
        let _ = reader.next().unwrap();
        assert_eq!(
            factory.global_extent(),
            Some(15),
            "extent should grow by the extension delta"
        );
    }

    #[test]
    fn extending_source_chunked_reservation_caps_at_current_end() {
        // A stride request that crosses the current `end` must
        // return a SHORT range up to `end`, not an over-claim.
        // Otherwise the next reservation would jump past the
        // extended range and lose ordinals.
        let policy = Arc::new(FixedExtensions::new(5, 1));
        let factory = ExtendingRangeSourceFactory::new("ext", 0, 3, policy);
        let mut reader = factory.create_reader();
        let range = reader.reserve(10).expect("first reserve");
        // Initial end was 3; the reservation must cap there.
        assert_eq!(range, 0..3);
        // Next reserve triggers the extension and consumes the
        // remainder.
        let range = reader.reserve(10).expect("second reserve");
        assert_eq!(range, 3..8);
    }

    fn ctx_at(elapsed_ms: u64, consumed: u64, base: u64) -> ExtensionContext {
        ExtensionContext {
            elapsed_ms,
            consumed,
            base,
        }
    }

    #[test]
    fn until_elapsed_policy_bootstrap_falls_back_to_delta_when_no_rate_signal() {
        // First end-reach: consumed or elapsed effectively zero —
        // no rate to project from. Policy returns `delta` exactly
        // so the next call has a measurement to work with.
        let policy = UntilElapsedPolicy {
            min_ms: 50,
            delta: 7,
        };
        assert_eq!(policy.next_extension(&ctx_at(0, 0, 0)), Some(7));
        assert_eq!(policy.next_extension(&ctx_at(49, 0, 0)), Some(7));
        assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(7));
    }

    #[test]
    fn until_elapsed_policy_stops_at_or_past_min_ms() {
        let policy = UntilElapsedPolicy {
            min_ms: 50,
            delta: 7,
        };
        assert_eq!(policy.next_extension(&ctx_at(50, 100, 100)), None);
        assert_eq!(policy.next_extension(&ctx_at(1000, 100, 100)), None);
    }

    #[test]
    fn until_elapsed_policy_projects_remaining_cycles_with_under_bias() {
        // base=100, consumed=100 (one pass), elapsed=200ms,
        // min_ms=1000 → remaining=800ms. Rate = 100/200 = 0.5
        // cycles/ms; project 0.5 * 800 = 400 cycles. Under-bias
        // 5% → 380. Round to multiples of base (100) → 3 repeats
        // → 300 cycles.
        let policy = UntilElapsedPolicy {
            min_ms: 1000,
            delta: 100,
        };
        assert_eq!(
            policy.next_extension(&ctx_at(200, 100, 100)),
            Some(300),
            "expected 3-pass batch (380 under-biased, floored to base multiples)",
        );
    }

    #[test]
    fn until_elapsed_policy_returns_none_when_remaining_is_under_one_pass() {
        // base=100, consumed=10000 over 990ms, min_ms=1000 →
        // remaining=10ms. Project 10000/990*10 ≈ 101 cycles.
        // Under-bias → 95. Rounded to base multiples → 0 repeats.
        // Policy terminates rather than under-shooting the budget
        // with a wasted partial pass.
        let policy = UntilElapsedPolicy {
            min_ms: 1000,
            delta: 100,
        };
        assert_eq!(policy.next_extension(&ctx_at(990, 10000, 100)), None);
    }

    #[test]
    fn until_elapsed_policy_converges_geometrically() {
        // Sanity: simulate the extension loop. Each iteration
        // commits ~95% of the remaining time budget; the
        // residual is bounded below by the base-pass rounding,
        // so convergence is one-base-coarse rather than
        // arbitrarily tight.
        let policy = UntilElapsedPolicy {
            min_ms: 1000,
            delta: 10,
        };
        // Pretend the first pass took 10ms (rate = 1 cycle/ms).
        let mut elapsed = 10u64;
        let mut consumed = 10u64;
        let mut iters = 0;
        while let Some(delta) = policy.next_extension(&ctx_at(elapsed, consumed, 10)) {
            iters += 1;
            assert!(iters < 10, "geometric series should converge fast");
            consumed += delta;
            // Same rate: 1 cycle/ms.
            elapsed += delta;
        }
        assert!(elapsed <= 1000, "must under-shoot, got elapsed={elapsed}");
        // Residual ≤ ~5% (under-bias) + 1 base pass (rounding) ≈ 6% of target.
        assert!(
            elapsed >= 940,
            "must come within ~6% of target, got {elapsed}"
        );
    }

    #[test]
    fn until_passes_policy_counts_in_base_multiples() {
        let policy = UntilPassesPolicy {
            min_passes: 3,
            delta: 100,
        };
        // 0 passes done — extend.
        assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(100));
        // 2 passes done (200 consumed @ base=100) — extend.
        assert_eq!(policy.next_extension(&ctx_at(0, 200, 100)), Some(100));
        // 3 passes done — stop.
        assert_eq!(policy.next_extension(&ctx_at(0, 300, 100)), None);
        // 4 passes done — stop.
        assert_eq!(policy.next_extension(&ctx_at(0, 400, 100)), None);
    }

    #[test]
    fn until_count_policy_uses_raw_consumed() {
        let policy = UntilCountPolicy {
            min_count: 250,
            delta: 50,
        };
        assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(50));
        assert_eq!(policy.next_extension(&ctx_at(0, 249, 100)), Some(50));
        assert_eq!(policy.next_extension(&ctx_at(0, 250, 100)), None);
    }

    #[test]
    fn and_policy_stops_when_any_child_stops() {
        // time<5000 AND passes<3.
        let time = Arc::new(UntilElapsedPolicy {
            min_ms: 5000,
            delta: 10,
        });
        let passes = Arc::new(UntilPassesPolicy {
            min_passes: 3,
            delta: 20,
        });
        let and = AndPolicy {
            policies: vec![time, passes],
        };
        // Both still want to continue: delta = min(10, 20) = 10.
        assert_eq!(and.next_extension(&ctx_at(0, 0, 100)), Some(10));
        // Time done — stop.
        assert_eq!(and.next_extension(&ctx_at(5000, 0, 100)), None);
        // Passes done — stop.
        assert_eq!(and.next_extension(&ctx_at(0, 300, 100)), None);
    }

    #[test]
    fn or_policy_continues_if_any_child_continues() {
        let time = Arc::new(UntilElapsedPolicy {
            min_ms: 5000,
            delta: 10,
        });
        let passes = Arc::new(UntilPassesPolicy {
            min_passes: 3,
            delta: 20,
        });
        let or = OrPolicy {
            policies: vec![time, passes],
        };
        // Both want to continue → delta = max(10, 20) = 20.
        assert_eq!(or.next_extension(&ctx_at(0, 0, 100)), Some(20));
        // Time done, passes still wants → delta = 20.
        assert_eq!(or.next_extension(&ctx_at(5000, 0, 100)), Some(20));
        // Passes done, time still wants → delta = 10.
        assert_eq!(or.next_extension(&ctx_at(0, 300, 100)), Some(10));
        // Both done → stop.
        assert_eq!(or.next_extension(&ctx_at(5000, 300, 100)), None);
    }

    #[test]
    fn time_elapsed_policy_compat_alias_still_works() {
        let p = TimeElapsedPolicy::new(7, 1);
        std::thread::sleep(std::time::Duration::from_millis(20));
        assert_eq!(p.next_extension(&ctx_at(20, 0, 0)), None);
    }
}

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

    fn drain(reader: &mut dyn DataSource) -> Vec<u64> {
        let mut out = Vec::new();
        while let Some(item) = reader.next() {
            out.push(item.ordinal);
        }
        out
    }

    /// A factory that rewinds hands out the same ordinal range again:
    /// what a host that re-runs a source between rounds relies on.
    #[test]
    fn a_range_factory_rewinds_to_the_same_ordinals() {
        let factory = RangeSourceFactory::new(3, 6);
        assert_eq!(drain(factory.create_reader().as_mut()), vec![3, 4, 5]);
        assert!(factory.create_reader().next().is_none());
        assert!(factory.rewind_for_poll());
        assert_eq!(drain(factory.create_reader().as_mut()), vec![3, 4, 5]);
    }

    /// The extending factory rewinds too: the cursor to the start and
    /// the extent to the base chunk, capped by a partition bound.
    #[test]
    fn an_extending_factory_rewinds_to_its_base_chunk() {
        struct Never;
        impl ExtensionPolicy for Never {
            fn next_extension(&self, _ctx: &ExtensionContext) -> Option<u64> {
                None
            }
        }
        let factory = ExtendingRangeSourceFactory::new("rows", 10, 4, Arc::new(Never));
        assert_eq!(
            drain(factory.create_reader().as_mut()),
            vec![10, 11, 12, 13]
        );
        assert_eq!(factory.global_consumed(), 4);
        assert!(factory.rewind_for_poll());
        assert_eq!(factory.global_consumed(), 0);
        assert_eq!(factory.global_extent(), Some(4));
        assert_eq!(
            drain(factory.create_reader().as_mut()),
            vec![10, 11, 12, 13]
        );
        let bounded = ExtendingRangeSourceFactory::new("rows", 0, 8, Arc::new(Never)).bounded(5);
        assert_eq!(drain(bounded.create_reader().as_mut()), vec![0, 1, 2, 3, 4]);
        assert!(bounded.rewind_for_poll());
        assert_eq!(drain(bounded.create_reader().as_mut()), vec![0, 1, 2, 3, 4]);
    }
}