Skip to main content

polydat_core/iteration/
source.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Data sources: typed sequences that drive workload iteration.
5//!
6//! A **source** is a data provider with identity. It knows what it yields
7//! (schema), how much it has (extent), where the consumer is (cursor),
8//! and how to partition across concurrent fibers.
9//!
10//! Sources replace the `cycles` counter as the workload iteration driver.
11//! The Polydat graph declares sources via the `source` keyword. The runtime
12//! pulls from sources to drive op dispatch. When a source is exhausted,
13//! the phase is done.
14//!
15//! ## Source Types
16//!
17//! - **Range**: `range(0, N)` — finite sequence of ordinals. Replaces `cycles: N`.
18//! - **Dataset**: `dataset_source("example:label_00", "base")` — vectors, queries, etc.
19//! - **Derived**: any Polydat binding promoted to a source via the `source` keyword.
20//!
21//! ## Crate Sovereignty
22//!
23//! All source API surface lives here in polydat. The runtime crates
24//! (host runtime, adapters) consume these types but don't define them.
25
26use std::sync::Arc;
27use std::sync::atomic::{AtomicU64, Ordering};
28
29use crate::ast::{PortType, Value};
30
31/// Whether a source item can be reconstructed from its ordinal without
32/// consulting or advancing mutable source state.
33///
34/// The default is deliberately conservative. A forward-only cursor is not
35/// sufficient evidence that an earlier item is safe to render again.
36#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
37pub enum SourceReplayStability {
38    /// Rendering is destructive, stateful, externally mutable, or otherwise
39    /// not proven stable for repeated calls at one ordinal.
40    #[default]
41    Consumptive,
42    /// `render_item(ordinal)` is pure, total, and byte-identical for the
43    /// lifetime of the advertised source generation.
44    StableByOrdinal,
45}
46
47/// A compact description of values which can be generated without rendering
48/// an opaque source item first.
49#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
50pub enum SourceValueForm {
51    /// The source may be replayable, but the item must still be rendered.
52    #[default]
53    Opaque,
54    /// The yielded scalar value is the ordinal itself. A typed consumer may
55    /// apply its declared narrowing or wrapping conversion while vectorizing.
56    Ordinal,
57}
58
59/// Runtime source capability used by offset-stamped batch plans.
60///
61/// `generation` must change whenever rendering the same ordinal may produce a
62/// different value. Activation identity remains executor-owned because a
63/// phase rewind can intentionally replay one unchanged source generation as a
64/// new logical evaluation.
65#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
66pub struct SourceReplayContract {
67    /// Whether rendering an ordinal again gives the same value.
68    pub stability: SourceReplayStability,
69    /// The generation of the source's values; changes when an ordinal may render differently.
70    pub generation: u64,
71    /// What a value is: an ordinal, or opaque.
72    pub value_form: SourceValueForm,
73}
74
75impl SourceReplayContract {
76    /// The contract of a source that is consumed as it advances: no replay.
77    pub const fn consumptive() -> Self {
78        Self {
79            stability: SourceReplayStability::Consumptive,
80            generation: 0,
81            value_form: SourceValueForm::Opaque,
82        }
83    }
84
85    /// The contract of a source stable by ordinal in `generation`.
86    pub const fn stable_ordinal(generation: u64) -> Self {
87        Self {
88            stability: SourceReplayStability::StableByOrdinal,
89            generation,
90            value_form: SourceValueForm::Ordinal,
91        }
92    }
93
94    /// Whether any owned ordinal renders again without advancing the source.
95    pub const fn is_stable_by_ordinal(self) -> bool {
96        matches!(self.stability, SourceReplayStability::StableByOrdinal)
97    }
98
99    /// True when the source ordinal is simultaneously a replay key, scalar
100    /// value, packet number, and low-order lane clock.
101    pub const fn is_perfect_ordinal(self) -> bool {
102        self.is_stable_by_ordinal() && matches!(self.value_form, SourceValueForm::Ordinal)
103    }
104}
105
106/// A single item yielded by a source.
107#[derive(Clone, Debug)]
108pub struct SourceItem {
109    /// Position in the source sequence.
110    pub ordinal: u64,
111    /// Named field values. Empty for range sources (ordinal IS the data).
112    /// For dataset sources: `[("vector", Value::Json(...)), ("metadata", Value::U64(...))]`.
113    pub fields: Vec<(String, Value)>,
114}
115
116impl SourceItem {
117    /// Create a range item (ordinal only, no fields).
118    pub fn ordinal(ordinal: u64) -> Self {
119        Self {
120            ordinal,
121            fields: Vec::new(),
122        }
123    }
124
125    /// Create an item with ordinal and named fields.
126    pub fn with_fields(ordinal: u64, fields: Vec<(String, Value)>) -> Self {
127        Self { ordinal, fields }
128    }
129
130    /// Get a field value by name.
131    pub fn field(&self, name: &str) -> Option<&Value> {
132        self.fields.iter().find(|(n, _)| n == name).map(|(_, v)| v)
133    }
134}
135
136/// Schema describing what a source yields.
137#[derive(Clone, Debug)]
138pub struct SourceSchema {
139    /// Source name as declared in the Polydat graph.
140    pub name: String,
141    /// Field names and types available for projection (e.g., `ordinal: U64`, `vector: Json`).
142    pub projections: Vec<(String, PortType)>,
143    /// Known extent, if finite. `None` for infinite sources.
144    ///
145    /// `None` may also mean the extent is computable but only at runtime
146    /// (e.g., the cursor's `range(...)` bounds depend on iteration-variable
147    /// externs). In that case `extent_outputs` carries the kernel output
148    /// names whose values yield `[start, end)` once externs are bound.
149    pub extent: Option<u64>,
150    /// Optional aux output names `(start, end)` that the runtime can pull
151    /// from the kernel after externs are populated to compute extent. Set
152    /// when `range(...)` bounds are non-literal (e.g., wire-bound dataset
153    /// function calls). The compiler also writes back to `extent` if both
154    /// values fold to constants at compile time.
155    pub extent_outputs: Option<(String, String)>,
156    /// Optional cursor limit clamp (from `--limit` / `limit=` param).
157    /// Applied after runtime extent evaluation.
158    pub extent_limit: Option<u64>,
159    /// What kind of cursor this is. Default `Range` (the legacy
160    /// bounded-finite cursor); `ExtendingTimed { ... }` declares
161    /// the cursor as runtime-extending under a wall-clock
162    /// minimum-duration policy. The executor branches on this
163    /// to pick between `RangeSourceFactory` and
164    /// `ExtendingRangeSourceFactory` at phase setup.
165    pub cursor_kind: CursorKind,
166    /// SRD 71: name of the kernel output that carries the
167    /// partition-narrowing source (the `over <expr>` clause on
168    /// the cursor declaration). The executor pulls this output
169    /// at phase setup and applies it to the source factory's
170    /// `[start, end)` range. The output value's type drives
171    /// the resolution:
172    /// - `Value::Str(s)`         — parse `s` as a partition spec, resolve, use partition 0.
173    /// - `Value::Ext(Partition)` — use the resolved partition's `start_ord` / `end_ord`.
174    /// - `Value::Ext(PartitionSpec)` — resolve against the cursor's extent, use partition 0.
175    /// - `Value::Ext(PartitionList)` — use partition 0.
176    /// - `Value::None`           — no narrowing (cursor uses its full extent).
177    ///
178    /// `None` means the cursor was declared without an `over`
179    /// clause; the cursor uses its full declared extent.
180    pub partition_output: Option<String>,
181    /// The partitions the `over` clause denotes, resolved by the
182    /// compiler when the clause is a literal spec and the extent is
183    /// known at build (engine_parity.md, step 3). A host reads them
184    /// without evaluating anything; `cursor_over_partitions` returns
185    /// them without a pull. `None` when the clause or the extent is
186    /// only known at run time, or the cursor has no `over` clause.
187    pub partitions: Option<Vec<crate::iteration::cursor_partition::Partition>>,
188}
189
190/// Cursor-construction discriminator. Set by the Polydat compiler
191/// when it recognises a cursor's constructor expression
192/// (`range(...)`, `until_elapsed(...)`, etc.); read by the
193/// executor at phase setup to instantiate the matching
194/// data-source factory + policy.
195///
196/// Each `Extending*` variant carries the Polydat output names the
197/// runtime should pull at phase setup to resolve the policy's
198/// parameters (same pattern as `extent_outputs` for `range`).
199/// `delta_output` is optional — when `None`, the cursor's base
200/// is used as the extension delta.
201#[derive(Clone, Debug, Default)]
202pub enum CursorKind {
203    /// Bounded `range(start, end)` cursor. Default — keeps
204    /// every existing call site working.
205    #[default]
206    Range,
207    /// `until_elapsed(base, min_ms[, delta])`.
208    /// Extends while elapsed_ms < min_ms.
209    ExtendingTimed {
210        /// The output holding the minimum elapsed milliseconds.
211        min_ms_output: String,
212        /// The output holding the extension step, if one was given.
213        delta_output: Option<String>,
214    },
215    /// `until_passes(base, min_passes[, delta])`.
216    /// Extends while completed passes < min_passes.
217    ExtendingPasses {
218        /// The output holding the minimum number of passes.
219        min_passes_output: String,
220        /// The output holding the extension step, if one was given.
221        delta_output: Option<String>,
222    },
223    /// `until_count(base, min_count[, delta])`.
224    /// Extends while raw consumed count < min_count.
225    ExtendingCount {
226        /// The output holding the minimum consumed count.
227        min_count_output: String,
228        /// The output holding the extension step, if one was given.
229        delta_output: Option<String>,
230    },
231    /// `until_elapsed_and_passes(base, min_ms, min_passes[, delta])`.
232    /// AND: stops when EITHER target is reached. Extends while
233    /// BOTH conditions are still below target.
234    ExtendingElapsedAndPasses {
235        /// The output holding the minimum elapsed milliseconds.
236        min_ms_output: String,
237        /// The output holding the minimum number of passes.
238        min_passes_output: String,
239        /// The output holding the extension step, if one was given.
240        delta_output: Option<String>,
241    },
242    /// `until_elapsed_or_passes(base, min_ms, min_passes[, delta])`.
243    /// OR: stops only when BOTH targets are reached. Extends
244    /// while EITHER condition is still below target.
245    ExtendingElapsedOrPasses {
246        /// The output holding the minimum elapsed milliseconds.
247        min_ms_output: String,
248        /// The output holding the minimum number of passes.
249        min_passes_output: String,
250        /// The output holding the extension step, if one was given.
251        delta_output: Option<String>,
252    },
253}
254
255/// Consumption API for data sources. One instance per fiber.
256///
257/// The interaction model has two phases:
258///
259/// 1. **Reserve** — `reserve(stride)` atomically claims a range of
260///    ordinals via CAS on the shared cursor. This touches shared
261///    state but is instantaneous (one atomic op). Returns `None`
262///    when the source is globally exhausted.
263///
264/// 2. **Render** — the fiber uses the reserved range with its own
265///    Polydat instance to produce field values. No shared state, no
266///    contention between fibers. For range sources, rendering is
267///    trivial (ordinal IS the data). For dataset sources, rendering
268///    reads vectors/metadata from mmap'd storage.
269///
270/// The `next_chunk` convenience method combines both phases. Use
271/// `reserve` directly when the rendering is handled by the
272/// executor's Polydat fiber.
273pub trait DataSource: Send {
274    /// Atomically reserve up to `stride` ordinals from the source.
275    ///
276    /// Returns the half-open range `[start..end)` of reserved
277    /// ordinals, or `None` if the source is exhausted. The range
278    /// may be shorter than `stride` at the tail of the source.
279    ///
280    /// This is the only method that touches shared state (the
281    /// global cursor). It must be lock-free — a single CAS or
282    /// fetch_add.
283    fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>>;
284
285    /// Pull the next item. `None` = source exhausted.
286    fn next(&mut self) -> Option<SourceItem> {
287        let range = self.reserve(1)?;
288        Some(self.render_item(range.start))
289    }
290
291    /// Pull up to `limit` items. Combines reserve + render.
292    /// Returns fewer than `limit` only when the source is globally
293    /// exhausted. Empty vec = exhausted.
294    fn next_chunk(&mut self, limit: usize) -> Vec<SourceItem> {
295        let range = match self.reserve(limit) {
296            Some(r) => r,
297            None => return Vec::new(),
298        };
299        (range.start..range.end)
300            .map(|ordinal| self.render_item(ordinal))
301            .collect()
302    }
303
304    /// Produce a source item for a previously reserved ordinal.
305    ///
306    /// This is the fiber-local rendering step — no shared state.
307    /// For range sources: returns `SourceItem::ordinal(ordinal)`.
308    /// For dataset sources: reads vector/metadata from storage.
309    fn render_item(&self, ordinal: u64) -> SourceItem;
310
311    /// Known extent, if finite.
312    fn extent(&self) -> Option<u64>;
313
314    /// Items consumed so far (for progress reporting).
315    fn consumed(&self) -> u64;
316
317    /// The schema of items this source yields.
318    fn schema(&self) -> &SourceSchema;
319
320    /// Replay/addressability capability for offset-stamped batch execution.
321    /// External source implementations inherit the safe consumptive default
322    /// until they explicitly prove the stronger contract.
323    fn replay_contract(&self) -> SourceReplayContract {
324        SourceReplayContract::consumptive()
325    }
326}
327
328/// Factory that creates per-fiber `DataSource` readers.
329///
330/// Holds shared state (atomic cursor, partition pool) that's
331/// distributed across readers. Each fiber gets its own reader.
332///
333/// ## Dispatch model
334///
335/// The **stride** is the stanza length — the number of source items
336/// a fiber acquires as an atomic unit. One stanza of ops processes
337/// one stride of source items. Strides are inseparable: a fiber
338/// that acquires a stride processes all items before acquiring the
339/// next.
340///
341/// The default implementation (`RangeSourceFactory`) uses a shared
342/// atomic cursor — all fibers pull strides from the same counter,
343/// producing natural monotonic striping. This is correct for range
344/// sources where items are independent ordinals.
345///
346/// For dataset sources with locality benefits (mmap prefetch),
347/// factories can implement partitioned allocation: each fiber gets
348/// a pre-assigned range of strides, and when exhausted, steals
349/// strides from a shared pool. The stride is the minimum unit of
350/// work stealing — a fiber never steals partial stanzas.
351pub trait DataSourceFactory: Send + Sync {
352    /// Create a new reader for a fiber.
353    fn create_reader(&self) -> Box<dyn DataSource>;
354
355    /// Schema for all readers from this factory.
356    fn schema(&self) -> &SourceSchema;
357
358    /// Global items consumed across all readers (for progress reporting).
359    fn global_consumed(&self) -> u64;
360
361    /// Known extent, if finite. Same as schema().extent but avoids clone.
362    fn global_extent(&self) -> Option<u64> {
363        self.schema().extent
364    }
365
366    /// Replay/addressability capability shared by readers from this factory.
367    fn replay_contract(&self) -> SourceReplayContract {
368        SourceReplayContract::consumptive()
369    }
370
371    /// Rewind the factory's shared cursor to the start so a
372    /// subsequent `create_reader()` produces a fresh stream
373    /// covering the same ordinal range. SRD-75 phase-poll uses
374    /// this between iterations: after each poll round
375    /// exhausts the source, the activity calls this and
376    /// re-creates the reader to drive another round of cycles
377    /// before the predicate is re-checked.
378    ///
379    /// Default impl: returns `false` to signal the factory
380    /// doesn't support rewinding. Factories that DO support it
381    /// (RangeSourceFactory, ExtendingRangeSourceFactory)
382    /// override and reset their internal cursor / extent
383    /// state. Phase-poll synthesis rejects workloads whose
384    /// source factory returns `false` so the operator gets a
385    /// clear diagnostic instead of a silent first-iteration
386    /// completion.
387    fn rewind_for_poll(&self) -> bool {
388        false
389    }
390}
391
392// =========================================================================
393// RangeSource: finite sequence of ordinals
394// =========================================================================
395
396/// Factory for range sources. Shared atomic cursor distributes
397/// ordinals across fibers. Replaces `CycleSource`.
398pub struct RangeSourceFactory {
399    cursor: Arc<AtomicU64>,
400    end: u64,
401    schema: SourceSchema,
402}
403
404impl RangeSourceFactory {
405    /// Create a range source from `[start, end)`.
406    pub fn new(start: u64, end: u64) -> Self {
407        Self {
408            cursor: Arc::new(AtomicU64::new(start)),
409            end,
410            schema: SourceSchema {
411                name: "_range".into(),
412                projections: vec![("ordinal".into(), PortType::U64)],
413                extent: Some(end.saturating_sub(start)),
414                extent_outputs: None,
415                extent_limit: None,
416                cursor_kind: CursorKind::Range,
417                partition_output: None,
418                partitions: None,
419            },
420        }
421    }
422
423    /// Create a range source with a named schema.
424    pub fn named(name: &str, start: u64, end: u64) -> Self {
425        let mut factory = Self::new(start, end);
426        factory.schema.name = name.to_string();
427        factory
428    }
429}
430
431impl DataSourceFactory for RangeSourceFactory {
432    fn create_reader(&self) -> Box<dyn DataSource> {
433        Box::new(RangeSource {
434            cursor: self.cursor.clone(),
435            end: self.end,
436            consumed: 0,
437            schema: self.schema.clone(),
438        })
439    }
440
441    fn schema(&self) -> &SourceSchema {
442        &self.schema
443    }
444
445    fn global_consumed(&self) -> u64 {
446        let pos = self.cursor.load(Ordering::Relaxed);
447        let start = self.end.saturating_sub(self.schema.extent.unwrap_or(0));
448        pos.saturating_sub(start)
449            .min(self.schema.extent.unwrap_or(u64::MAX))
450    }
451
452    fn replay_contract(&self) -> SourceReplayContract {
453        SourceReplayContract::stable_ordinal(0)
454    }
455
456    fn rewind_for_poll(&self) -> bool {
457        // Reset the shared atomic cursor back to the
458        // start-of-range. SRD-75 phase-poll uses this between
459        // iterations: each iteration covers ordinals
460        // `[start, end)` afresh, so the workload's ops see the
461        // same cycle space per iteration. With concurrency=1
462        // (mandated by phase-poll), there's a single fiber
463        // and no race on the reset.
464        let start = self.end.saturating_sub(self.schema.extent.unwrap_or(0));
465        self.cursor.store(start, Ordering::Relaxed);
466        true
467    }
468}
469
470/// Per-fiber range reader. Pulls ordinals from a shared atomic cursor.
471struct RangeSource {
472    cursor: Arc<AtomicU64>,
473    end: u64,
474    consumed: u64,
475    schema: SourceSchema,
476}
477
478impl DataSource for RangeSource {
479    fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>> {
480        let base = self.cursor.fetch_add(stride as u64, Ordering::Relaxed);
481        if base >= self.end {
482            return None;
483        }
484        let actual_end = (base + stride as u64).min(self.end);
485        let count = actual_end - base;
486        self.consumed += count;
487        Some(base..actual_end)
488    }
489
490    fn render_item(&self, ordinal: u64) -> SourceItem {
491        SourceItem::ordinal(ordinal)
492    }
493
494    fn extent(&self) -> Option<u64> {
495        self.schema.extent
496    }
497
498    fn consumed(&self) -> u64 {
499        self.consumed
500    }
501
502    fn schema(&self) -> &SourceSchema {
503        &self.schema
504    }
505
506    fn replay_contract(&self) -> SourceReplayContract {
507        SourceReplayContract::stable_ordinal(0)
508    }
509}
510
511// =========================================================================
512// ExtendingRangeSource: runtime-growable extent
513// =========================================================================
514
515/// Snapshot of cursor state passed to an [`ExtensionPolicy`]
516/// at each end-reach decision point. Policies are pure
517/// predicates over this context — they don't carry their own
518/// clocks or counters.
519#[derive(Clone, Copy, Debug)]
520pub struct ExtensionContext {
521    /// Wall-clock milliseconds since the source factory was
522    /// constructed (typically phase start).
523    pub elapsed_ms: u64,
524    /// Global ordinals consumed so far — `cursor.load() - start`.
525    /// Pass count is `consumed / base`.
526    pub consumed: u64,
527    /// The `base` chunk size declared by the cursor. Used by
528    /// policies that reason in passes rather than raw counts.
529    pub base: u64,
530}
531
532impl ExtensionContext {
533    /// Convenience: integer pass count (consumed / base).
534    /// Returns 0 when `base == 0` (degenerate cursor).
535    pub fn passes(&self) -> u64 {
536        self.consumed.checked_div(self.base).unwrap_or(0)
537    }
538}
539
540/// Policy that decides whether to extend an
541/// `ExtendingRangeSource` when its current end is reached.
542///
543/// Implementations are pure predicates over the
544/// [`ExtensionContext`] — no internal state, no side effects.
545/// The source provides elapsed time and consumed counts; the
546/// policy returns `Some(delta)` to grow the extent or `None`
547/// to terminate. Cheap-and-possibly-duplicated call contract
548/// (concurrent fibers may invoke under racing end-reach; only
549/// one CAS-wins extension takes effect per round).
550pub trait ExtensionPolicy: Send + Sync {
551    /// Decide how to extend (if at all) given the current
552    /// cursor context.
553    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64>;
554}
555
556/// Factory for ExtendingRangeSource. Differs from
557/// [`RangeSourceFactory`] in that `end` is an atomic that the
558/// extension policy may grow over the lifetime of the phase.
559/// `global_extent()` returns the CURRENT end so phase-status
560/// displays reflect any growth honestly.
561pub struct ExtendingRangeSourceFactory {
562    cursor: Arc<AtomicU64>,
563    end: Arc<AtomicU64>,
564    start: u64,
565    /// Per-pass chunk size — also the default extension delta
566    /// when the policy reports "continue". Exposed to the
567    /// policy via `ExtensionContext::base` so pass-count
568    /// predicates work.
569    base: u64,
570    /// SRD 71 — hard upper bound on growth. When a cursor is
571    /// narrowed by a partition (`until_elapsed(...) over p`),
572    /// the partition's end ordinal caps the extension: the
573    /// policy keeps making its time / pass / count decisions,
574    /// but the source terminates the moment the partition is
575    /// exhausted, whichever comes first. `None` = unbounded
576    /// (the policy alone decides).
577    max_end: Option<u64>,
578    /// Wall-clock baseline. Captured at factory construction
579    /// so per-phase factories yield per-phase elapsed numbers
580    /// without external clock plumbing.
581    started: std::time::Instant,
582    policy: Arc<dyn ExtensionPolicy>,
583    schema: SourceSchema,
584}
585
586impl ExtendingRangeSourceFactory {
587    /// Build with an initial extent `[start, start + initial_extent)`.
588    /// The extension policy is consulted only when the cursor
589    /// reaches the end — the first stride consumed produces
590    /// indices in the initial range.
591    pub fn new(
592        name: &str,
593        start: u64,
594        initial_extent: u64,
595        policy: Arc<dyn ExtensionPolicy>,
596    ) -> Self {
597        let end = start.saturating_add(initial_extent);
598        Self {
599            cursor: Arc::new(AtomicU64::new(start)),
600            end: Arc::new(AtomicU64::new(end)),
601            start,
602            base: initial_extent,
603            max_end: None,
604            started: std::time::Instant::now(),
605            policy,
606            schema: SourceSchema {
607                name: name.to_string(),
608                projections: vec![("ordinal".into(), PortType::U64)],
609                extent: Some(initial_extent),
610                extent_outputs: None,
611                cursor_kind: CursorKind::Range,
612                extent_limit: None,
613                partition_output: None,
614                partitions: None,
615            },
616        }
617    }
618
619    /// Cap growth at `max_end` (absolute ordinal, exclusive) —
620    /// the partition-narrowing bound from SRD 71. The initial
621    /// extent is clamped too, so a base chunk larger than the
622    /// partition never reserves past it.
623    pub fn bounded(mut self, max_end: u64) -> Self {
624        self.max_end = Some(max_end);
625        let clamped = self.end.load(Ordering::Acquire).min(max_end);
626        self.end.store(clamped, Ordering::Release);
627        self
628    }
629}
630
631impl DataSourceFactory for ExtendingRangeSourceFactory {
632    fn create_reader(&self) -> Box<dyn DataSource> {
633        Box::new(ExtendingRangeSource {
634            cursor: self.cursor.clone(),
635            end: self.end.clone(),
636            policy: self.policy.clone(),
637            start: self.start,
638            base: self.base,
639            max_end: self.max_end,
640            started: self.started,
641            consumed: 0,
642            schema: self.schema.clone(),
643        })
644    }
645
646    fn schema(&self) -> &SourceSchema {
647        &self.schema
648    }
649
650    fn global_consumed(&self) -> u64 {
651        let pos = self.cursor.load(Ordering::Relaxed);
652        pos.saturating_sub(self.start)
653    }
654
655    fn global_extent(&self) -> Option<u64> {
656        // Live extent — readers / status displays see growth.
657        Some(self.end.load(Ordering::Acquire).saturating_sub(self.start))
658    }
659
660    fn replay_contract(&self) -> SourceReplayContract {
661        SourceReplayContract::stable_ordinal(0)
662    }
663}
664
665/// Per-fiber reader for an extending range source. Reservation
666/// uses CAS so concurrent fibers race cleanly without lock
667/// contention; on end-reach the policy is consulted (possibly
668/// duplicated under concurrent end-reach), and the end atomic
669/// is grown via a single CAS (only one extension takes effect
670/// per round; losers retry and see the new end).
671struct ExtendingRangeSource {
672    cursor: Arc<AtomicU64>,
673    end: Arc<AtomicU64>,
674    policy: Arc<dyn ExtensionPolicy>,
675    start: u64,
676    base: u64,
677    /// SRD 71 partition cap — see
678    /// [`ExtendingRangeSourceFactory::bounded`].
679    max_end: Option<u64>,
680    started: std::time::Instant,
681    consumed: u64,
682    schema: SourceSchema,
683}
684
685impl DataSource for ExtendingRangeSource {
686    fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>> {
687        loop {
688            let cur = self.cursor.load(Ordering::Acquire);
689            let end = self.end.load(Ordering::Acquire);
690            if cur < end {
691                // Try to claim [cur, min(cur+stride, end)).
692                let target = (cur.saturating_add(stride as u64)).min(end);
693                match self
694                    .cursor
695                    .compare_exchange(cur, target, Ordering::AcqRel, Ordering::Acquire)
696                {
697                    Ok(_) => {
698                        let count = target - cur;
699                        self.consumed += count;
700                        return Some(cur..target);
701                    }
702                    Err(_) => continue, // raced; retry
703                }
704            }
705            // SRD 71: a partition-bound cursor terminates the
706            // moment the partition is exhausted, whether or not
707            // the policy's time / pass / count target was
708            // reached — no policy consultation past the cap.
709            if let Some(max) = self.max_end
710                && end >= max
711            {
712                return None;
713            }
714            // Cursor has caught up to end. Consult policy with
715            // a snapshot of the current state. Each fiber that
716            // races to this point gets its own context read;
717            // duplicate consultations are idempotent for pure
718            // predicates.
719            let ctx = ExtensionContext {
720                elapsed_ms: self.started.elapsed().as_millis() as u64,
721                consumed: end.saturating_sub(self.start),
722                base: self.base,
723            };
724            match self.policy.next_extension(&ctx) {
725                Some(delta) if delta > 0 => {
726                    // CAS the end forward, clamped at the
727                    // partition cap when one is set. If someone
728                    // else extended ahead of us, that's fine —
729                    // our duplicate policy consultation is
730                    // harmless and the new end is at least as
731                    // far as ours would have been.
732                    let mut new_end = end.saturating_add(delta);
733                    if let Some(max) = self.max_end {
734                        new_end = new_end.min(max);
735                    }
736                    if new_end == end {
737                        return None;
738                    }
739                    let _ = self.end.compare_exchange(
740                        end,
741                        new_end,
742                        Ordering::AcqRel,
743                        Ordering::Acquire,
744                    );
745                    continue;
746                }
747                _ => return None,
748            }
749        }
750    }
751
752    fn render_item(&self, ordinal: u64) -> SourceItem {
753        SourceItem::ordinal(ordinal)
754    }
755
756    fn extent(&self) -> Option<u64> {
757        // Live extent — same convention as the factory's
758        // `global_extent`: subscribers see the current ceiling
759        // even after it grows.
760        Some(self.end.load(Ordering::Acquire).saturating_sub(self.start))
761    }
762
763    fn consumed(&self) -> u64 {
764        self.consumed
765    }
766
767    fn schema(&self) -> &SourceSchema {
768        &self.schema
769    }
770
771    fn replay_contract(&self) -> SourceReplayContract {
772        SourceReplayContract::stable_ordinal(0)
773    }
774}
775
776// =========================================================================
777// Cursors: provenance-driven cursor targeting
778// =========================================================================
779
780/// A cursor target: a DataSource reader paired with its Polydat input index.
781struct CursorTarget {
782    /// The DataSource reader that provides values for this cursor.
783    reader: Box<dyn DataSource>,
784    /// The Polydat input index where the cursor's ordinal is injected.
785    input_index: usize,
786    /// Source name (for diagnostics).
787    source_name: String,
788    /// Stable process-independent identity used in ordinal packet stamps.
789    stream_id: u64,
790}
791
792/// Ownership token for one range reserved from a perfect ordinal source.
793///
794/// The token contains no reader borrow or rendered payload because the source
795/// contract proves that the scalar value is the ordinal itself. Moving this
796/// value transfers responsibility for every ordinal in `range`; the batch
797/// executor must drain all of them or process the remainder through its scalar
798/// recovery path. The shared source cursor is never rewound.
799#[derive(Debug)]
800pub struct OrdinalBatchLease {
801    source_name: String,
802    stream_id: u64,
803    source_generation: u64,
804    input_index: usize,
805    sequence: u64,
806    range: std::ops::Range<u64>,
807}
808
809impl OrdinalBatchLease {
810    /// The cursor the lease reserves from.
811    pub fn source_name(&self) -> &str {
812        &self.source_name
813    }
814
815    /// The stream the ordinals belong to.
816    pub const fn stream_id(&self) -> u64 {
817        self.stream_id
818    }
819
820    /// The generation of the source's values.
821    pub const fn source_generation(&self) -> u64 {
822        self.source_generation
823    }
824
825    /// The kernel input the cursor's ordinal is written to.
826    pub const fn input_index(&self) -> usize {
827        self.input_index
828    }
829
830    /// The lease's position in reservation order.
831    pub const fn sequence(&self) -> u64 {
832        self.sequence
833    }
834
835    /// The ordinals reserved.
836    pub fn range(&self) -> std::ops::Range<u64> {
837        self.range.clone()
838    }
839
840    /// How many ordinals the lease holds.
841    pub const fn len_u64(&self) -> u64 {
842        self.range.end - self.range.start
843    }
844
845    /// Whether the lease holds no ordinal.
846    pub const fn is_empty(&self) -> bool {
847        self.range.start == self.range.end
848    }
849}
850
851#[derive(Clone, Debug, PartialEq, Eq)]
852/// Why a cursor batch could not be reserved.
853pub enum CursorBatchError {
854    /// No ordinals were asked for.
855    ZeroDemand,
856    /// The reservation names several targets; a batch drives one.
857    RequiresSingleTarget {
858        /// The targets named.
859        targets: usize,
860    },
861    /// The source is not a perfect ordinal stream.
862    SourceNotPerfectOrdinal {
863        /// The source's name.
864        source: String,
865    },
866}
867
868impl std::fmt::Display for CursorBatchError {
869    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
870        match self {
871            Self::ZeroDemand => f.write_str("ordinal batch demand must be greater than zero"),
872            Self::RequiresSingleTarget { targets } => write!(
873                f,
874                "Tier-1 ordinal batching requires exactly one cursor target, found {targets}"
875            ),
876            Self::SourceNotPerfectOrdinal { source } => write!(
877                f,
878                "source '{source}' does not declare the perfect ordinal replay contract"
879            ),
880        }
881    }
882}
883
884impl std::error::Error for CursorBatchError {}
885
886/// Provenance-driven advancer that targets only the cursor nodes
887/// relevant to a specific set of output fields.
888///
889/// Built at phase setup by tracing Polydat provenance from the op template's
890/// referenced fields back to root cursor nodes. Only those cursors
891/// advance — unused cursors are left untouched.
892pub struct Cursors {
893    targets: Vec<CursorTarget>,
894    /// Last items read from each target (for injecting into Polydat state).
895    last_items: Vec<Option<SourceItem>>,
896    /// Total advances performed.
897    advances: u64,
898    /// Reservation order for owned ordinal leases from this cursor set.
899    next_batch_sequence: u64,
900}
901
902impl Cursors {
903    /// Build an advancer from a Polydat program, a set of output field names,
904    /// and a map of source name → DataSourceFactory.
905    ///
906    /// Traces provenance: for each field, finds the output node, gets its
907    /// input provenance bitmask, and identifies which inputs are cursor
908    /// ordinals (matching `{source}__ordinal` pattern). Creates a reader
909    /// for each targeted source.
910    pub fn for_fields(
911        program: &crate::kernel::PolydatProgram,
912        field_names: &[&str],
913        source_factories: &std::collections::HashMap<String, Arc<dyn DataSourceFactory>>,
914    ) -> Self {
915        // Collect the union of input provenance for all referenced
916        // fields — exact at any input count (multi-word ProvMask;
917        // the former one-word form over-matched cursor inputs >= 63).
918        let mut combined_provenance = crate::kernel::ProvMask::empty();
919        for name in field_names {
920            if let Some((node_idx, _)) = program.resolve_output(name)
921                && let Some(prov) = program.input_provenance_for(node_idx)
922            {
923                combined_provenance.union_with(prov);
924            }
925        }
926
927        // Find cursor inputs in the provenance
928        let input_names = program.input_names();
929        let mut targets = Vec::new();
930        let mut seen_sources = std::collections::HashSet::new();
931
932        for (idx, input_name) in input_names.iter().enumerate() {
933            if !combined_provenance.contains(idx) {
934                continue;
935            }
936
937            // Check if this input is a source projection ({source}__ordinal)
938            if let Some(source_name) = input_name.strip_suffix("__ordinal") {
939                if seen_sources.contains(source_name) {
940                    continue;
941                }
942                seen_sources.insert(source_name.to_string());
943
944                if let Some(factory) = source_factories.get(source_name) {
945                    let stream_id =
946                        xxhash_rust::xxh3::xxh3_64(format!("{source_name}@{idx}").as_bytes());
947                    targets.push(CursorTarget {
948                        reader: factory.create_reader(),
949                        input_index: idx,
950                        source_name: source_name.to_string(),
951                        stream_id,
952                    });
953                }
954            }
955        }
956
957        let target_count = targets.len();
958        Cursors {
959            targets,
960            last_items: vec![None; target_count],
961            advances: 0,
962            next_batch_sequence: 0,
963        }
964    }
965
966    /// Reserve one owned batch from a single perfect ordinal cursor.
967    ///
968    /// This is the source-side entry point for Tier-1 SIMD execution. It is
969    /// intentionally unavailable for multiple cursor targets or sources that
970    /// merely happen to be monotonic: only the explicit perfect-ordinal
971    /// contract permits payload-free replay from the returned range.
972    pub fn reserve_ordinal_batch(
973        &mut self,
974        demand: usize,
975    ) -> Result<Option<OrdinalBatchLease>, CursorBatchError> {
976        if demand == 0 {
977            return Err(CursorBatchError::ZeroDemand);
978        }
979        if self.targets.len() != 1 {
980            return Err(CursorBatchError::RequiresSingleTarget {
981                targets: self.targets.len(),
982            });
983        }
984
985        let target = &mut self.targets[0];
986        let contract = target.reader.replay_contract();
987        if !contract.is_perfect_ordinal() {
988            return Err(CursorBatchError::SourceNotPerfectOrdinal {
989                source: target.source_name.clone(),
990            });
991        }
992        let Some(range) = target.reader.reserve(demand) else {
993            return Ok(None);
994        };
995        let claimed = range.end - range.start;
996        let lease = OrdinalBatchLease {
997            source_name: target.source_name.clone(),
998            stream_id: target.stream_id,
999            source_generation: contract.generation,
1000            input_index: target.input_index,
1001            sequence: self.next_batch_sequence,
1002            range,
1003        };
1004        self.next_batch_sequence = self.next_batch_sequence.wrapping_add(1);
1005        self.advances = self.advances.saturating_add(claimed);
1006        Ok(Some(lease))
1007    }
1008
1009    /// Advance all targeted cursors. Returns `false` if any targeted
1010    /// cursor is exhausted (no more data).
1011    ///
1012    /// After advancing, the new ordinals and field projections are
1013    /// available via `inject_into_state()`.
1014    pub fn advance(&mut self) -> bool {
1015        for (i, target) in self.targets.iter_mut().enumerate() {
1016            match target.reader.next() {
1017                Some(item) => {
1018                    self.last_items[i] = Some(item);
1019                }
1020                None => return false, // this cursor exhausted
1021            }
1022        }
1023        self.advances += 1;
1024        true
1025    }
1026
1027    /// Inject the current cursor values into a Polydat state.
1028    ///
1029    /// Sets each cursor's ordinal at its input index, plus any
1030    /// field projections from the source item.
1031    pub fn inject_into_state(&self, state: &mut crate::kernel::PolydatState) {
1032        for (i, target) in self.targets.iter().enumerate() {
1033            if let Some(ref item) = self.last_items[i] {
1034                state.set_input(target.input_index, crate::ast::Value::U64(item.ordinal));
1035                // Inject field projections (e.g., base__vector)
1036                // These are handled by set_source_item on the FiberBuilder
1037            }
1038        }
1039    }
1040
1041    /// Get the last items read from all targets (for FiberBuilder injection).
1042    pub fn last_items(&self) -> &[Option<SourceItem>] {
1043        &self.last_items
1044    }
1045
1046    /// Known extent of the driving cursor (smallest among targeted
1047    /// cursors with known extent). Used for progress reporting.
1048    pub fn extent(&self) -> Option<u64> {
1049        self.targets.iter().filter_map(|t| t.reader.extent()).min()
1050    }
1051
1052    /// Total advances performed so far.
1053    pub fn consumed(&self) -> u64 {
1054        self.advances
1055    }
1056
1057    /// Number of targeted cursors.
1058    pub fn target_count(&self) -> usize {
1059        self.targets.len()
1060    }
1061
1062    /// Whether this advancer has any targets.
1063    pub fn is_empty(&self) -> bool {
1064        self.targets.is_empty()
1065    }
1066}
1067
1068// =========================================================================
1069// Built-in ExtensionPolicy implementations
1070// =========================================================================
1071//
1072// Each policy is a pure predicate over `ExtensionContext` and
1073// returns `Some(delta)` when the cursor should continue or
1074// `None` when it should stop. The delta is independent of the
1075// stop condition — workloads can extend in chunks unrelated to
1076// the cursor's base size (e.g. base=10000 but delta=1000 to
1077// check the condition more often).
1078
1079/// Extend while `ctx.elapsed_ms < min_ms`, projecting the
1080/// remaining work from the observed rate and committing it in
1081/// one batch.
1082///
1083/// Per call: estimate `repeats = floor((consumed * remaining_ms
1084/// / elapsed_ms) / base)` — the number of base-sized passes
1085/// needed to fill the time remaining at the current rate. Apply
1086/// a 5% under-bias (`repeats * 95 / 100`) so successive
1087/// extensions converge from below the target rather than
1088/// overshooting once. Commit `repeats * base` cycles.
1089///
1090/// The under-bias gives a geometric convergence: each call
1091/// covers ~95% of the time remaining, so 3–4 calls typically
1092/// land within 0.01% of the time budget. The first call (no
1093/// rate signal yet — `elapsed_ms == 0` or `consumed == 0`)
1094/// falls back to one base chunk via the `delta` field.
1095pub struct UntilElapsedPolicy {
1096    /// The elapsed milliseconds to reach.
1097    pub min_ms: u64,
1098    /// The extension step for the first call, before a rate is known.
1099    pub delta: u64,
1100}
1101
1102impl ExtensionPolicy for UntilElapsedPolicy {
1103    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1104        if ctx.elapsed_ms >= self.min_ms {
1105            return None;
1106        }
1107        // No rate signal yet — first end-reach with elapsed time
1108        // below the clock resolution, or a degenerate empty
1109        // cursor. Step by the declared `delta` (which the
1110        // executor sets to `base` when the workload didn't
1111        // supply an explicit `delta` arg) so the next call has
1112        // a rate measurement to project from.
1113        if ctx.elapsed_ms == 0 || ctx.consumed == 0 {
1114            return Some(self.delta.max(1));
1115        }
1116        let remaining_ms = self.min_ms - ctx.elapsed_ms;
1117        // u128 saturating math: consumed * remaining_ms can
1118        // overflow u64 for long-running phases with high
1119        // throughput (e.g. 10^9 ops over 10^4 ms).
1120        let est_remaining =
1121            (ctx.consumed as u128).saturating_mul(remaining_ms as u128) / (ctx.elapsed_ms as u128);
1122        let biased = est_remaining.saturating_mul(95) / 100;
1123        let base = ctx.base.max(1) as u128;
1124        let repeats = biased / base;
1125        if repeats == 0 {
1126            return None;
1127        }
1128        let delta = repeats.saturating_mul(base);
1129        Some(u64::try_from(delta).unwrap_or(u64::MAX))
1130    }
1131}
1132
1133/// Extend by `delta` while `ctx.passes() < min_passes`. Passes
1134/// are whole multiples of `ctx.base`.
1135pub struct UntilPassesPolicy {
1136    /// The passes to reach.
1137    pub min_passes: u64,
1138    /// The extension step.
1139    pub delta: u64,
1140}
1141
1142impl ExtensionPolicy for UntilPassesPolicy {
1143    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1144        if ctx.passes() < self.min_passes {
1145            Some(self.delta)
1146        } else {
1147            None
1148        }
1149    }
1150}
1151
1152/// Extend by `delta` while `ctx.consumed < min_count`.
1153pub struct UntilCountPolicy {
1154    /// The consumed count to reach.
1155    pub min_count: u64,
1156    /// The extension step.
1157    pub delta: u64,
1158}
1159
1160impl ExtensionPolicy for UntilCountPolicy {
1161    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1162        if ctx.consumed < self.min_count {
1163            Some(self.delta)
1164        } else {
1165            None
1166        }
1167    }
1168}
1169
1170/// Continue (extend) while ALL child policies say continue.
1171/// Stops the moment any child policy returns `None`. The
1172/// delta is the minimum of the children's deltas — conservative
1173/// step size keeps any single condition from over-shooting its
1174/// stop point.
1175pub struct AndPolicy {
1176    /// The policies that must all ask to continue.
1177    pub policies: Vec<Arc<dyn ExtensionPolicy>>,
1178}
1179
1180impl ExtensionPolicy for AndPolicy {
1181    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1182        let mut min_delta = u64::MAX;
1183        for p in &self.policies {
1184            let d = p.next_extension(ctx)?;
1185            min_delta = min_delta.min(d);
1186        }
1187        if min_delta == u64::MAX || min_delta == 0 {
1188            None
1189        } else {
1190            Some(min_delta)
1191        }
1192    }
1193}
1194
1195/// Continue (extend) while ANY child policy says continue.
1196/// Stops only when every child policy returns `None`. The
1197/// delta is the maximum of the policies that said continue —
1198/// matches the most aggressive child still pushing forward.
1199pub struct OrPolicy {
1200    /// The policies of which any may ask to continue.
1201    pub policies: Vec<Arc<dyn ExtensionPolicy>>,
1202}
1203
1204impl ExtensionPolicy for OrPolicy {
1205    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1206        let mut max_delta: Option<u64> = None;
1207        for p in &self.policies {
1208            if let Some(d) = p.next_extension(ctx) {
1209                max_delta = Some(max_delta.map(|m| m.max(d)).unwrap_or(d));
1210            }
1211        }
1212        max_delta
1213    }
1214}
1215
1216/// Back-compat alias for the original time-only policy.
1217/// Constructs an [`UntilElapsedPolicy`] with delta = base.
1218pub struct TimeElapsedPolicy {
1219    inner: UntilElapsedPolicy,
1220}
1221
1222impl TimeElapsedPolicy {
1223    /// A policy extending in steps of `base` until `min_ms` have elapsed.
1224    pub fn new(base: u64, min_ms: u64) -> Self {
1225        Self {
1226            inner: UntilElapsedPolicy {
1227                min_ms,
1228                delta: base,
1229            },
1230        }
1231    }
1232}
1233
1234impl ExtensionPolicy for TimeElapsedPolicy {
1235    fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1236        self.inner.next_extension(ctx)
1237    }
1238}
1239
1240#[cfg(test)]
1241mod tests {
1242    use super::*;
1243
1244    #[test]
1245    fn range_sources_declare_the_perfect_ordinal_replay_contract() {
1246        let factory = RangeSourceFactory::new(10, 20);
1247        assert!(factory.replay_contract().is_perfect_ordinal());
1248        let reader = factory.create_reader();
1249        assert!(reader.replay_contract().is_perfect_ordinal());
1250        let item = reader.render_item(13);
1251        assert_eq!(item.ordinal, 13);
1252        assert!(item.fields.is_empty());
1253    }
1254
1255    /// Policy that always extends by `base` — stands in for a
1256    /// time/pass policy whose target is far away, so the
1257    /// partition cap is the only thing that can stop growth.
1258    struct AlwaysExtend;
1259    impl ExtensionPolicy for AlwaysExtend {
1260        fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1261            Some(ctx.base.max(1))
1262        }
1263    }
1264
1265    #[test]
1266    fn extending_source_bounded_terminates_at_partition_end() {
1267        // SRD 71: `until_*(base, ...) over p` — base-sized chunks
1268        // walk within the partition; the cap stops growth even
1269        // though the policy would keep extending. Partition
1270        // [100, 125) with base 10 → 10 + 10 + 5, then exhausted.
1271        let factory =
1272            ExtendingRangeSourceFactory::new("q", 100, 10, Arc::new(AlwaysExtend)).bounded(125);
1273        let mut reader = factory.create_reader();
1274        let mut total = 0u64;
1275        let mut last_end = 100;
1276        while let Some(r) = reader.reserve(7) {
1277            assert!(r.end <= 125, "reservation past the partition cap: {r:?}");
1278            assert_eq!(r.start, last_end, "contiguous reservations");
1279            last_end = r.end;
1280            total += r.end - r.start;
1281        }
1282        assert_eq!(total, 25, "exactly the partition's cardinality");
1283        assert_eq!(last_end, 125);
1284    }
1285
1286    #[test]
1287    fn extending_source_bounded_clamps_oversized_base() {
1288        // A base chunk larger than the partition never reserves
1289        // past it.
1290        let factory =
1291            ExtendingRangeSourceFactory::new("q", 0, 1000, Arc::new(AlwaysExtend)).bounded(30);
1292        let mut reader = factory.create_reader();
1293        let r = reader.reserve(usize::MAX).unwrap();
1294        assert_eq!(r, 0..30);
1295        assert!(reader.reserve(1).is_none());
1296    }
1297
1298    #[test]
1299    fn extending_source_unbounded_keeps_policy_semantics() {
1300        // Without a cap the policy alone decides — three
1301        // extensions of a terminating policy.
1302        struct NTimes(std::sync::atomic::AtomicU64);
1303        impl ExtensionPolicy for NTimes {
1304            fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
1305                if self.0.fetch_add(1, Ordering::Relaxed) < 3 {
1306                    Some(ctx.base)
1307                } else {
1308                    None
1309                }
1310            }
1311        }
1312        let factory = ExtendingRangeSourceFactory::new(
1313            "q",
1314            0,
1315            10,
1316            Arc::new(NTimes(std::sync::atomic::AtomicU64::new(0))),
1317        );
1318        let mut reader = factory.create_reader();
1319        let mut total = 0u64;
1320        while let Some(r) = reader.reserve(64) {
1321            total += r.end - r.start;
1322        }
1323        assert_eq!(total, 40, "initial 10 + three 10-ordinal extensions");
1324    }
1325
1326    #[test]
1327    fn range_source_yields_ordinals() {
1328        let factory = RangeSourceFactory::new(0, 5);
1329        let mut reader = factory.create_reader();
1330        assert_eq!(reader.extent(), Some(5));
1331
1332        for i in 0..5 {
1333            let item = reader.next().unwrap();
1334            assert_eq!(item.ordinal, i);
1335            assert!(item.fields.is_empty());
1336        }
1337        assert!(reader.next().is_none());
1338        assert_eq!(reader.consumed(), 5);
1339    }
1340
1341    #[test]
1342    fn range_source_chunk() {
1343        let factory = RangeSourceFactory::new(0, 10);
1344        let mut reader = factory.create_reader();
1345
1346        let chunk = reader.next_chunk(3);
1347        assert_eq!(chunk.len(), 3);
1348        assert_eq!(chunk[0].ordinal, 0);
1349        assert_eq!(chunk[2].ordinal, 2);
1350
1351        let chunk = reader.next_chunk(100);
1352        assert_eq!(chunk.len(), 7); // only 7 remaining
1353        assert_eq!(chunk[0].ordinal, 3);
1354        assert_eq!(chunk[6].ordinal, 9);
1355
1356        let chunk = reader.next_chunk(1);
1357        assert!(chunk.is_empty()); // exhausted
1358    }
1359
1360    #[test]
1361    fn range_source_concurrent_readers() {
1362        let factory = RangeSourceFactory::new(0, 100);
1363        let mut r1 = factory.create_reader();
1364        let mut r2 = factory.create_reader();
1365
1366        // Each reader gets unique ordinals from the shared cursor
1367        let a = r1.next().unwrap().ordinal;
1368        let b = r2.next().unwrap().ordinal;
1369        assert_ne!(a, b);
1370
1371        // Drain both readers
1372        let mut total = 2;
1373        while r1.next().is_some() {
1374            total += 1;
1375        }
1376        while r2.next().is_some() {
1377            total += 1;
1378        }
1379        assert_eq!(total, 100);
1380    }
1381
1382    #[test]
1383    fn source_item_field_access() {
1384        let item = SourceItem::with_fields(
1385            42,
1386            vec![
1387                ("name".into(), Value::Str("test".into())),
1388                ("score".into(), Value::F64(0.95)),
1389            ],
1390        );
1391        assert_eq!(item.ordinal, 42);
1392        assert_eq!(item.field("name"), Some(&Value::Str("test".into())));
1393        assert_eq!(item.field("score"), Some(&Value::F64(0.95)));
1394        assert_eq!(item.field("missing"), None);
1395    }
1396
1397    #[test]
1398    fn range_source_named() {
1399        let factory = RangeSourceFactory::named("users", 0, 1000);
1400        assert_eq!(factory.schema().name, "users");
1401        assert_eq!(factory.schema().extent, Some(1000));
1402    }
1403
1404    // ── ExtendingRangeSource ─────────────────────────────────
1405
1406    /// Trivial extension policy for tests: lets the caller
1407    /// decide how many extensions remain. Each call to
1408    /// `next_extension` decrements the count.
1409    struct FixedExtensions {
1410        delta: u64,
1411        remaining: std::sync::atomic::AtomicU64,
1412    }
1413
1414    impl FixedExtensions {
1415        fn new(delta: u64, times: u64) -> Self {
1416            Self {
1417                delta,
1418                remaining: std::sync::atomic::AtomicU64::new(times),
1419            }
1420        }
1421    }
1422
1423    impl ExtensionPolicy for FixedExtensions {
1424        fn next_extension(&self, _ctx: &ExtensionContext) -> Option<u64> {
1425            let prev = self.remaining.fetch_sub(1, Ordering::Relaxed);
1426            if prev == 0 || prev > i64::MAX as u64 {
1427                self.remaining.store(0, Ordering::Relaxed);
1428                None
1429            } else {
1430                Some(self.delta)
1431            }
1432        }
1433    }
1434
1435    #[test]
1436    fn extending_source_consumes_initial_extent_then_extends() {
1437        // Initial extent = 5; policy extends once by 5 more.
1438        // Total expected: 10 ordinals consumed.
1439        let policy = Arc::new(FixedExtensions::new(5, 1));
1440        let factory = ExtendingRangeSourceFactory::new("ext", 0, 5, policy);
1441        let mut reader = factory.create_reader();
1442        let mut got: Vec<u64> = Vec::new();
1443        while let Some(item) = reader.next() {
1444            got.push(item.ordinal);
1445            if got.len() > 50 {
1446                panic!("runaway extension");
1447            }
1448        }
1449        assert_eq!(got, (0..10).collect::<Vec<u64>>());
1450        assert_eq!(reader.consumed(), 10);
1451    }
1452
1453    #[test]
1454    fn extending_source_zero_extensions_behaves_like_fixed_range() {
1455        let policy = Arc::new(FixedExtensions::new(0, 0));
1456        let factory = ExtendingRangeSourceFactory::new("ext", 0, 3, policy);
1457        let mut reader = factory.create_reader();
1458        let mut got: Vec<u64> = Vec::new();
1459        while let Some(item) = reader.next() {
1460            got.push(item.ordinal);
1461        }
1462        assert_eq!(got, vec![0, 1, 2]);
1463    }
1464
1465    #[test]
1466    fn extending_source_global_extent_grows_after_extension() {
1467        // global_extent must reflect the CURRENT end so phase
1468        // status displays show growth honestly.
1469        let policy = Arc::new(FixedExtensions::new(10, 2));
1470        let factory = ExtendingRangeSourceFactory::new("ext", 0, 5, policy);
1471        assert_eq!(factory.global_extent(), Some(5));
1472        let mut reader = factory.create_reader();
1473        // Drain the first 5 to force an extension.
1474        for _ in 0..5 {
1475            reader.next().unwrap();
1476        }
1477        // Trigger the extension by attempting one more pull.
1478        let _ = reader.next().unwrap();
1479        assert_eq!(
1480            factory.global_extent(),
1481            Some(15),
1482            "extent should grow by the extension delta"
1483        );
1484    }
1485
1486    #[test]
1487    fn extending_source_chunked_reservation_caps_at_current_end() {
1488        // A stride request that crosses the current `end` must
1489        // return a SHORT range up to `end`, not an over-claim.
1490        // Otherwise the next reservation would jump past the
1491        // extended range and lose ordinals.
1492        let policy = Arc::new(FixedExtensions::new(5, 1));
1493        let factory = ExtendingRangeSourceFactory::new("ext", 0, 3, policy);
1494        let mut reader = factory.create_reader();
1495        let range = reader.reserve(10).expect("first reserve");
1496        // Initial end was 3; the reservation must cap there.
1497        assert_eq!(range, 0..3);
1498        // Next reserve triggers the extension and consumes the
1499        // remainder.
1500        let range = reader.reserve(10).expect("second reserve");
1501        assert_eq!(range, 3..8);
1502    }
1503
1504    fn ctx_at(elapsed_ms: u64, consumed: u64, base: u64) -> ExtensionContext {
1505        ExtensionContext {
1506            elapsed_ms,
1507            consumed,
1508            base,
1509        }
1510    }
1511
1512    #[test]
1513    fn until_elapsed_policy_bootstrap_falls_back_to_delta_when_no_rate_signal() {
1514        // First end-reach: consumed or elapsed effectively zero —
1515        // no rate to project from. Policy returns `delta` exactly
1516        // so the next call has a measurement to work with.
1517        let policy = UntilElapsedPolicy {
1518            min_ms: 50,
1519            delta: 7,
1520        };
1521        assert_eq!(policy.next_extension(&ctx_at(0, 0, 0)), Some(7));
1522        assert_eq!(policy.next_extension(&ctx_at(49, 0, 0)), Some(7));
1523        assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(7));
1524    }
1525
1526    #[test]
1527    fn until_elapsed_policy_stops_at_or_past_min_ms() {
1528        let policy = UntilElapsedPolicy {
1529            min_ms: 50,
1530            delta: 7,
1531        };
1532        assert_eq!(policy.next_extension(&ctx_at(50, 100, 100)), None);
1533        assert_eq!(policy.next_extension(&ctx_at(1000, 100, 100)), None);
1534    }
1535
1536    #[test]
1537    fn until_elapsed_policy_projects_remaining_cycles_with_under_bias() {
1538        // base=100, consumed=100 (one pass), elapsed=200ms,
1539        // min_ms=1000 → remaining=800ms. Rate = 100/200 = 0.5
1540        // cycles/ms; project 0.5 * 800 = 400 cycles. Under-bias
1541        // 5% → 380. Round to multiples of base (100) → 3 repeats
1542        // → 300 cycles.
1543        let policy = UntilElapsedPolicy {
1544            min_ms: 1000,
1545            delta: 100,
1546        };
1547        assert_eq!(
1548            policy.next_extension(&ctx_at(200, 100, 100)),
1549            Some(300),
1550            "expected 3-pass batch (380 under-biased, floored to base multiples)",
1551        );
1552    }
1553
1554    #[test]
1555    fn until_elapsed_policy_returns_none_when_remaining_is_under_one_pass() {
1556        // base=100, consumed=10000 over 990ms, min_ms=1000 →
1557        // remaining=10ms. Project 10000/990*10 ≈ 101 cycles.
1558        // Under-bias → 95. Rounded to base multiples → 0 repeats.
1559        // Policy terminates rather than under-shooting the budget
1560        // with a wasted partial pass.
1561        let policy = UntilElapsedPolicy {
1562            min_ms: 1000,
1563            delta: 100,
1564        };
1565        assert_eq!(policy.next_extension(&ctx_at(990, 10000, 100)), None);
1566    }
1567
1568    #[test]
1569    fn until_elapsed_policy_converges_geometrically() {
1570        // Sanity: simulate the extension loop. Each iteration
1571        // commits ~95% of the remaining time budget; the
1572        // residual is bounded below by the base-pass rounding,
1573        // so convergence is one-base-coarse rather than
1574        // arbitrarily tight.
1575        let policy = UntilElapsedPolicy {
1576            min_ms: 1000,
1577            delta: 10,
1578        };
1579        // Pretend the first pass took 10ms (rate = 1 cycle/ms).
1580        let mut elapsed = 10u64;
1581        let mut consumed = 10u64;
1582        let mut iters = 0;
1583        while let Some(delta) = policy.next_extension(&ctx_at(elapsed, consumed, 10)) {
1584            iters += 1;
1585            assert!(iters < 10, "geometric series should converge fast");
1586            consumed += delta;
1587            // Same rate: 1 cycle/ms.
1588            elapsed += delta;
1589        }
1590        assert!(elapsed <= 1000, "must under-shoot, got elapsed={elapsed}");
1591        // Residual ≤ ~5% (under-bias) + 1 base pass (rounding) ≈ 6% of target.
1592        assert!(
1593            elapsed >= 940,
1594            "must come within ~6% of target, got {elapsed}"
1595        );
1596    }
1597
1598    #[test]
1599    fn until_passes_policy_counts_in_base_multiples() {
1600        let policy = UntilPassesPolicy {
1601            min_passes: 3,
1602            delta: 100,
1603        };
1604        // 0 passes done — extend.
1605        assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(100));
1606        // 2 passes done (200 consumed @ base=100) — extend.
1607        assert_eq!(policy.next_extension(&ctx_at(0, 200, 100)), Some(100));
1608        // 3 passes done — stop.
1609        assert_eq!(policy.next_extension(&ctx_at(0, 300, 100)), None);
1610        // 4 passes done — stop.
1611        assert_eq!(policy.next_extension(&ctx_at(0, 400, 100)), None);
1612    }
1613
1614    #[test]
1615    fn until_count_policy_uses_raw_consumed() {
1616        let policy = UntilCountPolicy {
1617            min_count: 250,
1618            delta: 50,
1619        };
1620        assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(50));
1621        assert_eq!(policy.next_extension(&ctx_at(0, 249, 100)), Some(50));
1622        assert_eq!(policy.next_extension(&ctx_at(0, 250, 100)), None);
1623    }
1624
1625    #[test]
1626    fn and_policy_stops_when_any_child_stops() {
1627        // time<5000 AND passes<3.
1628        let time = Arc::new(UntilElapsedPolicy {
1629            min_ms: 5000,
1630            delta: 10,
1631        });
1632        let passes = Arc::new(UntilPassesPolicy {
1633            min_passes: 3,
1634            delta: 20,
1635        });
1636        let and = AndPolicy {
1637            policies: vec![time, passes],
1638        };
1639        // Both still want to continue: delta = min(10, 20) = 10.
1640        assert_eq!(and.next_extension(&ctx_at(0, 0, 100)), Some(10));
1641        // Time done — stop.
1642        assert_eq!(and.next_extension(&ctx_at(5000, 0, 100)), None);
1643        // Passes done — stop.
1644        assert_eq!(and.next_extension(&ctx_at(0, 300, 100)), None);
1645    }
1646
1647    #[test]
1648    fn or_policy_continues_if_any_child_continues() {
1649        let time = Arc::new(UntilElapsedPolicy {
1650            min_ms: 5000,
1651            delta: 10,
1652        });
1653        let passes = Arc::new(UntilPassesPolicy {
1654            min_passes: 3,
1655            delta: 20,
1656        });
1657        let or = OrPolicy {
1658            policies: vec![time, passes],
1659        };
1660        // Both want to continue → delta = max(10, 20) = 20.
1661        assert_eq!(or.next_extension(&ctx_at(0, 0, 100)), Some(20));
1662        // Time done, passes still wants → delta = 20.
1663        assert_eq!(or.next_extension(&ctx_at(5000, 0, 100)), Some(20));
1664        // Passes done, time still wants → delta = 10.
1665        assert_eq!(or.next_extension(&ctx_at(0, 300, 100)), Some(10));
1666        // Both done → stop.
1667        assert_eq!(or.next_extension(&ctx_at(5000, 300, 100)), None);
1668    }
1669
1670    #[test]
1671    fn time_elapsed_policy_compat_alias_still_works() {
1672        let p = TimeElapsedPolicy::new(7, 1);
1673        std::thread::sleep(std::time::Duration::from_millis(20));
1674        assert_eq!(p.next_extension(&ctx_at(20, 0, 0)), None);
1675    }
1676}