Skip to main content

triplets_core/sampler/
mod.rs

1mod backends;
2
3use crate::chunking::{ChunkingAlgorithm, SlidingWindowChunker};
4use chrono::Duration;
5use indexmap::IndexMap;
6use nohash_hasher::IntSet;
7use rand::prelude::*;
8use rayon::prelude::*;
9use std::borrow::Cow;
10use std::collections::{HashMap, HashSet};
11use std::path::Path;
12use std::sync::atomic::{AtomicUsize, Ordering};
13use std::sync::mpsc;
14use std::sync::{Arc, Mutex};
15use std::thread;
16
17use crate::config::{
18    ChunkingStrategy, NegativeStrategy, SamplerConfig, Selector, TextRecipe, TripletRecipe,
19};
20use crate::constants::sampler::AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME;
21use crate::constants::sampler::{
22    ANCHOR_POSITIVE_SWAP_MASK, EPOCH_SEED_OFFSET, EXHAUSTION_RETRY_LIMIT, NEG_REASON_SAME_RECORD,
23    NEG_REASON_WRONG_ARTICLE, NEG_REASON_WRONG_DATE, NEG_REASON_WRONG_QA, PREFETCHER_SOURCE_ID,
24    PREFETCHER_STOPPED_REASON, RECIPE_LABEL_TEXT, RECIPE_LABEL_TRIPLETS,
25    RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER, ROLE_LABEL_ANCHOR, ROLE_LABEL_CONTEXT,
26    SAME_SELECTOR_PAIR_RETRY_LIMIT,
27};
28use crate::data::{
29    ChunkView, DataRecord, PairLabel, RecordChunk, RecordSection, SampleBatch, SamplePair,
30    SampleTriplet, SectionRole, TextBatch, TextSample, TripletBatch,
31};
32use crate::epoch::EpochTracker;
33use crate::errors::SamplerError;
34use crate::hash::{derive_epoch_seed, stable_hash_key, stable_hash_key_section, stable_hash_str};
35use crate::ingestion::IngestionManager;
36use crate::metadata::{META_FIELD_DATE, MetadataKey};
37use crate::metrics::{chunk_proximity_score, window_index_proximity};
38use crate::source::DataSource;
39use crate::splits::{
40    EpochStateStore, PersistedSamplerState, SamplerStateStore, SplitLabel, SplitStore,
41};
42use crate::tokenizer::{Tokenizer, WhitespaceTokenizer};
43use crate::types::{RecipeKey, RecordKey, SourceId};
44use crate::utils::platform_newline;
45
46// AUTO-RECIPE HANDLING OVERVIEW (end-to-end):
47// Stage A: Source-level injection eligibility ("should this source even get the recipe?")
48//   - `sync_records_from_cache` -> `record_has_long_anchor_or_context_section` marks
49//     `long_section_counts`.
50//   - `resolve_source_triplet_plan` -> `should_auto_inject_chunk_pair_recipe` appends
51//     `source_chunk_pair_recipe` for eligible sources.
52// Stage B: Record-level execution eligibility ("can this specific record run it now?")
53//   - `make_auto_chunk_pair_triplet_with_anchor` checks
54//     `record_has_at_least_two_window_chunks_for_selector` before sampling.
55// Stage C: Actual chunk window layout ("how are windows formed?")
56//   - `materialize_chunks` is the source of truth for token window construction.
57
58#[derive(Debug, Clone)]
59/// Small deterministic RNG used for reproducible sampler behavior.
60struct DeterministicRng {
61    state: u64,
62}
63
64impl DeterministicRng {
65    fn new(seed: u64) -> Self {
66        Self { state: seed }
67    }
68
69    fn from_state(state: u64) -> Self {
70        Self { state }
71    }
72
73    fn state(&self) -> u64 {
74        self.state
75    }
76
77    fn next_u64_internal(&mut self) -> u64 {
78        let mut z = self.state.wrapping_add(0x9E3779B97F4A7C15);
79        self.state = z;
80        z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
81        z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
82        z ^ (z >> 31)
83    }
84}
85
86impl rand::RngCore for DeterministicRng {
87    fn next_u32(&mut self) -> u32 {
88        self.next_u64_internal() as u32
89    }
90
91    fn next_u64(&mut self) -> u64 {
92        self.next_u64_internal()
93    }
94
95    fn fill_bytes(&mut self, dest: &mut [u8]) {
96        let mut offset = 0;
97        while offset < dest.len() {
98            let value = self.next_u64_internal();
99            let bytes = value.to_le_bytes();
100            let remaining = dest.len() - offset;
101            let copy_len = remaining.min(bytes.len());
102            dest[offset..offset + copy_len].copy_from_slice(&bytes[..copy_len]);
103            offset += copy_len;
104        }
105    }
106}
107
108/// Public helper so callers (and examples) can compute the same chunk weighting the sampler uses.
109///
110/// Definitions:
111/// - `trust` (range 0.0-1.0): comes from `RecordChunk.quality.trust`; scales the contribution of a chunk.
112/// - `base`: window chunks use proximity-to-head (`1 / (index + 1)`); summary chunks use their configured `summary_fallback_weight`.
113/// - `chunk_weight_floor`: minimum weight applied after scaling.
114///
115/// Formula: `max(chunk_weight_floor, base * trust)`.
116pub fn chunk_weight(strategy: &ChunkingStrategy, chunk: &RecordChunk) -> f32 {
117    let floor = strategy.chunk_weight_floor;
118    let trust = chunk.quality.trust.clamp(0.0, 1.0);
119    let base = match &chunk.view {
120        ChunkView::Window { index, .. } => window_index_proximity(*index),
121        ChunkView::SummaryFallback { weight, .. } => *weight,
122    };
123    (base * trust).max(floor)
124}
125
126/// Public sampling interface for pair, triplet, and text batch generation.
127pub trait Sampler {
128    /// Returns a batch of pairs. Consumes the shared epoch cursor for anchor selection.
129    ///
130    /// Samples all sources uniformly. Prefer [`Sampler::next_pair_batch_with_weights`]
131    /// with an explicit per-source weight map to honor a data mixture.
132    #[deprecated(
133        note = "samples all sources uniformly; call next_pair_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
134    )]
135    fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
136        self.next_pair_batch_with_weights(split, &HashMap::new())
137    }
138    /// Returns a batch of text samples. Consumes the shared epoch cursor for anchor selection.
139    ///
140    /// Samples all sources uniformly. Prefer [`Sampler::next_text_batch_with_weights`]
141    /// with an explicit per-source weight map to honor a data mixture.
142    #[deprecated(
143        note = "samples all sources uniformly; call next_text_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
144    )]
145    fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
146        self.next_text_batch_with_weights(split, &HashMap::new())
147    }
148    /// Returns a batch of triplets. Consumes the shared epoch cursor for anchor selection.
149    ///
150    /// Samples all sources uniformly. Prefer [`Sampler::next_triplet_batch_with_weights`]
151    /// with an explicit per-source weight map to honor a data mixture.
152    #[deprecated(
153        note = "samples all sources uniformly; call next_triplet_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
154    )]
155    fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
156        self.next_triplet_batch_with_weights(split, &HashMap::new())
157    }
158    /// Returns a batch of pairs with per-call source weights.
159    fn next_pair_batch_with_weights(
160        &self,
161        split: SplitLabel,
162        weights: &HashMap<SourceId, f32>,
163    ) -> Result<SampleBatch, SamplerError>;
164    /// Returns a batch of text samples with per-call source weights.
165    fn next_text_batch_with_weights(
166        &self,
167        split: SplitLabel,
168        weights: &HashMap<SourceId, f32>,
169    ) -> Result<TextBatch, SamplerError>;
170    /// Returns a batch of triplets with per-call source weights.
171    fn next_triplet_batch_with_weights(
172        &self,
173        split: SplitLabel,
174        weights: &HashMap<SourceId, f32>,
175    ) -> Result<TripletBatch, SamplerError>;
176    /// Persist sampler and split runtime state for restart-resume.
177    ///
178    /// Default is a no-op, allowing mock/test implementations to skip persistence.
179    fn save_sampler_state(&self, _save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
180        Ok(())
181    }
182}
183
184/// Background prefetcher that fills a bounded queue with sample batches.
185pub struct BatchPrefetcher<T> {
186    receiver: Option<mpsc::Receiver<Result<T, SamplerError>>>,
187    handle: Option<thread::JoinHandle<()>>,
188    stats: Arc<PrefetcherStats>,
189}
190
191#[derive(Default)]
192/// Prefetcher runtime counters.
193struct PrefetcherStats {
194    queued: AtomicUsize,
195    produced: AtomicUsize,
196    errors: AtomicUsize,
197}
198
199impl<T: Send + 'static> BatchPrefetcher<T> {
200    fn new<F>(capacity: usize, mut producer: F) -> Self
201    where
202        F: FnMut() -> Result<T, SamplerError> + Send + 'static,
203    {
204        let (sender, receiver) = mpsc::sync_channel(capacity.max(1));
205        let stats = Arc::new(PrefetcherStats::default());
206        let stats_thread = Arc::clone(&stats);
207        let handle = thread::spawn(move || {
208            loop {
209                let result = producer();
210                if result.is_err() {
211                    stats_thread.errors.fetch_add(1, Ordering::Relaxed);
212                }
213                if sender.send(result).is_err() {
214                    return;
215                }
216                stats_thread.queued.fetch_add(1, Ordering::Relaxed);
217                stats_thread.produced.fetch_add(1, Ordering::Relaxed);
218            }
219        });
220        Self {
221            receiver: Some(receiver),
222            handle: Some(handle),
223            stats,
224        }
225    }
226
227    /// Block until the next prefetched batch is available.
228    pub fn next(&self) -> Result<T, SamplerError> {
229        let receiver = self
230            .receiver
231            .as_ref()
232            .ok_or_else(|| SamplerError::SourceUnavailable {
233                source_id: PREFETCHER_SOURCE_ID.into(),
234                reason: PREFETCHER_STOPPED_REASON.into(),
235            })?;
236        let result = receiver.recv().unwrap_or_else(|_| {
237            Err(SamplerError::SourceUnavailable {
238                source_id: PREFETCHER_SOURCE_ID.into(),
239                reason: PREFETCHER_STOPPED_REASON.into(),
240            })
241        });
242        self.stats
243            .queued
244            .try_update(Ordering::Relaxed, Ordering::Relaxed, |value| {
245                Some(value.saturating_sub(1))
246            })
247            .ok();
248        result
249    }
250
251    /// Number of prefetched batches currently queued.
252    pub fn queue_len(&self) -> usize {
253        self.stats.queued.load(Ordering::Relaxed)
254    }
255
256    /// Total number of batches produced by the background worker.
257    pub fn produced_count(&self) -> usize {
258        self.stats.produced.load(Ordering::Relaxed)
259    }
260
261    /// Total number of errors produced by the background worker.
262    pub fn error_count(&self) -> usize {
263        self.stats.errors.load(Ordering::Relaxed)
264    }
265}
266
267impl<T> Drop for BatchPrefetcher<T> {
268    fn drop(&mut self) {
269        self.receiver.take();
270        if let Some(handle) = self.handle.take() {
271            let _ = handle.join();
272        }
273    }
274}
275
276/// Sampler that draws anchors from a single shared epoch cursor and then
277/// selects chunks from those records. Ingestion happens on demand when sampling.
278pub struct TripletSampler<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
279    inner: Mutex<TripletSamplerInner<S>>,
280}
281
282/// Internal sampler state implementation guarded by `TripletSampler`.
283struct TripletSamplerInner<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
284    /// Immutable sampler configuration (seed, batch size, recipes, splits, etc.).
285    config: SamplerConfig,
286    /// Active chunking implementation used to materialize section chunks.
287    chunker: Arc<dyn ChunkingAlgorithm>,
288    /// Split store backing train/val/test assignments and persisted sampler state.
289    split_store: Arc<S>,
290    /// On-demand ingestion manager that fills the batch-sized buffer.
291    ingestion: IngestionManager,
292    /// Current in-memory record pool keyed by composite record identity.
293    /// Composite keys (not bare ids) so overlapping ids from different
294    /// sources coexist instead of overwriting each other.
295    records: IndexMap<RecordKey, Arc<DataRecord>>,
296    /// Deficit counters for weighted source selection (smooth weighted
297    /// round-robin over `source_order`). Only engaged when the caller passes
298    /// explicit non-uniform per-source weights; uniform/empty maps take the
299    /// legacy uniform-cycling path bit-identically. New sources start at 0.0
300    /// via entry-or-insert (no sync hook needed); stale entries are pruned
301    /// with the source index. Intentionally NOT persisted: deficits
302    /// self-correct within a few draws, and no consumer persists sampler
303    /// state across processes with weights engaged.
304    source_deficits: HashMap<SourceId, f32>,
305    /// Deterministic RNG for per-batch shuffles and sampling.
306    rng: DeterministicRng,
307    /// Config-level triplet recipes used when sources do not supply their own.
308    triplet_recipes: Vec<TripletRecipe>,
309    /// Config-level text recipes used when sources do not supply their own.
310    text_recipes: Vec<TextRecipe>,
311    /// Per-source triplet recipes keyed by source id.
312    source_triplet_recipes: HashMap<SourceId, Vec<TripletRecipe>>,
313    /// Per-source count of pooled records containing at least one section
314    /// larger than the chunk window. A source is eligible for the
315    /// auto-injected chunk-pair recipe while its count is greater than zero.
316    /// Maintained incrementally on record insert/evict (strict O(Δ)) instead
317    /// of rescanning the pool on every sync.
318    long_section_counts: HashMap<SourceId, usize>,
319    /// Per-source text recipes keyed by source id.
320    source_text_recipes: HashMap<SourceId, Vec<TextRecipe>>,
321    /// True if triplet recipes came from config (no source defaults).
322    using_config_triplet_recipes: bool,
323    /// True if text recipes came from config (no derivation).
324    using_config_text_recipes: bool,
325    /// Last observed ingestion counter for cache updates.
326    last_observed_ingest: u64,
327    /// Epoch tracker for split-aware deterministic sampling.
328    epoch_tracker: EpochTracker,
329    /// Per-record, per-section chunk cursor to rotate through chunk windows.
330    chunk_cursors: HashMap<(RecordKey, usize), usize>,
331    /// Per-record, per-role cursor to rotate through role-specific sections.
332    role_cursors: HashMap<(RecordKey, String), usize>,
333    /// Pluggable negative-selection backend (uniform-random or BM25).
334    negative_backend: Box<dyn backends::NegativeBackend>,
335    /// Chunk key to parent record key lookup (used by epoch tracker).
336    chunk_index: HashMap<RecordKey, RecordKey>,
337    /// Round-robin order of source ids (deterministic).
338    source_order: Vec<SourceId>,
339    /// Current index into `source_order` for shuffled-cycle sampling.
340    source_cycle_idx: usize,
341    /// True once persisted source state is loaded.
342    source_state_loaded: bool,
343    /// True once persisted ingestion cursors are loaded.
344    ingestion_cursors_loaded: bool,
345    /// Marks whether source state needs persistence.
346    source_state_dirty: bool,
347    /// Record keys per source for round-robin within a source.
348    source_record_indices: HashMap<SourceId, Vec<RecordKey>>,
349    /// Per-source cursor into `source_record_indices`.
350    source_record_cursors: HashMap<SourceId, usize>,
351    /// Round-robin index for triplet recipe cycling.
352    triplet_recipe_rr_idx: usize,
353    /// Hashes of text already emitted by the text batch path while the record
354    /// pool is stable.  Pruned on sync_records_from_cache only when records
355    /// are evicted.  Prevents the same text from being sampled again within
356    /// the same ingestion window, regardless of source wrapping or epoch.
357    ///
358    /// Per-record tracking lets us surgically remove hashes only for evicted
359    /// records, preserving cross-batch dedup across small window advances
360    /// (e.g. window=512, batch_size=4 where only 4/512 records are swapped).
361    emitted_text_hashes: HashMap<RecordKey, HashSet<u64>>,
362    /// Round-robin index for text recipe cycling.
363    text_recipe_rr_idx: usize,
364    /// Epoch counter for per-source deterministic shuffling (seed ^ epoch).
365    epoch: u64,
366
367    /// Tracks whether each source has wrapped its cursor in the current epoch.
368    source_wrapped: HashMap<SourceId, bool>,
369
370    /// Last sync version watermark for incremental delta synchronization.
371    last_sync_version: u64,
372
373    /// Monotonic counter advanced only when record-pool membership actually
374    /// changes (genuinely new IDs added or IDs truly evicted). Pure version
375    /// refreshes of already-pooled IDs do not advance it, so the cached
376    /// `records_by_split` map below stays valid across steady-state advances.
377    pool_generation: u64,
378
379    /// Cached `records_by_split` output, valid while
380    /// `cached_records_by_split_gen == pool_generation`. Avoids rebuilding the
381    /// full per-split map (O(N) clones + label lookups) on every batch.
382    cached_records_by_split: Option<HashMap<SplitLabel, Vec<RecordKey>>>,
383    /// Pool generation the cached map was built from.
384    cached_records_by_split_gen: u64,
385
386    /// Last observed force-refresh generation. When the ingestion manager
387    /// reports a newer force-refresh generation, the next sync falls back to
388    /// a full pool resync (exact full-snapshot semantics) instead of the
389    /// O(Δ) delta path, because a force refresh replaces the cache contents
390    /// wholesale. Steady-state advances keep the delta fast path.
391    last_force_refresh_generation: u64,
392
393    /// Persistent cache of split labels keyed by composite record identity.
394    /// Eliminates redundant DataStore::read() calls for split label lookups.
395    split_labels: HashMap<RecordKey, SplitLabel>,
396}
397
398impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSamplerInner<S> {
399    fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
400        Self::new_with_chunker(config, split_store, Arc::new(SlidingWindowChunker))
401    }
402
403    fn new_with_chunker(
404        config: SamplerConfig,
405        split_store: Arc<S>,
406        chunker: Arc<dyn ChunkingAlgorithm>,
407    ) -> Self {
408        let buffer_size = config.ingestion_max_records.max(config.batch_size).max(2);
409        let using_config_triplet_recipes = !config.recipes.is_empty();
410        let using_config_text_recipes = !config.text_recipes.is_empty();
411        let triplet_recipes = if using_config_triplet_recipes {
412            config.recipes.clone()
413        } else {
414            Vec::new()
415        };
416        let text_recipes = if using_config_text_recipes {
417            config.text_recipes.clone()
418        } else if !triplet_recipes.is_empty() {
419            Self::build_derived_text_recipes(&triplet_recipes)
420        } else {
421            Vec::new()
422        };
423        let ingestion = IngestionManager::new(buffer_size, config.clone());
424        let epoch_backend = Some(Arc::clone(&split_store) as Arc<dyn EpochStateStore>);
425        let epoch_tracker = EpochTracker::new(
426            true,
427            epoch_backend,
428            derive_epoch_seed(config.seed, EPOCH_SEED_OFFSET),
429        );
430        let mut sampler = Self {
431            rng: DeterministicRng::new(config.seed),
432            config,
433            chunker,
434            split_store,
435            ingestion,
436            records: IndexMap::new(),
437            triplet_recipes,
438            text_recipes,
439            source_triplet_recipes: HashMap::new(),
440            long_section_counts: HashMap::new(),
441            source_text_recipes: HashMap::new(),
442            using_config_triplet_recipes,
443            using_config_text_recipes,
444            last_observed_ingest: 0,
445            epoch_tracker,
446            chunk_cursors: HashMap::new(),
447            role_cursors: HashMap::new(),
448            negative_backend: {
449                #[cfg(feature = "bm25-mining")]
450                {
451                    Box::new(backends::Bm25Backend::new())
452                }
453                #[cfg(not(feature = "bm25-mining"))]
454                {
455                    Box::new(backends::DefaultBackend)
456                }
457            },
458            chunk_index: HashMap::new(),
459            source_order: Vec::new(),
460            source_cycle_idx: 0,
461            source_deficits: HashMap::new(),
462            source_state_loaded: false,
463            ingestion_cursors_loaded: false,
464            source_state_dirty: false,
465            source_record_indices: HashMap::new(),
466            source_record_cursors: HashMap::new(),
467            emitted_text_hashes: HashMap::new(),
468            triplet_recipe_rr_idx: 0,
469            text_recipe_rr_idx: 0,
470            epoch: 0,
471            source_wrapped: HashMap::new(),
472            last_sync_version: 0,
473            last_force_refresh_generation: 0,
474            pool_generation: 0,
475            cached_records_by_split: None,
476            cached_records_by_split_gen: 0,
477            split_labels: HashMap::new(),
478        };
479        if !sampler.using_config_text_recipes {
480            sampler.rebuild_derived_text_recipes();
481        }
482        sampler
483    }
484
485    fn text_recipes(&self) -> &[TextRecipe] {
486        &self.text_recipes
487    }
488
489    /// Current epoch-adjusted seed: mixes `epoch` into `config.seed` so every epoch
490    /// produces a distinct permutation across all seed-dependent operations.
491    fn epoch_seed(&self) -> u64 {
492        derive_epoch_seed(self.config.seed, self.epoch)
493    }
494
495    /// Get or compute the split label for a record, using the persistent cache.
496    fn get_or_insert_split_label(&mut self, key: &RecordKey) -> Result<SplitLabel, SamplerError> {
497        if let Some(&label) = self.split_labels.get(key) {
498            return Ok(label);
499        }
500        let label = match self.split_store.label_for(key) {
501            Some(label) => label,
502            None => self.split_store.ensure(key.clone())?,
503        };
504        self.split_labels.insert(key.clone(), label);
505        Ok(label)
506    }
507
508    fn register_source(
509        &mut self,
510        source: Box<dyn DataSource + 'static>,
511    ) -> Result<(), SamplerError> {
512        let source_id = source.id().to_string();
513        if !self.using_config_triplet_recipes {
514            let triplets = source.default_triplet_recipes();
515            if !triplets.is_empty() {
516                self.source_triplet_recipes
517                    .insert(source_id.clone(), triplets.clone());
518                if !self.using_config_text_recipes {
519                    let derived = Self::build_derived_text_recipes(&triplets);
520                    self.source_text_recipes
521                        .insert(source_id.clone(), derived.clone());
522                    self.extend_text_recipes_unique(&derived);
523                }
524            }
525        }
526        self.ingestion.register_source(source)?;
527        Ok(())
528    }
529
530    fn set_epoch(&mut self, epoch: u64) -> Result<(), SamplerError> {
531        self.epoch_tracker.ensure_loaded()?;
532        self.epoch_tracker.force_epoch(epoch);
533        self.epoch = epoch;
534        self.ingestion.set_epoch(epoch);
535        self.ingestion.reset_stream_cursors();
536        // Reset epoch step at epoch boundary so each epoch
537        // starts with step=0, giving deterministic step sequences.
538        self.ingestion.reset_epoch_step();
539        self.source_record_cursors.clear();
540        self.source_cycle_idx = 0;
541        for source in &self.source_order {
542            self.source_wrapped.insert(source.clone(), false);
543        }
544        self.rebuild_source_index()?;
545        self.source_state_dirty = self.source_order.len() > 1;
546        Ok(())
547    }
548
549    fn next_chunk_from_pool(
550        &mut self,
551        key: &RecordKey,
552        section_idx: usize,
553        pool: Vec<RecordChunk>,
554    ) -> Option<RecordChunk> {
555        if pool.is_empty() {
556            return None;
557        }
558        let map_key = (key.clone(), section_idx);
559        if !self.chunk_cursors.contains_key(&map_key) {
560            // First touch should **not** always start at window 0.
561            // We derive a deterministic per-(record,section) starting offset so:
562            // - repeated runs with the same seed/epoch are reproducible,
563            // - first sampled window is spread across records/sections,
564            // - subsequent calls still rotate cyclically through the pool.
565            let cursor_key = stable_hash_key_section(self.epoch_seed(), key, section_idx);
566            let start = (cursor_key as usize) % pool.len();
567            self.chunk_cursors.insert(map_key.clone(), start);
568        }
569        let cursor = self.chunk_cursors.entry(map_key).or_insert(0);
570        if *cursor >= pool.len() {
571            *cursor = 0;
572        }
573        let chunk = pool.get(*cursor).cloned();
574        *cursor = (*cursor + 1) % pool.len();
575        chunk
576    }
577
578    fn prune_cursor_state(&mut self) {
579        if self.chunk_cursors.is_empty()
580            && self.role_cursors.is_empty()
581            && self.negative_backend.cursors_empty()
582        {
583            return;
584        }
585        let valid_keys: HashSet<RecordKey> = self.records.keys().cloned().collect();
586        self.chunk_cursors
587            .retain(|(record_key, _), _| valid_keys.contains(record_key));
588        self.role_cursors
589            .retain(|(record_key, _), _| valid_keys.contains(record_key));
590        self.negative_backend.prune_cursors(&valid_keys);
591    }
592
593    /// Ordered removal of a record id from one source's index vector.
594    /// Empty source entries are left in place; the caller prunes them after
595    /// all evictions and additions are processed (mirroring
596    /// `rebuild_source_index`, which only contains sources with at least one
597    /// allowed-split record).
598    fn remove_id_from_source_index(&mut self, source: &SourceId, record_key: &RecordKey) {
599        if let Some(ids) = self.source_record_indices.get_mut(source)
600            && let Some(pos) = ids.iter().position(|key| key == record_key)
601        {
602            ids.remove(pos);
603        }
604    }
605
606    /// Insert a record key into one source's index vector, preserving the
607    /// deterministic `stable_hash_key(epoch_seed, key)` sort order via binary
608    /// search. Records whose split is not allowed are not indexed, mirroring
609    /// `rebuild_source_index`. Creates the source entry when needed; the
610    /// caller keeps `source_order` sorted afterwards.
611    fn insert_id_into_source_index(
612        &mut self,
613        source: SourceId,
614        record_key: RecordKey,
615        label: SplitLabel,
616        allowed: &HashSet<SplitLabel>,
617    ) {
618        if !allowed.contains(&label) {
619            return;
620        }
621        let seed = self.epoch_seed();
622        let key = stable_hash_key(seed, &record_key);
623        // Upper-bound insertion: with distinct hashes this is the unique
624        // sorted position; on a hash collision it places the key after
625        // existing equals, matching stable-sort semantics.
626        let pos = match self.source_record_indices.get(&source) {
627            Some(ids) => ids.partition_point(|k| stable_hash_key(seed, k) <= key),
628            None => 0,
629        };
630        self.source_record_indices
631            .entry(source)
632            .or_default()
633            .insert(pos, record_key);
634    }
635
636    /// Increment the long-section record count for a source.
637    fn increment_long_section_count(&mut self, source: &SourceId) {
638        *self.long_section_counts.entry(source.clone()).or_default() += 1;
639    }
640
641    /// Decrement the long-section record count for a source, removing the
642    /// entry when no counted records remain.
643    fn decrement_long_section_count(&mut self, source: &SourceId) {
644        let remove = if let Some(count) = self.long_section_counts.get_mut(source) {
645            *count = count.saturating_sub(1);
646            *count == 0
647        } else {
648            false
649        };
650        if remove {
651            self.long_section_counts.remove(source);
652        }
653    }
654
655    fn rebuild_chunk_index(&mut self) {
656        self.chunk_index.clear();
657        for record in self.records.values() {
658            let key = RecordKey::from(record);
659            self.chunk_index.insert(key.clone(), key);
660        }
661    }
662
663    fn rebuild_source_index(&mut self) -> Result<(), SamplerError> {
664        self.source_record_indices.clear();
665        let allowed = self.allowed_target_splits();
666        let allowed_set: HashSet<SplitLabel> = allowed.into_iter().collect();
667        // Collect record keys first to avoid borrowing issues
668        let record_keys: Vec<RecordKey> = self.records.values().map(RecordKey::of_arc).collect();
669        for record_key in record_keys {
670            let label = self.get_or_insert_split_label(&record_key)?;
671            if !allowed_set.contains(&label) {
672                continue;
673            }
674            self.source_record_indices
675                .entry(record_key.source.clone())
676                .or_default()
677                .push(record_key);
678        }
679
680        let shuffle_seed = self.epoch_seed();
681        for record_keys in self.source_record_indices.values_mut() {
682            record_keys.sort_by_key(|key| stable_hash_key(shuffle_seed, key));
683        }
684
685        self.source_order = self.source_record_indices.keys().cloned().collect();
686        self.source_order.sort();
687        self.refresh_source_wrapped();
688        // Drop deficit counters for sources that left the pool.
689        self.source_deficits
690            .retain(|s, _| self.source_record_indices.contains_key(s));
691
692        self.source_record_cursors
693            .retain(|source, _| self.source_record_indices.contains_key(source));
694        if self.source_state_loaded {
695            if self.source_order.is_empty() {
696                self.source_cycle_idx = 0;
697            }
698            self.source_state_dirty = self.source_order.len() > 1;
699        }
700        Ok(())
701    }
702
703    fn refresh_source_wrapped(&mut self) {
704        self.source_wrapped.clear();
705        for source in &self.source_order {
706            let len = self
707                .source_record_indices
708                .get(source)
709                .map(|items| items.len())
710                .unwrap_or(0);
711            if len == 0 {
712                self.source_wrapped.insert(source.clone(), false);
713                continue;
714            }
715            let cursor = self.source_record_cursors.get(source).copied().unwrap_or(0);
716            let wrapped = cursor > 0 && cursor % len == 0;
717            self.source_wrapped.insert(source.clone(), wrapped);
718        }
719    }
720
721    fn shuffled_source_cycle(&self, cycle: u64) -> Vec<SourceId> {
722        let mut sources = self.source_order.clone();
723        let seed = self.epoch_seed() ^ cycle;
724        sources.sort_by_key(|source| stable_hash_str(seed, source));
725        sources
726    }
727
728    /// True when `weights` carries no selective signal: empty or all equal.
729    /// Uniform maps take the legacy uniform-cycling path bit-identically.
730    fn weights_are_selective(weights: &HashMap<SourceId, f32>) -> bool {
731        if weights.is_empty() {
732            return false;
733        }
734        let first = *weights.values().next().unwrap();
735        weights.values().any(|&w| w != first)
736    }
737
738    /// Effective quantum for `source`: its weight when finite and
739    /// non-negative, else 0. Unlisted sources default to 1.0 (mirrors the
740    /// ingestion drain).
741    fn weight_quantum(weights: &HashMap<SourceId, f32>, source: &str) -> f32 {
742        weights.get(source).copied().unwrap_or(1.0).max(0.0)
743    }
744
745    /// Weighted source choice: smooth weighted round-robin over `candidates`.
746    /// Credits each eligible source its weight quantum, picks the highest
747    /// deficit, deducts the round total. Deterministic given pool + deficits
748    /// (no RNG): same state, same choice. New sources start at 0.0 via
749    /// entry-or-insert. Returns `None` when no candidate has records.
750    fn select_weighted_source(
751        &mut self,
752        candidates: &[SourceId],
753        weights: &HashMap<SourceId, f32>,
754    ) -> Option<SourceId> {
755        let eligible: Vec<&SourceId> = candidates
756            .iter()
757            .filter(|s| {
758                self.source_record_indices
759                    .get(*s)
760                    .map(|v| !v.is_empty())
761                    .unwrap_or(false)
762            })
763            .collect();
764        if eligible.is_empty() {
765            return None;
766        }
767        let mut total = 0.0f32;
768        for source in &eligible {
769            let q = Self::weight_quantum(weights, source);
770            total += q;
771            *self.source_deficits.entry((*source).clone()).or_insert(0.0) += q;
772        }
773        if total <= 0.0 {
774            return None;
775        }
776        let mut best: Option<(f32, SourceId)> = None;
777        for source in &eligible {
778            let d = self.source_deficits.get(*source).copied().unwrap_or(0.0);
779            if best.as_ref().map(|(bd, _)| d > *bd).unwrap_or(true) {
780                best = Some((d, (*source).clone()));
781            }
782        }
783        let (_, source) = best?;
784        *self
785            .source_deficits
786            .get_mut(&source)
787            .expect("deficit credited above") -= total;
788        Some(source)
789    }
790
791    fn ensure_source_state(&mut self) -> Result<(), SamplerError> {
792        if self.source_state_loaded {
793            return Ok(());
794        }
795        let persisted = self.split_store.load_sampler_state()?;
796        self.source_cycle_idx = persisted
797            .as_ref()
798            .map(|state| state.source_cycle_idx as usize)
799            .unwrap_or(0);
800        if let Some(state) = persisted {
801            for (source, cursor) in state.source_record_cursors {
802                if self.source_record_indices.contains_key(&source) {
803                    self.source_record_cursors.insert(source, cursor as usize);
804                }
805            }
806            self.epoch = state.epoch;
807            self.ingestion.set_epoch(state.epoch);
808            self.rng = DeterministicRng::from_state(state.rng_state);
809            self.triplet_recipe_rr_idx = state.triplet_recipe_rr_idx as usize;
810            self.text_recipe_rr_idx = state.text_recipe_rr_idx as usize;
811        }
812        self.refresh_source_wrapped();
813        self.source_state_loaded = true;
814        self.source_state_dirty = true;
815        Ok(())
816    }
817
818    fn persist_source_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
819        if !self.source_state_loaded {
820            return Ok(());
821        }
822        let state = PersistedSamplerState {
823            source_cycle_idx: self.source_cycle_idx as u64,
824            source_record_cursors: self
825                .source_record_cursors
826                .iter()
827                .map(|(source, cursor)| (source.clone(), *cursor as u64))
828                .collect(),
829            epoch: self.epoch,
830            epoch_step: self.ingestion.epoch_step(),
831            rng_state: self.rng.state(),
832            triplet_recipe_rr_idx: self.triplet_recipe_rr_idx as u64,
833            text_recipe_rr_idx: self.text_recipe_rr_idx as u64,
834            source_stream_cursors: self.ingestion.snapshot_cursors(),
835        };
836        self.split_store.save_sampler_state(&state, save_to)?;
837        self.source_state_dirty = false;
838        Ok(())
839    }
840
841    fn rebuild_derived_text_recipes(&mut self) {
842        if self.using_config_text_recipes {
843            return;
844        }
845        if self.triplet_recipes.is_empty() {
846            self.text_recipes.clear();
847        } else {
848            self.text_recipes = Self::build_derived_text_recipes(&self.triplet_recipes);
849        }
850    }
851
852    fn extend_text_recipes_unique(&mut self, recipes: &[TextRecipe]) {
853        for recipe in recipes {
854            if self
855                .text_recipes
856                .iter()
857                .any(|existing| existing.name == recipe.name)
858            {
859                continue;
860            }
861            self.text_recipes.push(recipe.clone());
862        }
863    }
864
865    fn configured_triplet_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TripletRecipe] {
866        if self.using_config_triplet_recipes {
867            return &self.triplet_recipes;
868        }
869        self.source_triplet_recipes
870            .get(source)
871            .map(|recipes| recipes.as_slice())
872            .unwrap_or(&[])
873    }
874
875    /// Returns true when `recipes` already contains the long-section auto recipe.
876    fn contains_auto_chunk_pair_recipe(recipes: &[TripletRecipe]) -> bool {
877        recipes
878            .iter()
879            .any(|recipe| recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME)
880    }
881
882    fn source_supports_chunk_pair_recipe(&self, source: &str) -> bool {
883        if self.config.chunking.max_window_tokens == 0 {
884            return false;
885        }
886        self.long_section_counts
887            .get(source)
888            .is_some_and(|&count| count > 0)
889    }
890
891    /// Stage A (source-level): decide whether to append auto recipe for `source`.
892    ///
893    /// Decision criteria:
894    /// - source must be eligible (`source_supports_chunk_pair_recipe`), and
895    /// - configured pool must not already include the auto recipe name.
896    fn should_auto_inject_chunk_pair_recipe(
897        &self,
898        source: &str,
899        recipes: &[TripletRecipe],
900    ) -> bool {
901        self.source_supports_chunk_pair_recipe(source)
902            && !Self::contains_auto_chunk_pair_recipe(recipes)
903    }
904
905    /// Build the auto-injected long-section recipe.
906    ///
907    /// Semantics:
908    /// - Anchor selector: `Context`
909    /// - Positive selector: `Context`
910    /// - Negative selector: `Context` on a different record (`WrongArticle`)
911    ///
912    /// Anchor and positive are selected by two independent `select_chunk` calls
913    /// from the same context chunk candidate pool for the chosen record.
914    /// They are not concatenated and one is not derived from the other's text.
915    fn source_chunk_pair_recipe() -> TripletRecipe {
916        TripletRecipe {
917            name: Cow::Borrowed(AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME),
918            anchor: Selector::Role(SectionRole::Context),
919            positive_selector: Selector::Role(SectionRole::Context),
920            negative_selector: Selector::Role(SectionRole::Context),
921            negative_strategy: NegativeStrategy::WrongArticle,
922            weight: 1.0,
923            instruction: None,
924            allow_same_anchor_positive: false,
925        }
926    }
927
928    /// Resolve the source triplet plan.
929    ///
930    /// This is Stage A (source-level injection), not record-level execution.
931    ///
932    /// Algorithm for the auto-triplet injection path:
933    /// 1) Start from the configured source recipe pool (`configured_triplet_recipes_for_source`).
934    /// 2) Check auto-injection eligibility (`should_auto_inject_chunk_pair_recipe`), which is true only when:
935    ///    - the source has long sections discovered during ingest sync, and
936    ///    - the pool does not already include the auto recipe name.
937    /// 3) If eligible, append `source_chunk_pair_recipe`.
938    /// 4) Return both the effective recipe pool and whether step (3) happened.
939    fn resolve_source_triplet_plan(&self, source: &str) -> (Vec<TripletRecipe>, bool) {
940        let mut recipes = self.configured_triplet_recipes_for_source(source).to_vec();
941        let mut auto_injected = false;
942        if self.should_auto_inject_chunk_pair_recipe(source, &recipes) {
943            recipes.push(Self::source_chunk_pair_recipe());
944            auto_injected = true;
945        }
946        (recipes, auto_injected)
947    }
948
949    #[cfg(test)]
950    fn triplet_recipes_for_source(&self, source: &str) -> Vec<TripletRecipe> {
951        self.resolve_source_triplet_plan(source).0
952    }
953
954    fn triplet_recipe_count_for_source(&self, source: &str) -> usize {
955        let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
956        recipes.len()
957    }
958
959    fn text_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TextRecipe] {
960        if self.using_config_text_recipes || self.using_config_triplet_recipes {
961            return &self.text_recipes;
962        }
963        self.source_text_recipes
964            .get(source)
965            .map(|recipes| recipes.as_slice())
966            .unwrap_or(&[])
967    }
968
969    /// Build a weighted, shuffled selection order from a slice of recipe weights.
970    ///
971    /// Each recipe with `weight > 0.0` receives a number of slots proportional to its weight
972    /// relative to the smallest positive weight, capped at
973    /// [`RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER`].  The slot list is then shuffled so that
974    /// deterministic round-robin cycling naturally draws recipes at the requested frequency:
975    ///
976    /// * `weight = [1.0, 1.0, 1.0]` → one slot each — identical to uniform round-robin.
977    /// * `weight = [3.0, 1.0]`       → three slots for recipe 0, one for recipe 1.
978    /// * `weight = [1.0, 0.0]`       → recipe 1 is excluded entirely (zero slots).
979    fn recipe_order_weighted_shuffled(
980        &mut self,
981        weights: &[f32],
982        rng: &mut DeterministicRng,
983    ) -> Vec<usize> {
984        weighted_recipe_order(weights, rng)
985    }
986
987    /// Weighted shuffle rotated by `rr_idx` for deterministic round-robin cycling.
988    ///
989    /// Calls [`Self::recipe_order_weighted_shuffled`] and rotates the result by
990    /// `rr_idx % order.len()` so successive batches start from different recipes,
991    /// preserving both the proportional frequency and the cross-batch cycling.
992    fn recipe_order_weighted_cycled(
993        &mut self,
994        weights: &[f32],
995        rr_idx: usize,
996        rng: &mut DeterministicRng,
997    ) -> Vec<usize> {
998        let base = self.recipe_order_weighted_shuffled(weights, rng);
999        if base.is_empty() {
1000            return base;
1001        }
1002        let start = rr_idx % base.len();
1003        let mut order = Vec::with_capacity(base.len());
1004        order.extend_from_slice(&base[start..]);
1005        order.extend_from_slice(&base[..start]);
1006        order
1007    }
1008
1009    /// Weighted shuffled selection order for text recipes — same algorithm as
1010    /// [`Self::recipe_order_weighted_shuffled`], kept separate so the text and triplet
1011    /// round-robin counters can advance independently.
1012    fn text_recipe_order_weighted_shuffled(
1013        &mut self,
1014        weights: &[f32],
1015        rng: &mut DeterministicRng,
1016    ) -> Vec<usize> {
1017        weighted_recipe_order(weights, rng)
1018    }
1019
1020    fn text_recipe_order_weighted_cycled(
1021        &mut self,
1022        weights: &[f32],
1023        rr_idx: usize,
1024        rng: &mut DeterministicRng,
1025    ) -> Vec<usize> {
1026        let base = self.text_recipe_order_weighted_shuffled(weights, rng);
1027        if base.is_empty() {
1028            return base;
1029        }
1030        let start = rr_idx % base.len();
1031        let mut order = Vec::with_capacity(base.len());
1032        order.extend_from_slice(&base[start..]);
1033        order.extend_from_slice(&base[..start]);
1034        order
1035    }
1036
1037    fn allowed_target_splits(&self) -> Vec<SplitLabel> {
1038        self.config.allowed_splits.clone()
1039    }
1040
1041    fn ensure_split_allowed(&self, split: SplitLabel) -> Result<(), SamplerError> {
1042        let allowed = self.allowed_target_splits();
1043        if allowed.contains(&split) {
1044            return Ok(());
1045        }
1046        Err(SamplerError::Configuration(format!(
1047            "requested split {:?} is not in allowed_splits {:?}",
1048            split, allowed
1049        )))
1050    }
1051
1052    fn ensure_split_has_records(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
1053        self.refresh_records_by_split_cache()?;
1054        let has_records = self
1055            .cached_records_by_split
1056            .as_ref()
1057            .and_then(|map| map.get(&target_split))
1058            .is_some_and(|records| !records.is_empty());
1059        if has_records {
1060            return Ok(());
1061        }
1062        Err(SamplerError::Exhausted(
1063            "no records available for target split".into(),
1064        ))
1065    }
1066
1067    /// Rebuild the cached `records_by_split` map when the record pool changed
1068    /// since the last build. Steady-state batches hit the cache (O(1)) instead
1069    /// of rebuilding the full map (O(N) clones + label lookups).
1070    fn refresh_records_by_split_cache(&mut self) -> Result<(), SamplerError> {
1071        if self.cached_records_by_split_gen == self.pool_generation
1072            && self.cached_records_by_split.is_some()
1073        {
1074            return Ok(());
1075        }
1076        let mut map: HashMap<SplitLabel, Vec<RecordKey>> = HashMap::new();
1077        // Collect chunk_index entries first to avoid borrowing issues
1078        let entries: Vec<(RecordKey, RecordKey)> = self
1079            .chunk_index
1080            .iter()
1081            .map(|(k, v)| (k.clone(), v.clone()))
1082            .collect();
1083        for (chunk_key, record_key) in entries {
1084            if !self.records.contains_key(&record_key) {
1085                continue;
1086            }
1087            let label = self.get_or_insert_split_label(&record_key)?;
1088            map.entry(label).or_default().push(chunk_key);
1089        }
1090        self.cached_records_by_split = Some(map);
1091        self.cached_records_by_split_gen = self.pool_generation;
1092        Ok(())
1093    }
1094
1095    /// Owned-map convenience used by tests that hand-mutate the pool.
1096    /// Production paths use [`Self::refresh_records_by_split_cache`] plus a
1097    /// direct reference to avoid the per-call clone.
1098    #[cfg(test)]
1099    fn records_by_split(&mut self) -> Result<HashMap<SplitLabel, Vec<RecordKey>>, SamplerError> {
1100        self.refresh_records_by_split_cache()?;
1101        Ok(self.cached_records_by_split.clone().unwrap_or_default())
1102    }
1103
1104    fn choose_anchor_record(
1105        &mut self,
1106        source: Option<&str>,
1107        split: SplitLabel,
1108    ) -> Option<Arc<DataRecord>> {
1109        if let Some(source) = source {
1110            let record_ids = self.source_record_indices.get(source)?;
1111            if record_ids.is_empty() {
1112                return None;
1113            }
1114            let len = record_ids.len();
1115            let mut cursor = *self.source_record_cursors.get(source).unwrap_or(&0);
1116            let cycle = cursor / len;
1117            let offset_seed = self.epoch_seed() ^ (cycle as u64);
1118            let offset = (stable_hash_str(offset_seed, source) as usize) % len;
1119            let mut wrapped = false;
1120            let mut selected_key: Option<RecordKey> = None;
1121
1122            for _ in 0..len {
1123                let pos = (cursor % len).saturating_add(offset) % len;
1124                let record_key = &record_ids[pos];
1125                cursor = cursor.saturating_add(1);
1126                if cursor.is_multiple_of(len) {
1127                    wrapped = true;
1128                }
1129                if self.split_labels.get(record_key) == Some(&split) {
1130                    selected_key = Some(record_key.clone());
1131                    break;
1132                }
1133            }
1134
1135            self.source_record_cursors
1136                .insert(source.to_string(), cursor);
1137            if wrapped {
1138                self.mark_source_wrapped(source);
1139            }
1140
1141            selected_key.and_then(|key| self.records.get(&key).map(Arc::clone))
1142        } else {
1143            while let Some(chunk_key) = self.epoch_tracker.next_record(split) {
1144                if let Some(record_key) = self.chunk_index.get(&chunk_key)
1145                    && let Some(record) = self.records.get(record_key)
1146                {
1147                    return Some(Arc::clone(record));
1148                }
1149            }
1150            None
1151        }
1152    }
1153
1154    fn save_sampler_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
1155        if self.epoch_tracker.is_enabled() {
1156            self.epoch_tracker.persist()?;
1157        }
1158        self.persist_source_state(save_to)?;
1159        Ok(())
1160    }
1161
1162    fn mark_source_wrapped(&mut self, source: &str) {
1163        self.source_wrapped.insert(source.to_string(), true);
1164        if self.source_order.is_empty() {
1165            return;
1166        }
1167        let all_wrapped = self
1168            .source_order
1169            .iter()
1170            .all(|name| self.source_wrapped.get(name).copied().unwrap_or(false));
1171        if all_wrapped {
1172            self.advance_epoch();
1173        }
1174    }
1175
1176    fn advance_epoch(&mut self) {
1177        // The epoch advanced, changing the sampling permutation.
1178        // emitted_text_hashes is NOT cleared: the record pool is unchanged,
1179        // so texts already emitted are still off-limits.  The epoch
1180        // changes the shuffling permutation, not the dedup window.
1181        // Advance the epoch counter so source refreshes receive a different
1182        // permutation seed (seed ^ epoch), producing a fresh traversal order.
1183        self.epoch = self.epoch.saturating_add(1);
1184        self.ingestion.set_epoch(self.epoch);
1185        // Do NOT reset the epoch step here: it tracks the number of
1186        // refresh batches, which is orthogonal to the sampling epoch.
1187        // Resetting it would corrupt the persisted cursor value on resume.
1188        // Reset per-source record cursors so the new epoch starts from the
1189        // beginning of each source's permuted index space (not from where the
1190        // previous epoch left off).
1191        self.source_record_cursors.clear();
1192        self.source_cycle_idx = 0;
1193        // Reset the wrapped-tracking so all sources are eligible for sampling
1194        // in the new epoch.
1195        for source in &self.source_order {
1196            self.source_wrapped.insert(source.clone(), false);
1197        }
1198        let _ = self.rebuild_source_index();
1199        self.source_state_dirty = self.source_order.len() > 1;
1200    }
1201
1202    fn select_temporal_neighbor(
1203        &'_ self,
1204        record: &DataRecord,
1205        offset_days: i32,
1206    ) -> Option<Arc<DataRecord>> {
1207        let target = record.created_at + Duration::days(offset_days.into());
1208        let key = record.taxonomy.first().cloned();
1209        let record_split = self.split_store.label_for(&RecordKey::from(record))?;
1210        self.records
1211            .values()
1212            .filter(|candidate| {
1213                RecordKey::from(*candidate) != RecordKey::from(record)
1214                    && self
1215                        .split_store
1216                        .label_for(&RecordKey::from(*candidate))
1217                        .map(|label| label == record_split)
1218                        .unwrap_or(false)
1219                    && (candidate.source == record.source
1220                        || key
1221                            .as_ref()
1222                            .zip(candidate.taxonomy.first())
1223                            .map(|(a, b)| a == b)
1224                            .unwrap_or(false))
1225            })
1226            .min_by_key(|candidate| (candidate.created_at - target).num_seconds().abs())
1227            .cloned()
1228    }
1229
1230    fn select_negative_record(
1231        &self,
1232        anchor_record: &DataRecord,
1233        strategy: &NegativeStrategy,
1234        anchor_query_text: Option<&str>,
1235        rng: &mut dyn rand::RngCore,
1236    ) -> Option<(Arc<DataRecord>, bool)> {
1237        let anchor_split = self
1238            .split_store
1239            .label_for(&RecordKey::from(anchor_record))?;
1240
1241        let in_anchor_split = |candidate: &DataRecord| {
1242            self.split_store
1243                .label_for(&RecordKey::from(candidate))
1244                .map(|label| label == anchor_split)
1245                .unwrap_or(false)
1246        };
1247
1248        match strategy {
1249            NegativeStrategy::WrongArticle => {
1250                let anchor_date =
1251                    taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1252                let mut same_date: Vec<Arc<DataRecord>> = self
1253                    .records
1254                    .values()
1255                    .filter(|candidate| {
1256                        candidate.source == anchor_record.source
1257                            && candidate.id != anchor_record.id
1258                            && in_anchor_split(candidate)
1259                    })
1260                    .filter(|candidate| {
1261                        anchor_date
1262                            .as_deref()
1263                            .zip(taxonomy_value(candidate, META_FIELD_DATE))
1264                            .map(|(a, b)| a == b)
1265                            .unwrap_or(false)
1266                    })
1267                    .cloned()
1268                    .collect();
1269                if same_date.is_empty() {
1270                    same_date = self
1271                        .records
1272                        .values()
1273                        .filter(|candidate| {
1274                            candidate.source == anchor_record.source
1275                                && candidate.id != anchor_record.id
1276                                && in_anchor_split(candidate)
1277                        })
1278                        .cloned()
1279                        .collect();
1280                }
1281                if !same_date.is_empty() {
1282                    return self.negative_backend.choose_negative(
1283                        anchor_record,
1284                        anchor_split,
1285                        same_date,
1286                        false,
1287                        anchor_query_text,
1288                        rng,
1289                    );
1290                }
1291                let pool = self
1292                    .records
1293                    .values()
1294                    .filter(|candidate| {
1295                        candidate.id != anchor_record.id && in_anchor_split(candidate)
1296                    })
1297                    .cloned()
1298                    .collect::<Vec<_>>();
1299                self.negative_backend.choose_negative(
1300                    anchor_record,
1301                    anchor_split,
1302                    pool,
1303                    true,
1304                    anchor_query_text,
1305                    rng,
1306                )
1307            }
1308            NegativeStrategy::WrongPublicationDate => {
1309                let anchor_date =
1310                    taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1311                let pool: Vec<Arc<DataRecord>> = self
1312                    .records
1313                    .values()
1314                    .filter(|candidate| {
1315                        candidate.source == anchor_record.source
1316                            && candidate.id != anchor_record.id
1317                            && in_anchor_split(candidate)
1318                    })
1319                    .filter(|candidate| {
1320                        match (
1321                            anchor_date.as_deref(),
1322                            taxonomy_value(candidate, META_FIELD_DATE),
1323                        ) {
1324                            (Some(anchor), Some(candidate_date)) => anchor != candidate_date,
1325                            (Some(_), None) => true,
1326                            (None, Some(_)) => true,
1327                            (None, None) => false,
1328                        }
1329                    })
1330                    .cloned()
1331                    .collect();
1332                if pool.is_empty() {
1333                    // Fallback to any other record in the same split so split boundaries
1334                    // remain strictly isolated.
1335                    let fallback_pool = self
1336                        .records
1337                        .values()
1338                        .filter(|candidate| {
1339                            candidate.id != anchor_record.id && in_anchor_split(candidate)
1340                        })
1341                        .cloned()
1342                        .collect::<Vec<_>>();
1343
1344                    return self.negative_backend.choose_negative(
1345                        anchor_record,
1346                        anchor_split,
1347                        fallback_pool,
1348                        true,
1349                        anchor_query_text,
1350                        rng,
1351                    );
1352                }
1353
1354                self.negative_backend.choose_negative(
1355                    anchor_record,
1356                    anchor_split,
1357                    pool,
1358                    false,
1359                    anchor_query_text,
1360                    rng,
1361                )
1362            }
1363            NegativeStrategy::QuestionAnswerMismatch => {
1364                let pool: Vec<Arc<DataRecord>> = self
1365                    .records
1366                    .values()
1367                    .filter(|candidate| {
1368                        candidate.source == anchor_record.source
1369                            && candidate.id != anchor_record.id
1370                            && in_anchor_split(candidate)
1371                    })
1372                    .cloned()
1373                    .collect();
1374                if pool.is_empty() {
1375                    // Fallback to any other record in the same split so split boundaries
1376                    // remain strictly isolated.
1377                    let fallback_pool = self
1378                        .records
1379                        .values()
1380                        .filter(|candidate| {
1381                            candidate.id != anchor_record.id && in_anchor_split(candidate)
1382                        })
1383                        .cloned()
1384                        .collect::<Vec<_>>();
1385
1386                    return self.negative_backend.choose_negative(
1387                        anchor_record,
1388                        anchor_split,
1389                        fallback_pool,
1390                        true,
1391                        anchor_query_text,
1392                        rng,
1393                    );
1394                }
1395
1396                self.negative_backend.choose_negative(
1397                    anchor_record,
1398                    anchor_split,
1399                    pool,
1400                    false,
1401                    anchor_query_text,
1402                    rng,
1403                )
1404            }
1405            NegativeStrategy::SameRecord => {
1406                // Return the anchor record itself.  The caller
1407                // (finalize_triplet_with_negative) will call select_chunk on
1408                // this record using the recipe's negative_selector, which
1409                // picks from Context sections.  The dedup guard rejects
1410                // triplets where any two slots share identical text.
1411                Some((Arc::new(anchor_record.clone()), false))
1412            }
1413        }
1414    }
1415
1416    /// True when the recipe is the special auto-injected long-section chunk-pair recipe.
1417    fn is_auto_chunk_pair_recipe(recipe: &TripletRecipe) -> bool {
1418        recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME
1419    }
1420
1421    /// Select anchor/positive chunk pair from selectors for a single record.
1422    ///
1423    /// When selectors are the same, retries until a valid pair is found or retry limit is hit.
1424    fn select_anchor_positive_pair(
1425        &mut self,
1426        record: &DataRecord,
1427        anchor_selector: &Selector,
1428        positive_selector: &Selector,
1429        enforce_window_pair: bool,
1430    ) -> Option<(RecordChunk, RecordChunk)> {
1431        let mut anchor_chunk = self.select_chunk(record, anchor_selector)?;
1432        let mut positive_chunk = self.select_chunk(record, positive_selector)?;
1433        if anchor_selector == positive_selector {
1434            let mut retries = 0usize;
1435            while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1436                && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1437            {
1438                let Some(redraw_anchor) = self.select_chunk(record, anchor_selector) else {
1439                    break;
1440                };
1441                let Some(redraw_positive) = self.select_chunk(record, positive_selector) else {
1442                    break;
1443                };
1444                anchor_chunk = redraw_anchor;
1445                positive_chunk = redraw_positive;
1446                retries += 1;
1447            }
1448            if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1449                return None;
1450            }
1451        }
1452        Some((anchor_chunk, positive_chunk))
1453    }
1454
1455    /// Select two distinct *window* chunks for the auto-injected recipe using recipe selectors.
1456    ///
1457    /// This method intentionally uses the recipe's selectors (no hardcoded selector values).
1458    /// For the current auto recipe those selectors are both `Role(Context)`, but the behavior
1459    /// remains tied to recipe configuration, not this helper body.
1460    ///
1461    /// NOTE: this does *not* decide whether the auto recipe is enabled for a source.
1462    /// Source-level enablement happens in `resolve_source_triplet_plan` via
1463    /// `should_auto_inject_chunk_pair_recipe`.
1464    fn select_distinct_window_pair_for_auto_recipe(
1465        &mut self,
1466        recipe: &TripletRecipe,
1467        record: &DataRecord,
1468    ) -> Option<(RecordChunk, RecordChunk)> {
1469        if recipe.anchor != recipe.positive_selector {
1470            return None;
1471        }
1472        self.select_anchor_positive_pair(record, &recipe.anchor, &recipe.positive_selector, true)
1473    }
1474
1475    /// Stage B (record-level): guard for the auto-injected recipe.
1476    ///
1477    /// Requires at least two `ChunkView::Window` candidates across sections addressed by
1478    /// `selector`. This makes the auto recipe eligibility check explicit at execution time.
1479    fn record_has_at_least_two_window_chunks_for_selector(
1480        &self,
1481        record: &DataRecord,
1482        selector: &Selector,
1483    ) -> bool {
1484        let section_indices: Vec<usize> = match selector {
1485            Selector::Role(role) => record
1486                .sections
1487                .iter()
1488                .enumerate()
1489                .filter(|(_, section)| roles_match(role, &section.role))
1490                .map(|(idx, _)| idx)
1491                .collect(),
1492            Selector::Paragraph(idx) => {
1493                if *idx < record.sections.len() {
1494                    vec![*idx]
1495                } else {
1496                    Vec::new()
1497                }
1498            }
1499            Selector::Random => (0..record.sections.len()).collect(),
1500            Selector::TemporalOffset(_) => return false,
1501        };
1502
1503        let mut window_count = 0usize;
1504        for section_idx in section_indices {
1505            let Some(section) = record.sections.get(section_idx) else {
1506                continue;
1507            };
1508            let chunks = self.materialize_chunks(record, section_idx, section);
1509            window_count += chunks
1510                .iter()
1511                .filter(|chunk| matches!(chunk.view, ChunkView::Window { .. }))
1512                .count();
1513            if window_count >= 2 {
1514                return true;
1515            }
1516        }
1517        false
1518    }
1519
1520    /// Serial candidate builder with per-batch negative-text dedup (Option A).
1521    ///
1522    /// Draws the anchor/positive pair once, hoists the 50 % swap + (A) guard
1523    /// out of the retry loop, then redraws **only the negative** up to
1524    /// `NEG_DEDUP_RETRIES` times. The coin-flip swap (one RNG draw, negative
1525    /// unaffected) defeats positional shortcuts so the model treats both
1526    /// positive slots symmetrically. All validation reads use references;
1527    /// ownership of `final_anchor`/`final_positive` moves into the triplet
1528    /// only after a passing negative is found (then return). Text guard (A)
1529    /// (`a == p` when forbidden) is checked before the loop — unrecoverable.
1530    /// Loop checks (B) (`fresh_tuple`, `fresh_neg`, `n == a/p`) on references.
1531    fn make_triplet_with_unique_negative(
1532        &mut self,
1533        recipe: &TripletRecipe,
1534        record: &DataRecord,
1535        rng: &mut DeterministicRng,
1536        seen: &mut HashSet<(RecordKey, RecordKey, RecordKey)>,
1537        seen_negatives: &mut IntSet<u64>,
1538    ) -> Option<SampleTriplet> {
1539        // Draw anchor/positive once (serial cursor semantics).
1540        let (mut anchor_chunk, mut positive_chunk) = if Self::is_auto_chunk_pair_recipe(recipe) {
1541            if !self.record_has_at_least_two_window_chunks_for_selector(record, &recipe.anchor) {
1542                return None;
1543            }
1544            self.select_distinct_window_pair_for_auto_recipe(recipe, record)?
1545        } else {
1546            self.select_anchor_positive_pair(
1547                record,
1548                &recipe.anchor,
1549                &recipe.positive_selector,
1550                false,
1551            )?
1552        };
1553        let anchor_raw_text = anchor_chunk.text.clone();
1554        self.decorate_chunk(record, &mut anchor_chunk, rng);
1555        self.decorate_chunk(record, &mut positive_chunk, rng);
1556
1557        // Hoisted swap (single RNG draw) + guard (A).
1558        let (final_anchor, final_positive) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0 {
1559            (positive_chunk, anchor_chunk)
1560        } else {
1561            (anchor_chunk, positive_chunk)
1562        };
1563        if !recipe.allow_same_anchor_positive && final_anchor.text == final_positive.text {
1564            return None;
1565        }
1566
1567        // Negative-only retries; references only, no anchor/positive clones.
1568        let mut accepted: Option<(RecordChunk, f32, bool)> = None;
1569        for _ in 0..NEG_DEDUP_RETRIES {
1570            let (negative_record, fallback_used) = self.select_negative_record(
1571                record,
1572                &recipe.negative_strategy,
1573                Some(&anchor_raw_text),
1574                rng,
1575            )?;
1576            let mut negative_chunk =
1577                match self.select_chunk(&negative_record, &recipe.negative_selector) {
1578                    Some(c) => c,
1579                    None => continue,
1580                };
1581            self.decorate_chunk(&negative_record, &mut negative_chunk, rng);
1582            if negative_chunk.text == final_positive.text
1583                || negative_chunk.text == final_anchor.text
1584            {
1585                continue;
1586            }
1587            let key = (
1588                RecordKey::of_chunk(&final_anchor),
1589                RecordKey::of_chunk(&final_positive),
1590                RecordKey::of_chunk(&negative_chunk),
1591            );
1592            let neg_h = negative_text_hash(&negative_chunk);
1593            // Atomic check-then-commit: no partial `seen` mutation on reject.
1594            if seen.contains(&key) || seen_negatives.contains(&neg_h) {
1595                continue;
1596            }
1597            let chunk_weight =
1598                self.triplet_chunk_weight(recipe, &final_anchor, &final_positive, &negative_chunk);
1599            let weight = recipe.weight * chunk_weight;
1600            seen.insert(key);
1601            seen_negatives.insert(neg_h);
1602            accepted = Some((negative_chunk, weight, fallback_used));
1603            break;
1604        }
1605        let (negative_chunk, weight, fallback_used) = accepted?;
1606        let recipe_name = if fallback_used {
1607            format!("{}_fallback_same_split", recipe.name)
1608        } else {
1609            recipe.name.to_string()
1610        };
1611        Some(SampleTriplet {
1612            recipe: recipe_name,
1613            anchor: final_anchor,
1614            positive: final_positive,
1615            negative: negative_chunk,
1616            weight,
1617            instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1618        })
1619    }
1620
1621    fn make_text_sample_for_split(
1622        &mut self,
1623        recipe: &TextRecipe,
1624        source: Option<&str>,
1625        split: SplitLabel,
1626        rng: &mut DeterministicRng,
1627    ) -> Option<TextSample> {
1628        let record = self.choose_anchor_record(source, split)?;
1629        let mut chunk = self.select_chunk(&record, &recipe.selector)?;
1630        self.decorate_chunk(&record, &mut chunk, rng);
1631        let weight = recipe.weight * self.chunk_weight(&chunk);
1632        Some(TextSample {
1633            recipe: recipe.name.to_string(),
1634            chunk,
1635            weight,
1636            instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1637        })
1638    }
1639
1640    fn chunk_weight(&self, chunk: &RecordChunk) -> f32 {
1641        chunk_weight(&self.config.chunking, chunk)
1642    }
1643
1644    fn triplet_chunk_weight(
1645        &self,
1646        recipe: &TripletRecipe,
1647        anchor: &RecordChunk,
1648        positive: &RecordChunk,
1649        negative: &RecordChunk,
1650    ) -> f32 {
1651        let floor = self.config.chunking.chunk_weight_floor;
1652        let negative_weight = negative.quality.trust.clamp(0.0, 1.0).max(floor);
1653        if Self::is_auto_chunk_pair_recipe(recipe) {
1654            // For the auto long-section recipe, use one coherence signal (proximity)
1655            // for the anchor/positive pair.
1656            let pair_trust = ((anchor.quality.trust.clamp(0.0, 1.0)
1657                + positive.quality.trust.clamp(0.0, 1.0))
1658                / 2.0)
1659                .clamp(0.0, 1.0);
1660            let pair_weight = (chunk_proximity_score(anchor, positive) * pair_trust).max(floor);
1661            // Keep negative weighting simple in this recipe: trust + floor only.
1662            return (pair_weight + pair_weight + negative_weight) / 3.0;
1663        }
1664        // Non-auto recipes also apply anchor-positive proximity, while keeping
1665        // negative weighting trust-only.
1666        let pair_proximity = chunk_proximity_score(anchor, positive);
1667        let anchor_weight = (self.chunk_weight(anchor) * pair_proximity).max(floor);
1668        let positive_weight = (self.chunk_weight(positive) * pair_proximity).max(floor);
1669        (anchor_weight + positive_weight + negative_weight) / 3.0
1670    }
1671
1672    fn decorate_chunk(
1673        &mut self,
1674        record: &DataRecord,
1675        chunk: &mut RecordChunk,
1676        rng: &mut DeterministicRng,
1677    ) {
1678        chunk.kvp_meta = record
1679            .meta_prefix
1680            .as_ref()
1681            .map(|s| s.all_metadata())
1682            .unwrap_or_default();
1683        if let Some(spec) = record.meta_prefix.as_ref()
1684            && let Some(prefix) = spec.sample(rng)
1685        {
1686            let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1687            let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1688            let total_tokens = prefix_tokens.len() + body_tokens.len();
1689            let max_window = self.config.chunking.max_window_tokens;
1690            if max_window > 0 && total_tokens > max_window {
1691                if prefix_tokens.len() >= max_window {
1692                    chunk.text = prefix_tokens
1693                        .into_iter()
1694                        .take(max_window)
1695                        .collect::<Vec<_>>()
1696                        .join(" ");
1697                    chunk.tokens_estimate = max_window;
1698                } else {
1699                    let remaining = max_window - prefix_tokens.len();
1700                    let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1701                    chunk.text =
1702                        format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1703                    chunk.tokens_estimate = max_window;
1704                }
1705            } else {
1706                chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1707                chunk.tokens_estimate = total_tokens;
1708            }
1709        }
1710    }
1711
1712    // ── parallel-batch helpers (&self, rng-based, no cursor maps) ────────────
1713
1714    /// Select a chunk using rng instead of cursor maps (for parallel execution).
1715    fn select_chunk_parallel(
1716        &self,
1717        record: &DataRecord,
1718        selector: &Selector,
1719        rng: &mut DeterministicRng,
1720    ) -> Option<RecordChunk> {
1721        match selector {
1722            Selector::Role(role) => self.select_role_parallel(record, role, rng),
1723            Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1724                let pool = self.materialize_chunks(record, *idx, section);
1725                if pool.is_empty() {
1726                    return None;
1727                }
1728                let i = rng.random_range(0..pool.len());
1729                pool.into_iter().nth(i)
1730            }),
1731            Selector::TemporalOffset(offset) => self
1732                .select_temporal_neighbor(record, *offset)
1733                .and_then(|neighbor| {
1734                    self.select_role_parallel(&neighbor, &SectionRole::Context, rng)
1735                }),
1736            Selector::Random => {
1737                if record.sections.is_empty() {
1738                    return None;
1739                }
1740                let idx = rng.random_range(0..record.sections.len());
1741                record.sections.get(idx).and_then(|section| {
1742                    let pool = self.materialize_chunks(record, idx, section);
1743                    if pool.is_empty() {
1744                        return None;
1745                    }
1746                    let i = rng.random_range(0..pool.len());
1747                    pool.into_iter().nth(i)
1748                })
1749            }
1750        }
1751    }
1752
1753    /// Select a chunk by role using rng (parallel path — no role_cursors written).
1754    fn select_role_parallel(
1755        &self,
1756        record: &DataRecord,
1757        role: &SectionRole,
1758        rng: &mut DeterministicRng,
1759    ) -> Option<RecordChunk> {
1760        let indices: Vec<usize> = record
1761            .sections
1762            .iter()
1763            .enumerate()
1764            .filter(|(_, s)| roles_match(role, &s.role))
1765            .map(|(i, _)| i)
1766            .collect();
1767        if indices.is_empty() {
1768            return None;
1769        }
1770        let start = rng.random_range(0..indices.len());
1771        for offset in 0..indices.len() {
1772            let section_idx = indices[(start + offset) % indices.len()];
1773            let section = &record.sections[section_idx];
1774            let pool = self.materialize_chunks(record, section_idx, section);
1775            if !pool.is_empty() {
1776                let i = rng.random_range(0..pool.len());
1777                return pool.into_iter().nth(i);
1778            }
1779        }
1780        None
1781    }
1782
1783    /// Decorate a chunk using a provided rng (parallel-safe, &self).
1784    fn decorate_chunk_parallel(
1785        &self,
1786        record: &DataRecord,
1787        chunk: &mut RecordChunk,
1788        rng: &mut DeterministicRng,
1789    ) {
1790        chunk.kvp_meta = record
1791            .meta_prefix
1792            .as_ref()
1793            .map(|s| s.all_metadata())
1794            .unwrap_or_default();
1795        if let Some(spec) = record.meta_prefix.as_ref()
1796            && let Some(prefix) = spec.sample(rng)
1797        {
1798            let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1799            let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1800            let total_tokens = prefix_tokens.len() + body_tokens.len();
1801            let max_window = self.config.chunking.max_window_tokens;
1802            if max_window > 0 && total_tokens > max_window {
1803                if prefix_tokens.len() >= max_window {
1804                    chunk.text = prefix_tokens
1805                        .into_iter()
1806                        .take(max_window)
1807                        .collect::<Vec<_>>()
1808                        .join(" ");
1809                    chunk.tokens_estimate = max_window;
1810                } else {
1811                    let remaining = max_window - prefix_tokens.len();
1812                    let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1813                    chunk.text =
1814                        format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1815                    chunk.tokens_estimate = max_window;
1816                }
1817            } else {
1818                chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1819                chunk.tokens_estimate = total_tokens;
1820            }
1821        }
1822    }
1823
1824    /// Select an anchor/positive pair for the parallel path (rng-based).
1825    fn select_anchor_positive_parallel(
1826        &self,
1827        record: &DataRecord,
1828        anchor_selector: &Selector,
1829        positive_selector: &Selector,
1830        enforce_window_pair: bool,
1831        rng: &mut DeterministicRng,
1832    ) -> Option<(RecordChunk, RecordChunk)> {
1833        let anchor_chunk = self.select_chunk_parallel(record, anchor_selector, rng)?;
1834        let mut positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1835        if anchor_selector == positive_selector {
1836            let mut retries = 0usize;
1837            while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1838                && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1839            {
1840                positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1841                retries += 1;
1842            }
1843            if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1844                return None;
1845            }
1846        }
1847        // No swap here — finalize_triplet_parallel applies the coin-flip swap.
1848        Some((anchor_chunk, positive_chunk))
1849    }
1850
1851    /// Select (anchor_chunk, positive_chunk, anchor_raw_text) for a recipe — parallel-safe.
1852    /// Does NOT select a negative; that happens sequentially in Phase 3.
1853    fn select_anchor_positive_for_recipe(
1854        &self,
1855        recipe: &TripletRecipe,
1856        anchor_record: &DataRecord,
1857        rng: &mut DeterministicRng,
1858    ) -> Option<(RecordChunk, RecordChunk, String)> {
1859        if Self::is_auto_chunk_pair_recipe(recipe) {
1860            if !self
1861                .record_has_at_least_two_window_chunks_for_selector(anchor_record, &recipe.anchor)
1862            {
1863                return None;
1864            }
1865            let mut anchor_chunk =
1866                self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1867            let mut positive_chunk =
1868                self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1869            let mut tries = 0usize;
1870            while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, true) {
1871                tries += 1;
1872                if tries >= SAME_SELECTOR_PAIR_RETRY_LIMIT {
1873                    return None;
1874                }
1875                anchor_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1876                positive_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1877            }
1878            let anchor_raw_text = anchor_chunk.text.clone();
1879            self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1880            self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1881            return Some((anchor_chunk, positive_chunk, anchor_raw_text));
1882        }
1883        let (mut anchor_chunk, mut positive_chunk) = self.select_anchor_positive_parallel(
1884            anchor_record,
1885            &recipe.anchor,
1886            &recipe.positive_selector,
1887            false,
1888            rng,
1889        )?;
1890        let anchor_raw_text = anchor_chunk.text.clone();
1891        self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1892        self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1893        Some((anchor_chunk, positive_chunk, anchor_raw_text))
1894    }
1895
1896    fn select_chunk(&mut self, record: &DataRecord, selector: &Selector) -> Option<RecordChunk> {
1897        match selector {
1898            Selector::Role(role) => self.select_by_role(record, role),
1899            Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1900                let pool = self.materialize_chunks(record, *idx, section);
1901                self.next_chunk_from_pool(&RecordKey::from(record), *idx, pool)
1902            }),
1903            Selector::TemporalOffset(offset) => self
1904                .select_temporal_neighbor(record, *offset)
1905                .and_then(|neighbor| self.select_by_role(&neighbor, &SectionRole::Context)),
1906            Selector::Random => {
1907                if record.sections.is_empty() {
1908                    return None;
1909                }
1910                let idx = self.rng.random_range(0..record.sections.len());
1911                record.sections.get(idx).and_then(|section| {
1912                    let pool = self.materialize_chunks(record, idx, section);
1913                    self.next_chunk_from_pool(&RecordKey::from(record), idx, pool)
1914                })
1915            }
1916        }
1917    }
1918
1919    fn select_by_role(&mut self, record: &DataRecord, role: &SectionRole) -> Option<RecordChunk> {
1920        let indices: Vec<usize> = record
1921            .sections
1922            .iter()
1923            .enumerate()
1924            .filter(|(_, section)| roles_match(role, &section.role))
1925            .map(|(idx, _)| idx)
1926            .collect();
1927        if indices.is_empty() {
1928            return None;
1929        }
1930        let key = role_cursor_key(&RecordKey::from(record), role);
1931        let start_offset = self
1932            .role_cursors
1933            .get(&key)
1934            .and_then(|last_idx| indices.iter().position(|idx| idx == last_idx))
1935            .map(|pos| (pos + 1) % indices.len())
1936            .unwrap_or_else(|| {
1937                // On first use for this (record,role), choose a deterministic hashed
1938                // section offset instead of always starting at the first matching section.
1939                // This avoids systematic head-bias while preserving reproducibility.
1940                // Preimage replicates the historical "{id}::{role}" form so
1941                // sequences are unchanged for existing corpora.
1942                let record_key = RecordKey::from(record);
1943                let seed_key = format!("{}::{}", record_key.id, role_label(role));
1944                (stable_hash_str(self.epoch_seed(), &seed_key) as usize) % indices.len()
1945            });
1946        for offset in 0..indices.len() {
1947            let section_idx = indices[(start_offset + offset) % indices.len()];
1948            let section = &record.sections[section_idx];
1949            let pool = self.materialize_chunks(record, section_idx, section);
1950            if let Some(chunk) =
1951                self.next_chunk_from_pool(&RecordKey::from(record), section_idx, pool)
1952            {
1953                self.role_cursors.insert(key.clone(), section_idx);
1954                return Some(chunk);
1955            }
1956        }
1957        None
1958    }
1959
1960    /// Materialize chunk windows for one section according to `ChunkingStrategy`.
1961    ///
1962    /// Window layout algorithm:
1963    /// 1) Tokenize section text with `split_whitespace`.
1964    /// 2) Set `span = min(max_window_tokens, total_tokens)`.
1965    /// 3) If `span == total_tokens`, emit one full-section `ChunkView::Window`.
1966    /// 4) Otherwise, for each configured overlap:
1967    ///    - `stride = max(1, span - overlap)`
1968    ///    - emit windows `[start, min(start + span, total_tokens))` while advancing by `stride`.
1969    /// 5) Optionally append one `SummaryFallback` chunk when section is longer than `max_window_tokens`.
1970    ///
1971    /// This is the source of truth for how chunk windows are "laid" in a record.
1972    fn materialize_chunks(
1973        &self,
1974        record: &DataRecord,
1975        section_idx: usize,
1976        section: &RecordSection,
1977    ) -> Vec<RecordChunk> {
1978        let mut chunks =
1979            self.chunker
1980                .materialize(&self.config.chunking, record, section_idx, section);
1981        // Propagate the section's precomputed embedding (if any) into every
1982        // chunk materialized from it. The chunker stays agnostic; Arc makes
1983        // this a pointer copy, never a vector copy.
1984        for chunk in &mut chunks {
1985            chunk.embedding = section.embedding.clone();
1986            chunk.source = record.source.clone();
1987        }
1988        chunks
1989    }
1990
1991    fn build_derived_text_recipes(recipes: &[TripletRecipe]) -> Vec<TextRecipe> {
1992        let mut derived = Vec::new();
1993        for recipe in recipes {
1994            let base = recipe.name.as_ref();
1995            derived.push(TextRecipe {
1996                name: Cow::Owned(format!("{base}_anchor")),
1997                selector: recipe.anchor.clone(),
1998                weight: recipe.weight,
1999                instruction: None,
2000            });
2001            derived.push(TextRecipe {
2002                name: Cow::Owned(format!("{base}_positive")),
2003                selector: recipe.positive_selector.clone(),
2004                weight: recipe.weight,
2005                instruction: None,
2006            });
2007            derived.push(TextRecipe {
2008                name: Cow::Owned(format!("{base}_negative")),
2009                selector: recipe.negative_selector.clone(),
2010                weight: recipe.weight,
2011                instruction: None,
2012            });
2013        }
2014        derived
2015    }
2016
2017    /// Stage A helper: true when record contains an Anchor/Context section whose token
2018    /// count exceeds `chunking.max_window_tokens`.
2019    fn record_has_long_anchor_or_context_section(&self, record: &DataRecord) -> bool {
2020        let window = self.config.chunking.max_window_tokens;
2021        if window == 0 {
2022            return false;
2023        }
2024        record.sections.iter().any(|s| {
2025            matches!(s.role, SectionRole::Anchor | SectionRole::Context) && s.token_count > window
2026        })
2027    }
2028    fn sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
2029        if self.last_sync_version == 0 {
2030            self.full_sync_records_from_cache()?;
2031            return Ok(());
2032        }
2033
2034        // Force refreshes replace cache contents wholesale; fall back to a
2035        // full pool resync with exact full-snapshot semantics (evicted =
2036        // pool minus snapshot, dedup pruned only for records that left).
2037        // Steady-state advances keep the O(Δ) delta fast path below.
2038        if self.ingestion.force_refresh_generation() != self.last_force_refresh_generation {
2039            self.last_force_refresh_generation = self.ingestion.force_refresh_generation();
2040            self.full_sync_records_from_cache()?;
2041            return Ok(());
2042        }
2043
2044        // Use sync_delta to get only added and evicted records.
2045        // No full snapshot, no sort, no HashSet diff — O(ΔN) path.
2046        let (new_version, added, evicted) = self.ingestion.sync_delta(self.last_sync_version);
2047
2048        if added.is_empty() && evicted.is_empty() {
2049            self.negative_backend.on_sync_start();
2050            self.prune_cursor_state();
2051            self.last_sync_version = new_version;
2052            return Ok(());
2053        }
2054
2055        // --- Process evictions (strict O(Δ)): prune pool state and the
2056        // per-source index only for IDs that truly left. IDs cleared and
2057        // immediately re-added in this same sync (e.g. force-refresh
2058        // returning the same records) keep their cross-batch dedup and
2059        // split-label state. This preserves the
2060        // `emitted_text_hashes_blocks_repeats` contract while still freeing
2061        // state for evicted records.
2062        let mut truly_evicted = false;
2063        if !evicted.is_empty() {
2064            let added_keys: HashSet<RecordKey> = added.iter().map(RecordKey::of_arc).collect();
2065            for evicted_key in &evicted {
2066                if added_keys.contains(evicted_key) {
2067                    continue;
2068                }
2069                let Some(record) = self.records.get(evicted_key) else {
2070                    continue;
2071                };
2072                let source = record.source.clone();
2073                let is_long = self.record_has_long_anchor_or_context_section(record);
2074                self.records.shift_remove(evicted_key);
2075                self.emitted_text_hashes.remove(evicted_key);
2076                self.split_labels.remove(evicted_key);
2077                self.chunk_index.remove(evicted_key);
2078                self.remove_id_from_source_index(&source, evicted_key);
2079                if is_long {
2080                    self.decrement_long_section_count(&source);
2081                }
2082                truly_evicted = true;
2083            }
2084        }
2085
2086        self.negative_backend.on_sync_start();
2087
2088        // --- Process additions (strict O(Δ log N) via binary search) ---
2089        // `pool_changed` tracks actual membership change: genuinely new IDs
2090        // or true evictions. Pure version refreshes of already-pooled IDs
2091        // leave the pool (and the cached records_by_split map) intact.
2092        let mut pool_changed = truly_evicted;
2093        if !added.is_empty() {
2094            let allowed: HashSet<SplitLabel> = self.allowed_target_splits().into_iter().collect();
2095            for record in &added {
2096                // Re-ingested keys (version bump for records already pooled,
2097                // e.g. a wrapped source re-draining known records) must not
2098                // create duplicate index entries or double-count long
2099                // sections. A full rebuild derives each key exactly once from
2100                // the pool, so skip index/count work here and only refresh
2101                // the stored Arc.
2102                let key = RecordKey::from(record);
2103                if self.records.contains_key(&key) {
2104                    self.records.insert(key, Arc::clone(record));
2105                    continue;
2106                }
2107                pool_changed = true;
2108                let label = self.get_or_insert_split_label(&key)?;
2109                if self.record_has_long_anchor_or_context_section(record) {
2110                    self.increment_long_section_count(&record.source);
2111                }
2112                self.records.insert(key.clone(), Arc::clone(record));
2113                self.chunk_index.insert(key.clone(), key.clone());
2114                self.insert_id_into_source_index(record.source.clone(), key, label, &allowed);
2115            }
2116        }
2117
2118        // --- Refresh derived source state (mirrors `rebuild_source_index` tail) ---
2119        // Drop index entries for sources left with no allowed-split records,
2120        // then refresh ordering, wrap flags, and cursors exactly as a full
2121        // rebuild would.
2122        let empty_sources: Vec<SourceId> = self
2123            .source_record_indices
2124            .iter()
2125            .filter(|(_, ids)| ids.is_empty())
2126            .map(|(source, _)| source.clone())
2127            .collect();
2128        let mut order_changed = !empty_sources.is_empty();
2129        for source in &empty_sources {
2130            self.source_record_indices.remove(source);
2131        }
2132        // New sources may have appeared via the insertions above.
2133        if !order_changed {
2134            for source in self.source_record_indices.keys() {
2135                if !self.source_order.contains(source) {
2136                    order_changed = true;
2137                    break;
2138                }
2139            }
2140        }
2141        if order_changed {
2142            self.source_order = self.source_record_indices.keys().cloned().collect();
2143            self.source_order.sort();
2144        }
2145        // Drop deficit counters for sources that left the pool.
2146        self.source_deficits
2147            .retain(|s, _| self.source_record_indices.contains_key(s));
2148        self.refresh_source_wrapped();
2149        self.source_record_cursors
2150            .retain(|source, _| self.source_record_indices.contains_key(source));
2151        if self.source_state_loaded {
2152            if self.source_order.is_empty() {
2153                self.source_cycle_idx = 0;
2154            }
2155            self.source_state_dirty = self.source_order.len() > 1;
2156        }
2157        // Cursor pruning is a no-op unless the valid pool shrank.
2158        if truly_evicted {
2159            self.prune_cursor_state();
2160        }
2161
2162        if pool_changed {
2163            self.pool_generation = self.pool_generation.saturating_add(1);
2164        }
2165        self.last_sync_version = new_version;
2166        Ok(())
2167    }
2168
2169    /// Full snapshot sync: used only for the initial bootstrap (last_sync_version == 0).
2170    /// Inserts records in ID-sorted order so subsequent delta syncs can append
2171    /// new records in sorted order without breaking the IndexMap ordering.
2172    fn full_sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
2173        let mut snapshot = self.ingestion.all_records_snapshot();
2174        snapshot.sort_by_key(RecordKey::of_arc);
2175
2176        // Determine which records are being evicted BEFORE clearing self.records.
2177        let old_keys: HashSet<RecordKey> = self.records.keys().cloned().collect();
2178        let incoming_keys: HashSet<RecordKey> = snapshot.iter().map(RecordKey::of_arc).collect();
2179        let evicted: HashSet<&RecordKey> = old_keys.difference(&incoming_keys).collect();
2180
2181        self.records.clear();
2182        self.long_section_counts.clear();
2183        for evicted_key in &evicted {
2184            self.emitted_text_hashes.remove(*evicted_key);
2185            self.split_labels.remove(*evicted_key);
2186        }
2187        self.negative_backend.on_sync_start();
2188        for record in snapshot {
2189            let key = RecordKey::from(&record);
2190            let _ = self.get_or_insert_split_label(&key)?;
2191            if self.record_has_long_anchor_or_context_section(&record) {
2192                self.increment_long_section_count(&record.source);
2193            }
2194            self.records.insert(key, Arc::clone(&record));
2195        }
2196        self.prune_cursor_state();
2197        self.rebuild_chunk_index();
2198        self.rebuild_source_index()?;
2199        self.last_sync_version = self.ingestion.global_version();
2200        // The pool was fully replaced; invalidate the cached split map.
2201        self.pool_generation = self.pool_generation.saturating_add(1);
2202        Ok(())
2203    }
2204
2205    fn ingest_internal_for_split(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
2206        if !self.ingestion.has_sources() {
2207            return Ok(());
2208        }
2209        if !self.ingestion_cursors_loaded {
2210            if let Some(state) = self.split_store.load_sampler_state()? {
2211                self.ingestion.load_cursors(&state.source_stream_cursors);
2212                self.ingestion.set_epoch(state.epoch);
2213                self.ingestion.set_epoch_step(state.epoch_step);
2214            }
2215            self.ingestion_cursors_loaded = true;
2216        }
2217        if self.ingestion.all_caches_empty() {
2218            self.ingestion.refresh_all();
2219        } else {
2220            self.ingestion.advance(self.config.batch_size);
2221        }
2222        let observed = self.ingestion.total_ingest_count();
2223        if observed == self.last_observed_ingest && !self.records.is_empty() {
2224            return Ok(());
2225        }
2226        self.last_observed_ingest = observed;
2227        self.sync_records_from_cache()?;
2228        let max_window_tokens = self.config.chunking.max_window_tokens;
2229        self.negative_backend.on_records_refreshed(
2230            &self.records,
2231            max_window_tokens,
2232            &|id| self.split_labels.get(id).copied(),
2233            self.ingestion.last_refreshed_sources(),
2234        );
2235        // Epoch tracking and source-state management must run every batch regardless
2236        // of whether the record pool changed — reconcile advances the sampling cursor.
2237        self.epoch_tracker.ensure_loaded()?;
2238        self.refresh_records_by_split_cache()?;
2239        let cached_records_by_split = self
2240            .cached_records_by_split
2241            .as_ref()
2242            .expect("records_by_split cache populated by refresh");
2243        self.epoch_tracker
2244            .reconcile(target_split, cached_records_by_split);
2245        self.ensure_source_state()?;
2246        Ok(())
2247    }
2248
2249    /// Ensure ingestion cursors (including epoch step) are loaded from the
2250    /// persisted state so the step counter reflects the saved value before
2251    /// the wrapper increments it for the current call.
2252    ///
2253    /// Only loads ingestion cursors and source epoch.  The sampling state
2254    /// (RNG, source_record_cursors, …) is restored later by
2255    /// `ingest_internal_for_split` → `ensure_source_state`, after
2256    /// `sync_records_from_cache` has populated `source_record_indices`.
2257    fn ensure_ingestion_cursors_loaded(&mut self) -> Result<(), SamplerError> {
2258        if !self.ingestion_cursors_loaded {
2259            if let Some(state) = self.split_store.load_sampler_state()? {
2260                self.ingestion.load_cursors(&state.source_stream_cursors);
2261                self.ingestion.set_epoch(state.epoch);
2262                self.ingestion.set_epoch_step(state.epoch_step);
2263            }
2264            self.ingestion_cursors_loaded = true;
2265        }
2266        Ok(())
2267    }
2268
2269    #[cfg(test)]
2270    fn ingest_internal(&mut self, split: SplitLabel) -> Result<(), SamplerError> {
2271        self.ingest_internal_for_split(split)
2272    }
2273
2274    fn ingest_internal_with_weights_for_split(
2275        &mut self,
2276        target_split: SplitLabel,
2277        weights: &HashMap<SourceId, f32>,
2278    ) -> Result<(), SamplerError> {
2279        if !self.ingestion.has_sources() {
2280            return Ok(());
2281        }
2282        if !self.ingestion_cursors_loaded {
2283            if let Some(state) = self.split_store.load_sampler_state()? {
2284                self.ingestion.load_cursors(&state.source_stream_cursors);
2285                self.ingestion.set_epoch(state.epoch);
2286                self.ingestion.set_epoch_step(state.epoch_step);
2287            }
2288            self.ingestion_cursors_loaded = true;
2289        }
2290        if self.ingestion.all_caches_empty() {
2291            self.ingestion.refresh_all_with_weights(weights)?;
2292        } else {
2293            self.ingestion
2294                .advance_with_weights(self.config.batch_size, weights)?;
2295        }
2296        let observed = self.ingestion.total_ingest_count();
2297        if observed == self.last_observed_ingest && !self.records.is_empty() {
2298            return Ok(());
2299        }
2300        self.last_observed_ingest = observed;
2301        self.sync_records_from_cache()?;
2302        let max_window_tokens = self.config.chunking.max_window_tokens;
2303        self.negative_backend.on_records_refreshed(
2304            &self.records,
2305            max_window_tokens,
2306            &|id| self.split_labels.get(id).copied(),
2307            self.ingestion.last_refreshed_sources(),
2308        );
2309        self.epoch_tracker.ensure_loaded()?;
2310        self.refresh_records_by_split_cache()?;
2311        let cached_records_by_split = self
2312            .cached_records_by_split
2313            .as_ref()
2314            .expect("records_by_split cache populated by refresh");
2315        self.epoch_tracker
2316            .reconcile(target_split, cached_records_by_split);
2317        self.ensure_source_state()?;
2318        Ok(())
2319    }
2320
2321    /// Helper that centralises the weight-fallback logic.
2322    ///
2323    /// If `weights` is `Some` and non-uniform or if `weights` is `None`,
2324    /// delegates to the appropriate ingest method.
2325    fn ingest_with_weights_fallback(
2326        &mut self,
2327        target_split: SplitLabel,
2328        weights: Option<&HashMap<SourceId, f32>>,
2329    ) -> Result<(), SamplerError> {
2330        match weights {
2331            Some(weights)
2332                if !weights.is_empty()
2333                    && !weights
2334                        .values()
2335                        .all(|&w| w == *weights.values().next().unwrap()) =>
2336            {
2337                self.ingest_internal_with_weights_for_split(target_split, weights)?
2338            }
2339            _ => self.ingest_internal_for_split(target_split)?,
2340        }
2341        Ok(())
2342    }
2343
2344    fn force_ingest_refresh_with_weights_for_split(
2345        &mut self,
2346        target_split: SplitLabel,
2347        weights: &HashMap<SourceId, f32>,
2348    ) -> Result<(), SamplerError> {
2349        if !self.ingestion.has_sources() {
2350            return Ok(());
2351        }
2352        if !self.ingestion_cursors_loaded {
2353            if let Some(state) = self.split_store.load_sampler_state()? {
2354                self.ingestion.load_cursors(&state.source_stream_cursors);
2355                self.ingestion.set_epoch(state.epoch);
2356                self.ingestion.set_epoch_step(state.epoch_step);
2357            }
2358            self.ingestion_cursors_loaded = true;
2359        }
2360        self.ingestion.force_refresh_all_with_weights(weights)?;
2361        self.last_observed_ingest = self.ingestion.total_ingest_count();
2362        self.sync_records_from_cache()?;
2363        let max_window_tokens = self.config.chunking.max_window_tokens;
2364        self.negative_backend.on_records_refreshed(
2365            &self.records,
2366            max_window_tokens,
2367            &|id| self.split_labels.get(id).copied(),
2368            self.ingestion.last_refreshed_sources(),
2369        );
2370        self.epoch_tracker.ensure_loaded()?;
2371        self.refresh_records_by_split_cache()?;
2372        let cached_records_by_split = self
2373            .cached_records_by_split
2374            .as_ref()
2375            .expect("records_by_split cache populated by refresh");
2376        self.epoch_tracker
2377            .reconcile(target_split, cached_records_by_split);
2378        self.ensure_source_state()?;
2379        Ok(())
2380    }
2381
2382    /// Select one triplet candidate for a specific source and split.
2383    ///
2384    /// This helper is the source-level recipe execution path used by both pair and
2385    /// triplet batch builders. It always starts from `triplet_recipes_for_source`,
2386    /// which already includes any auto-injected long-section recipe when eligible.
2387    ///
2388    /// Returns:
2389    /// - candidate triplet (with recipe) if one was sampled,
2390    /// - number of recipe attempts consumed (for round-robin bookkeeping).
2391    ///
2392    /// Takes the batch `seen` sets so negative-text dedup commits atomically.
2393    #[allow(clippy::too_many_arguments)]
2394    fn sample_source_triplet_candidate(
2395        &mut self,
2396        source: &str,
2397        target_split: SplitLabel,
2398        recipe_orders: &mut HashMap<RecipeKey, Vec<usize>>,
2399        recipe_positions: &mut HashMap<RecipeKey, usize>,
2400        rng: &mut DeterministicRng,
2401        seen: &mut HashSet<(RecordKey, RecordKey, RecordKey)>,
2402        seen_negatives: &mut IntSet<u64>,
2403    ) -> (Option<(TripletRecipe, SampleTriplet)>, usize) {
2404        // Stage A (source-level injection): resolve effective recipe pool,
2405        // including auto-injected long-section recipe when source is eligible.
2406        let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
2407        if recipes.is_empty() {
2408            return (None, 0);
2409        }
2410        if !recipe_orders.contains_key(source) {
2411            let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2412            let order =
2413                self.recipe_order_weighted_cycled(&recipe_weights, self.triplet_recipe_rr_idx, rng);
2414            recipe_orders.insert(source.to_string(), order);
2415        }
2416        let order = recipe_orders
2417            .get(source)
2418            .expect("recipe order missing for source");
2419        let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2420        let Some(anchor) = self.choose_anchor_record(Some(source), target_split) else {
2421            return (None, 0);
2422        };
2423
2424        let mut attempts = 0usize;
2425        for offset in 0..order.len() {
2426            let idx = order[(*pos + offset) % order.len()];
2427            attempts = attempts.saturating_add(1);
2428            let recipe = recipes[idx].clone();
2429            // Stage B/C happen inside the unique-negative builder, which draws
2430            // anchor/positive once then retries only the negative (per-batch
2431            // negative-text dedup, best-effort pre-pad). Commit is atomic
2432            // inside the helper.
2433            if let Some(sample) =
2434                self.make_triplet_with_unique_negative(&recipe, &anchor, rng, seen, seen_negatives)
2435            {
2436                *pos = (*pos + offset + 1) % order.len();
2437                return (Some((recipe, sample)), attempts);
2438            }
2439        }
2440
2441        (None, attempts)
2442    }
2443
2444    fn next_pair_batch_inner_with_weights(
2445        &mut self,
2446        target_split: SplitLabel,
2447        weights: Option<&HashMap<SourceId, f32>>,
2448    ) -> Result<SampleBatch, SamplerError> {
2449        self.ingest_with_weights_fallback(target_split, weights)?;
2450        self.ensure_split_has_records(target_split)?;
2451        let sources = self.source_order.clone();
2452        if sources.is_empty() {
2453            if self.triplet_recipes.is_empty() {
2454                return Err(SamplerError::Configuration(
2455                    "no triplet recipes available".into(),
2456                ));
2457            }
2458            let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2459            let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2460            let recipe_order = self.recipe_order_weighted_cycled(
2461                &recipe_weights,
2462                self.triplet_recipe_rr_idx,
2463                &mut rng,
2464            );
2465            let mut pairs = Vec::new();
2466            let mut seen = HashSet::new();
2467            // Per-batch negative-text dedup (Option A, best-effort pre-pad).
2468            // `pad_with_reuse` below may still duplicate when the unique pool
2469            // is smaller than `batch_size`.
2470            let mut seen_negatives: IntSet<u64> = IntSet::default();
2471            let mut last_recipe_name = None;
2472            let mut recipe_pos = 0usize;
2473            let mut recipe_steps = 0usize;
2474            let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2475            for _ in 0..attempts {
2476                if pairs.len() >= self.config.batch_size {
2477                    break;
2478                }
2479                let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2480                    break;
2481                };
2482                let mut triplet = None;
2483                for offset in 0..recipe_order.len() {
2484                    let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2485                    recipe_steps = recipe_steps.saturating_add(1);
2486                    let recipe = self.triplet_recipes[idx].clone();
2487                    last_recipe_name = Some(recipe.name.clone());
2488                    if let Some(sample) = self.make_triplet_with_unique_negative(
2489                        &recipe,
2490                        &anchor,
2491                        &mut rng,
2492                        &mut seen,
2493                        &mut seen_negatives,
2494                    ) {
2495                        triplet = Some((recipe, sample));
2496                        recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2497                        break;
2498                    }
2499                }
2500                if let Some((recipe, triplet)) = triplet {
2501                    let SampleTriplet {
2502                        recipe: triplet_recipe_name,
2503                        anchor,
2504                        positive,
2505                        negative,
2506                        weight,
2507                        instruction,
2508                    } = triplet;
2509
2510                    // Look up the anchor record's label to override structural defaults.
2511                    // For SRD pair-mode data, the label is stored in DataRecord.
2512                    let record_label = self
2513                        .records
2514                        .get(&RecordKey::of_chunk(&anchor))
2515                        .and_then(|r| r.label.clone());
2516
2517                    if pairs.len() < self.config.batch_size {
2518                        pairs.push(SamplePair {
2519                            recipe: triplet_recipe_name.clone(),
2520                            anchor: anchor.clone(),
2521                            positive: positive.clone(),
2522                            weight,
2523                            instruction: instruction.clone(),
2524                            label: record_label.clone().unwrap_or(PairLabel::Positive),
2525                            reason: None,
2526                        });
2527                    }
2528                    if pairs.len() < self.config.batch_size {
2529                        pairs.push(SamplePair {
2530                            recipe: triplet_recipe_name,
2531                            anchor,
2532                            positive: negative,
2533                            weight,
2534                            instruction,
2535                            label: record_label.unwrap_or(PairLabel::Negative),
2536                            reason: Some(strategy_reason(&recipe.negative_strategy).to_string()),
2537                        });
2538                    }
2539                }
2540            }
2541            if recipe_steps > 0 {
2542                self.triplet_recipe_rr_idx =
2543                    self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2544            }
2545            self.rng = rng;
2546            pad_with_reuse(&mut pairs, self.config.batch_size);
2547            if pairs.len() == self.config.batch_size {
2548                return Ok(SampleBatch { pairs });
2549            }
2550            return Err(SamplerError::Exhausted(
2551                last_recipe_name
2552                    .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2553                    .to_string(),
2554            ));
2555        }
2556
2557        let mut pairs = Vec::new();
2558        let mut seen = HashSet::new();
2559        // Per-batch negative-text dedup (Option A, best-effort pre-pad).
2560        let mut seen_negatives: IntSet<u64> = IntSet::default();
2561        let mut source_steps = 0usize;
2562        let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2563        let mut source_idx = self.source_cycle_idx % sources.len();
2564        let mut cycle_sources = self.shuffled_source_cycle(cycle);
2565        // Explicit non-uniform weights engage deficit round-robin source
2566        // choice; uniform/empty maps keep legacy uniform cycling bit-identical.
2567        let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2568            Some(w) if Self::weights_are_selective(w) => Some(w),
2569            _ => None,
2570        };
2571        let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2572        let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2573        let mut recipe_steps = 0usize;
2574        let max_recipe_len = sources
2575            .iter()
2576            .map(|source| self.triplet_recipe_count_for_source(source))
2577            .max()
2578            .unwrap_or(1)
2579            .max(1);
2580        let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2581        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2582        for _ in 0..attempts {
2583            if pairs.len() >= self.config.batch_size {
2584                break;
2585            }
2586            let weighted_source;
2587            let source = if let Some(w) = weighted {
2588                match self.select_weighted_source(&sources, w) {
2589                    Some(s) => {
2590                        weighted_source = s;
2591                        weighted_source.as_str()
2592                    }
2593                    None => break,
2594                }
2595            } else {
2596                cycle_sources[source_idx].as_str()
2597            };
2598            let (triplet, attempts_used) = self.sample_source_triplet_candidate(
2599                source,
2600                target_split,
2601                &mut recipe_orders,
2602                &mut recipe_positions,
2603                &mut rng,
2604                &mut seen,
2605                &mut seen_negatives,
2606            );
2607            recipe_steps = recipe_steps.saturating_add(attempts_used);
2608            if let Some((recipe, triplet)) = triplet {
2609                let SampleTriplet {
2610                    recipe: triplet_recipe_name,
2611                    anchor,
2612                    positive,
2613                    negative,
2614                    weight,
2615                    instruction,
2616                } = triplet;
2617
2618                // Look up the anchor record's label to override structural defaults.
2619                let record_label = self
2620                    .records
2621                    .get(&RecordKey::of_chunk(&anchor))
2622                    .and_then(|r| r.label.clone());
2623
2624                if pairs.len() < self.config.batch_size {
2625                    pairs.push(SamplePair {
2626                        recipe: triplet_recipe_name.clone(),
2627                        anchor: anchor.clone(),
2628                        positive: positive.clone(),
2629                        weight,
2630                        instruction: instruction.clone(),
2631                        label: record_label.clone().unwrap_or(PairLabel::Positive),
2632                        reason: None,
2633                    });
2634                }
2635                if pairs.len() < self.config.batch_size {
2636                    pairs.push(SamplePair {
2637                        recipe: triplet_recipe_name,
2638                        anchor,
2639                        positive: negative,
2640                        weight,
2641                        instruction,
2642                        label: record_label.unwrap_or(PairLabel::Negative),
2643                        reason: Some(strategy_reason(&recipe.negative_strategy).to_string()),
2644                    });
2645                }
2646            }
2647            source_idx += 1;
2648            source_steps += 1;
2649            if source_idx >= cycle_sources.len() {
2650                source_idx = 0;
2651                cycle = cycle.saturating_add(1);
2652                cycle_sources = self.shuffled_source_cycle(cycle);
2653            }
2654        }
2655        if recipe_steps > 0 {
2656            self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2657        }
2658        self.rng = rng;
2659        pad_with_reuse(&mut pairs, self.config.batch_size);
2660        if pairs.len() == self.config.batch_size {
2661            self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2662            self.source_state_dirty = sources.len() > 1;
2663            return Ok(SampleBatch { pairs });
2664        }
2665        Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
2666    }
2667
2668    /// Check whether `sample` passes both per-batch and per-window text dedup.
2669    /// The per-batch set is never cleared mid-batch (survives epoch advances),
2670    /// while `emitted_text_hashes` is cleared on epoch advance to allow fresh
2671    /// permutations.  Both must pass for the sample to be emitted.
2672    fn try_emit_text_sample(
2673        &mut self,
2674        sample: &TextSample,
2675        batch_texts: &mut HashSet<String>,
2676    ) -> bool {
2677        batch_texts.insert(sample.chunk.text.clone())
2678            && self
2679                .emitted_text_hashes
2680                .entry(RecordKey::of_chunk(&sample.chunk))
2681                .or_default()
2682                .insert(stable_hash_str(0, &sample.chunk.text))
2683    }
2684
2685    fn next_text_batch_inner_with_weights(
2686        &mut self,
2687        target_split: SplitLabel,
2688        weights: Option<&HashMap<SourceId, f32>>,
2689    ) -> Result<TextBatch, SamplerError> {
2690        self.ingest_with_weights_fallback(target_split, weights)?;
2691        self.ensure_split_has_records(target_split)?;
2692        let sources = self.source_order.clone();
2693        if sources.is_empty() {
2694            if self.text_recipes.is_empty() {
2695                return Err(SamplerError::Configuration(
2696                    "no text recipes configured".into(),
2697                ));
2698            }
2699            let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2700            let recipe_weights: Vec<f32> = self.text_recipes.iter().map(|r| r.weight).collect();
2701            let recipe_order = self.text_recipe_order_weighted_cycled(
2702                &recipe_weights,
2703                self.text_recipe_rr_idx,
2704                &mut rng,
2705            );
2706            let mut samples = Vec::new();
2707            let mut batch_texts = HashSet::new();
2708            let mut last_recipe_name = None;
2709            let mut recipe_pos = 0usize;
2710            let mut recipe_steps = 0usize;
2711            let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2712            for _ in 0..attempts {
2713                if samples.len() >= self.config.batch_size {
2714                    break;
2715                }
2716                let recipe_idx = recipe_order[recipe_pos];
2717                recipe_pos = (recipe_pos + 1) % recipe_order.len();
2718                recipe_steps = recipe_steps.saturating_add(1);
2719                let recipe = self.text_recipes[recipe_idx].clone();
2720                last_recipe_name = Some(recipe.name.clone());
2721                if let Some(sample) =
2722                    self.make_text_sample_for_split(&recipe, None, target_split, &mut rng)
2723                    && self.try_emit_text_sample(&sample, &mut batch_texts)
2724                {
2725                    samples.push(sample);
2726                }
2727            }
2728            if recipe_steps > 0 {
2729                self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2730            }
2731            self.rng = rng;
2732            pad_with_reuse(&mut samples, self.config.batch_size);
2733            if samples.len() == self.config.batch_size {
2734                return Ok(TextBatch { samples });
2735            }
2736            return Err(SamplerError::Exhausted(
2737                last_recipe_name
2738                    .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TEXT))
2739                    .to_string(),
2740            ));
2741        }
2742
2743        let mut samples = Vec::new();
2744        let mut batch_texts = HashSet::new();
2745        let mut source_steps = 0usize;
2746        let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2747        let mut idx = self.source_cycle_idx % sources.len();
2748        let mut cycle_sources = self.shuffled_source_cycle(cycle);
2749        let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2750        let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2751        let mut recipe_steps = 0usize;
2752        let max_recipe_len = sources
2753            .iter()
2754            .map(|source| self.text_recipes_for_source(source).len())
2755            .max()
2756            .unwrap_or(1)
2757            .max(1);
2758        let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2759        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2760        // Explicit non-uniform weights engage deficit round-robin source
2761        // choice; uniform/empty maps keep legacy uniform cycling bit-identical.
2762        let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2763            Some(w) if Self::weights_are_selective(w) => Some(w),
2764            _ => None,
2765        };
2766        for _ in 0..attempts {
2767            if samples.len() >= self.config.batch_size {
2768                break;
2769            }
2770            let weighted_source;
2771            let source = if let Some(w) = weighted {
2772                match self.select_weighted_source(&sources, w) {
2773                    Some(s) => {
2774                        weighted_source = s;
2775                        weighted_source.as_str()
2776                    }
2777                    None => break,
2778                }
2779            } else {
2780                cycle_sources[idx].as_str()
2781            };
2782            let recipes = self.text_recipes_for_source(source).to_vec();
2783            if recipes.is_empty() {
2784                idx += 1;
2785                source_steps += 1;
2786                if idx >= cycle_sources.len() {
2787                    idx = 0;
2788                    cycle = cycle.saturating_add(1);
2789                    cycle_sources = self.shuffled_source_cycle(cycle);
2790                }
2791                continue;
2792            }
2793            if !recipe_orders.contains_key(source) {
2794                let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2795                let order = self.text_recipe_order_weighted_cycled(
2796                    &recipe_weights,
2797                    self.text_recipe_rr_idx,
2798                    &mut rng,
2799                );
2800                recipe_orders.insert(source.to_string(), order);
2801            }
2802            let order = recipe_orders
2803                .get(source)
2804                .expect("recipe order missing for source");
2805            let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2806            let mut sample: Option<(TextRecipe, TextSample)> = None;
2807            for offset in 0..order.len() {
2808                let recipe_idx = order[(*pos + offset) % order.len()];
2809                let recipe = recipes[recipe_idx].clone();
2810                if let Some(item) =
2811                    self.make_text_sample_for_split(&recipe, Some(source), target_split, &mut rng)
2812                    && self.try_emit_text_sample(&item, &mut batch_texts)
2813                {
2814                    recipe_steps = recipe_steps.saturating_add(offset + 1);
2815                    *pos = (*pos + offset + 1) % order.len();
2816                    sample = Some((recipe, item));
2817                    break;
2818                }
2819            }
2820            if sample.is_none() {
2821                recipe_steps = recipe_steps.saturating_add(order.len());
2822            }
2823            if let Some((_recipe, sample)) = sample {
2824                samples.push(sample);
2825            }
2826            idx += 1;
2827            source_steps += 1;
2828            if idx >= cycle_sources.len() {
2829                idx = 0;
2830                cycle = cycle.saturating_add(1);
2831                cycle_sources = self.shuffled_source_cycle(cycle);
2832            }
2833        }
2834        if samples.len() != self.config.batch_size {
2835            pad_with_reuse(&mut samples, self.config.batch_size);
2836        }
2837        if samples.len() != self.config.batch_size {
2838            self.rng = rng;
2839            return Err(SamplerError::Exhausted(RECIPE_LABEL_TEXT.into()));
2840        }
2841        self.rng = rng;
2842        self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2843        self.source_state_dirty = sources.len() > 1;
2844        if recipe_steps > 0 {
2845            self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2846        }
2847        Ok(TextBatch { samples })
2848    }
2849
2850    fn next_triplet_batch_inner_with_weights(
2851        &mut self,
2852        target_split: SplitLabel,
2853        weights: Option<&HashMap<SourceId, f32>>,
2854    ) -> Result<TripletBatch, SamplerError> {
2855        self.ingest_with_weights_fallback(target_split, weights)?;
2856        self.ensure_split_has_records(target_split)?;
2857        let sources = self.source_order.clone();
2858        if sources.is_empty() {
2859            if self.triplet_recipes.is_empty() {
2860                return Err(SamplerError::Configuration(
2861                    "no triplet recipes configured".into(),
2862                ));
2863            }
2864            let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2865            let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2866            let recipe_order = self.recipe_order_weighted_cycled(
2867                &recipe_weights,
2868                self.triplet_recipe_rr_idx,
2869                &mut rng,
2870            );
2871            let mut triplets = Vec::new();
2872            let mut seen = HashSet::new();
2873            // Per-batch negative-text dedup (Option A, best-effort pre-pad).
2874            let mut seen_negatives: IntSet<u64> = IntSet::default();
2875            let mut last_recipe_name = None;
2876            let mut recipe_pos = 0usize;
2877            let mut recipe_steps = 0usize;
2878            let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2879            for _ in 0..attempts {
2880                if triplets.len() >= self.config.batch_size {
2881                    break;
2882                }
2883                let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2884                    break;
2885                };
2886                let mut triplet = None;
2887                for offset in 0..recipe_order.len() {
2888                    let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2889                    recipe_steps = recipe_steps.saturating_add(1);
2890                    let recipe = self.triplet_recipes[idx].clone();
2891                    last_recipe_name = Some(recipe.name.clone());
2892                    if let Some(sample) = self.make_triplet_with_unique_negative(
2893                        &recipe,
2894                        &anchor,
2895                        &mut rng,
2896                        &mut seen,
2897                        &mut seen_negatives,
2898                    ) {
2899                        triplet = Some(sample);
2900                        recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2901                        break;
2902                    }
2903                }
2904                if let Some(triplet) = triplet {
2905                    triplets.push(triplet);
2906                }
2907            }
2908            if recipe_steps > 0 {
2909                self.triplet_recipe_rr_idx =
2910                    self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2911            }
2912            self.rng = rng;
2913            pad_with_reuse(&mut triplets, self.config.batch_size);
2914            if triplets.len() == self.config.batch_size {
2915                return Ok(TripletBatch { triplets });
2916            }
2917            return Err(SamplerError::Exhausted(
2918                last_recipe_name
2919                    .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2920                    .to_string(),
2921            ));
2922        }
2923
2924        // ── Declarations for multi-source parallel path ────────────────────────
2925        let mut triplets: Vec<SampleTriplet> = Vec::new();
2926        let mut seen: HashSet<(RecordKey, RecordKey, RecordKey)> = HashSet::new();
2927        // Per-batch negative-text dedup (Option A, best-effort pre-pad).
2928        let mut seen_negatives: IntSet<u64> = IntSet::default();
2929        let mut source_steps = 0usize;
2930        let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2931        let mut source_idx = self.source_cycle_idx % sources.len();
2932        let mut cycle_sources = self.shuffled_source_cycle(cycle);
2933        // Explicit non-uniform weights engage deficit round-robin source
2934        // choice; uniform/empty maps keep legacy uniform cycling bit-identical.
2935        let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2936            Some(w) if Self::weights_are_selective(w) => Some(w),
2937            _ => None,
2938        };
2939        let mut recipe_steps = 0usize;
2940        let max_recipe_len = sources
2941            .iter()
2942            .map(|source| self.triplet_recipe_count_for_source(source.as_str()))
2943            .max()
2944            .unwrap_or(1)
2945            .max(1);
2946        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2947
2948        // ── Phase 1: sequential anchor pre-assignment ──────────────────────────
2949        // Advance cursors and pick anchors deterministically. Capture a fork seed
2950        // for each slot so that Phase 2 can reconstruct independent rngs in
2951        // deterministic order.
2952        struct SlotPlan {
2953            anchor: Arc<DataRecord>,
2954            recipes: Vec<TripletRecipe>,
2955            fork_seed: u64,
2956        }
2957        let target_slots = self.config.batch_size * 4 * max_recipe_len;
2958        let mut slot_plans: Vec<SlotPlan> = Vec::with_capacity(target_slots);
2959
2960        for _ in 0..target_slots {
2961            if slot_plans.len() >= target_slots {
2962                break;
2963            }
2964            let weighted_source;
2965            let source = if let Some(w) = weighted {
2966                match self.select_weighted_source(&sources, w) {
2967                    Some(s) => {
2968                        weighted_source = s;
2969                        weighted_source.as_str()
2970                    }
2971                    None => break,
2972                }
2973            } else {
2974                cycle_sources[source_idx].as_str()
2975            };
2976            let (recipes, _) = self.resolve_source_triplet_plan(source);
2977            if !recipes.is_empty() {
2978                let fork_seed = rng.next_u64();
2979                if let Some(anchor) = self.choose_anchor_record(Some(source), target_split) {
2980                    slot_plans.push(SlotPlan {
2981                        anchor,
2982                        recipes,
2983                        fork_seed,
2984                    });
2985                    recipe_steps = recipe_steps.saturating_add(1);
2986                }
2987            }
2988            source_idx += 1;
2989            source_steps += 1;
2990            if source_idx >= cycle_sources.len() {
2991                source_idx = 0;
2992                cycle = cycle.saturating_add(1);
2993                cycle_sources = self.shuffled_source_cycle(cycle);
2994            }
2995        }
2996        self.rng = rng;
2997
2998        // ── Phase 2: parallel anchor+positive selection ────────────────────────
2999        // Only chunk selection runs in parallel; negative selection and the
3000        // coin-flip swap remain sequential in Phase 3 so that the BM25 cursor
3001        // rotation (shared RwLock state) is deterministic across runs.
3002        struct SlotCandidate {
3003            recipe: TripletRecipe,
3004            anchor: Arc<DataRecord>,
3005            anchor_chunk: RecordChunk,
3006            positive_chunk: RecordChunk,
3007            anchor_raw_text: String,
3008        }
3009        let mut raw_candidates: Vec<(usize, Option<SlotCandidate>)> = slot_plans
3010            .par_iter()
3011            .enumerate()
3012            .map(|(slot_idx, plan)| {
3013                let mut fork_rng = DeterministicRng::new(plan.fork_seed);
3014                // Use weighted shuffle so recipes with higher weight are tried more
3015                // often and zero-weight recipes are excluded entirely.
3016                let weights: Vec<f32> = plan.recipes.iter().map(|r| r.weight).collect();
3017                let order = weighted_recipe_order(&weights, &mut fork_rng);
3018                let mut candidate = None;
3019                for &idx in &order {
3020                    let recipe = &plan.recipes[idx];
3021                    if let Some((ac, pc, raw)) =
3022                        self.select_anchor_positive_for_recipe(recipe, &plan.anchor, &mut fork_rng)
3023                    {
3024                        candidate = Some(SlotCandidate {
3025                            recipe: recipe.clone(),
3026                            anchor: Arc::clone(&plan.anchor),
3027                            anchor_chunk: ac,
3028                            positive_chunk: pc,
3029                            anchor_raw_text: raw,
3030                        });
3031                        break;
3032                    }
3033                }
3034                (slot_idx, candidate)
3035            })
3036            .collect();
3037
3038        // ── Phase 3: sequential finalization (negative selection + coin-flip) ───
3039        // Sort by slot_idx to restore deterministic ordering, then finalize each
3040        // candidate sequentially so that BM25 cursor rotation advances in a fixed
3041        // order regardless of rayon's thread scheduling.
3042        raw_candidates.sort_unstable_by_key(|(i, _)| *i);
3043        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3044        for (_, candidate) in raw_candidates {
3045            if triplets.len() >= self.config.batch_size {
3046                break;
3047            }
3048            let Some(sc) = candidate else { continue };
3049            // Hoisted 50 % coin-flip swap (single RNG draw) + guard (A).
3050            // Unrecoverable a==p pairs skip before any negative work.
3051            let (final_anchor, final_positive) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0
3052            {
3053                (sc.positive_chunk, sc.anchor_chunk)
3054            } else {
3055                (sc.anchor_chunk, sc.positive_chunk)
3056            };
3057            if !sc.recipe.allow_same_anchor_positive && final_anchor.text == final_positive.text {
3058                continue;
3059            }
3060            // Negative-only retries: references only, no anchor/positive clones.
3061            // (B) checks inside loop are all negative-resolvable.
3062            let mut accepted: Option<(RecordChunk, f32, String)> = None;
3063            for _ in 0..NEG_DEDUP_RETRIES {
3064                let (negative_record, fallback_used) = match self.select_negative_record(
3065                    &sc.anchor,
3066                    &sc.recipe.negative_strategy,
3067                    Some(&sc.anchor_raw_text),
3068                    &mut rng,
3069                ) {
3070                    Some(pair) => pair,
3071                    None => break,
3072                };
3073                let mut negative_chunk = match self.select_chunk_parallel(
3074                    &negative_record,
3075                    &sc.recipe.negative_selector,
3076                    &mut rng,
3077                ) {
3078                    Some(c) => c,
3079                    None => continue,
3080                };
3081                self.decorate_chunk_parallel(&negative_record, &mut negative_chunk, &mut rng);
3082                if negative_chunk.text == final_positive.text
3083                    || negative_chunk.text == final_anchor.text
3084                {
3085                    continue;
3086                }
3087                let key = (
3088                    RecordKey::of_chunk(&final_anchor),
3089                    RecordKey::of_chunk(&final_positive),
3090                    RecordKey::of_chunk(&negative_chunk),
3091                );
3092                let neg_h = negative_text_hash(&negative_chunk);
3093                // Atomic check-then-commit: no partial `seen` mutation on reject.
3094                if seen.contains(&key) || seen_negatives.contains(&neg_h) {
3095                    continue;
3096                }
3097                let chunk_weight = self.triplet_chunk_weight(
3098                    &sc.recipe,
3099                    &final_anchor,
3100                    &final_positive,
3101                    &negative_chunk,
3102                );
3103                let weight = sc.recipe.weight * chunk_weight;
3104                let recipe_name = if fallback_used {
3105                    format!("{}_fallback_same_split", sc.recipe.name)
3106                } else {
3107                    sc.recipe.name.to_string()
3108                };
3109                seen.insert(key);
3110                seen_negatives.insert(neg_h);
3111                accepted = Some((negative_chunk, weight, recipe_name));
3112                break;
3113            }
3114            let Some((negative_chunk, weight, recipe_name)) = accepted else {
3115                continue;
3116            };
3117            let triplet = SampleTriplet {
3118                recipe: recipe_name,
3119                anchor: final_anchor,
3120                positive: final_positive,
3121                negative: negative_chunk,
3122                weight,
3123                instruction: sc.recipe.instruction.as_ref().map(|s| s.to_string()),
3124            };
3125            if triplets.len() < self.config.batch_size {
3126                triplets.push(triplet);
3127            }
3128        }
3129        self.rng = rng;
3130
3131        if recipe_steps > 0 {
3132            self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
3133        }
3134        pad_with_reuse(&mut triplets, self.config.batch_size);
3135        if triplets.len() == self.config.batch_size {
3136            self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
3137            self.source_state_dirty = sources.len() > 1;
3138            let batch = TripletBatch { triplets };
3139            return Ok(batch);
3140        }
3141        Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
3142    }
3143
3144    // ── test-only helpers ─────────────────────────────────────────────────────
3145
3146    /// Return a mutable reference to the concrete [`backends::Bm25Backend`], panicking
3147    /// when the feature is disabled (which won't happen — call sites are
3148    /// equally gated on `#[cfg(feature = "bm25-mining")]`).
3149    #[cfg(test)]
3150    #[cfg(feature = "bm25-mining")]
3151    fn bm25_backend_mut(&mut self) -> &mut backends::Bm25Backend {
3152        self.negative_backend
3153            .as_any_mut()
3154            .downcast_mut::<backends::Bm25Backend>()
3155            .expect("bm25_backend_mut: negative_backend is Bm25Backend when bm25-mining feature is active")
3156    }
3157
3158    /// Test shim: `recipe_order_weighted_shuffled` using the inner rng.
3159    #[cfg(test)]
3160    fn recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
3161        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3162        let result = self.recipe_order_weighted_shuffled(weights, &mut rng);
3163        self.rng = rng;
3164        result
3165    }
3166
3167    /// Test shim: `recipe_order_weighted_cycled` using the inner rng.
3168    #[cfg(test)]
3169    fn recipe_order_weighted_cycled_seeded(
3170        &mut self,
3171        weights: &[f32],
3172        rr_idx: usize,
3173    ) -> Vec<usize> {
3174        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3175        let result = self.recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
3176        self.rng = rng;
3177        result
3178    }
3179
3180    /// Test shim: `text_recipe_order_weighted_shuffled` using the inner rng.
3181    #[cfg(test)]
3182    fn text_recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
3183        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3184        let result = self.text_recipe_order_weighted_shuffled(weights, &mut rng);
3185        self.rng = rng;
3186        result
3187    }
3188
3189    /// Test shim: `text_recipe_order_weighted_cycled` using the inner rng.
3190    #[cfg(test)]
3191    fn text_recipe_order_weighted_cycled_seeded(
3192        &mut self,
3193        weights: &[f32],
3194        rr_idx: usize,
3195    ) -> Vec<usize> {
3196        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3197        let result = self.text_recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
3198        self.rng = rng;
3199        result
3200    }
3201
3202    /// Test shim: `select_negative_record` using the inner rng.
3203    #[cfg(test)]
3204    fn select_negative_record_seeded(
3205        &mut self,
3206        anchor_record: &DataRecord,
3207        strategy: &NegativeStrategy,
3208        anchor_query_text: Option<&str>,
3209    ) -> Option<(Arc<DataRecord>, bool)> {
3210        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3211        let result =
3212            self.select_negative_record(anchor_record, strategy, anchor_query_text, &mut rng);
3213        self.rng = rng;
3214        result
3215    }
3216
3217    /// Test shim: `make_triplet_with_unique_negative` using the inner rng and
3218    /// fresh per-batch dedup sets (single-candidate semantics).
3219    #[cfg(test)]
3220    fn make_triplet_with_unique_negative_seeded(
3221        &mut self,
3222        recipe: &TripletRecipe,
3223        anchor: &DataRecord,
3224    ) -> Option<SampleTriplet> {
3225        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3226        let mut seen = HashSet::new();
3227        let mut seen_negatives: IntSet<u64> = IntSet::default();
3228        let result = self.make_triplet_with_unique_negative(
3229            recipe,
3230            anchor,
3231            &mut rng,
3232            &mut seen,
3233            &mut seen_negatives,
3234        );
3235        self.rng = rng;
3236        result
3237    }
3238
3239    /// Test shim: `decorate_chunk` using the inner rng.
3240    #[cfg(test)]
3241    fn decorate_chunk_seeded(&mut self, record: &DataRecord, chunk: &mut RecordChunk) {
3242        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3243        self.decorate_chunk(record, chunk, &mut rng);
3244        self.rng = rng;
3245    }
3246
3247    // ── extended-metrics helpers ───────────────────────────────────────────────
3248
3249    /// Return cumulative `(fallback_count, selection_count)` from the BM25
3250    /// backend since it was created.  Only available when both `bm25-mining`
3251    /// and `extended-metrics` features are active.
3252    #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3253    fn bm25_fallback_stats(&self) -> (u64, u64) {
3254        self.negative_backend.bm25_fallback_stats()
3255    }
3256
3257    /// Look up `anchor`'s split label and return BM25-ranked candidate IDs.
3258    ///
3259    /// Mirrors the old `bm25_ranked_candidates` method that tests called on
3260    /// `TripletSamplerInner` before BM25 state moved into `Bm25Backend`.
3261    #[cfg(test)]
3262    #[cfg(feature = "bm25-mining")]
3263    fn bm25_ranked_candidates(&mut self, anchor: &crate::data::DataRecord) -> Vec<RecordKey> {
3264        let split = self
3265            .split_store
3266            .label_for(&RecordKey::from(anchor))
3267            .unwrap_or(SplitLabel::Train);
3268        self.negative_backend
3269            .as_any_mut()
3270            .downcast_mut::<backends::Bm25Backend>()
3271            .expect("bm25_ranked_candidates: Bm25Backend")
3272            .ranked_candidates_pub(anchor, split)
3273    }
3274}
3275
3276/// Build a shuffled selection order from a slice of recipe weights.
3277///
3278/// Each recipe with `weight > 0.0` receives a number of slots proportional to its weight
3279/// relative to the smallest positive weight, capped at
3280/// [`RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER`].  The resulting expanded slot list is then
3281/// shuffled, giving each recipe approximately the requested frequency in round-robin
3282/// cycling.
3283///
3284/// This is a free function (no `self`) so it can be called from rayon parallel closures as
3285/// well as from within the `TripletSamplerInner` instance methods.
3286fn weighted_recipe_order(weights: &[f32], rng: &mut DeterministicRng) -> Vec<usize> {
3287    let nonzero: Vec<(usize, f32)> = weights
3288        .iter()
3289        .enumerate()
3290        .filter(|(_, w)| **w > 0.0)
3291        .map(|(i, &w)| (i, w))
3292        .collect();
3293    if nonzero.is_empty() {
3294        return Vec::new();
3295    }
3296    let w_min = nonzero
3297        .iter()
3298        .map(|(_, w)| *w)
3299        .fold(f32::INFINITY, f32::min);
3300    let mut order: Vec<usize> = Vec::new();
3301    for (recipe_idx, w) in &nonzero {
3302        let tickets = ((w / w_min).round() as usize).clamp(1, RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER);
3303        for _ in 0..tickets {
3304            order.push(*recipe_idx);
3305        }
3306    }
3307    order.shuffle(rng);
3308    order
3309}
3310
3311fn same_selector_pair_is_valid(
3312    anchor_chunk: &RecordChunk,
3313    positive_chunk: &RecordChunk,
3314    enforce_window_pair: bool,
3315) -> bool {
3316    if triplet_chunk_key(anchor_chunk) == triplet_chunk_key(positive_chunk) {
3317        return false;
3318    }
3319    if !enforce_window_pair {
3320        return true;
3321    }
3322    matches!(
3323        (&anchor_chunk.view, &positive_chunk.view),
3324        (ChunkView::Window { .. }, ChunkView::Window { .. })
3325    )
3326}
3327
3328impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSampler<S> {
3329    /// Create a sampler from config and a split-state backend.
3330    pub fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
3331        let inner = TripletSamplerInner::new(config, split_store);
3332        Self {
3333            inner: Mutex::new(inner),
3334        }
3335    }
3336
3337    /// Create a sampler from config with a custom chunking implementation.
3338    pub fn new_with_chunker(
3339        config: SamplerConfig,
3340        split_store: Arc<S>,
3341        chunker: Arc<dyn ChunkingAlgorithm>,
3342    ) -> Self {
3343        let inner = TripletSamplerInner::new_with_chunker(config, split_store, chunker);
3344        Self {
3345            inner: Mutex::new(inner),
3346        }
3347    }
3348
3349    /// Return an unweighted pair batch for `split`.
3350    ///
3351    /// Samples all sources uniformly. Prefer
3352    /// [`Self::next_pair_batch_with_weights_for_split`] with an explicit
3353    /// per-source weight map to honor a data mixture.
3354    #[deprecated(
3355        note = "samples all sources uniformly; call next_pair_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3356    )]
3357    pub fn next_pair_batch_for_split(
3358        &self,
3359        split: SplitLabel,
3360    ) -> Result<SampleBatch, SamplerError> {
3361        self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3362    }
3363
3364    /// Return an unweighted text batch for `split`.
3365    ///
3366    /// Samples all sources uniformly. Prefer
3367    /// [`Self::next_text_batch_with_weights_for_split`] with an explicit
3368    /// per-source weight map to honor a data mixture.
3369    #[deprecated(
3370        note = "samples all sources uniformly; call next_text_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3371    )]
3372    pub fn next_text_batch_for_split(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3373        self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3374    }
3375
3376    /// Return an unweighted triplet batch for `split`.
3377    ///
3378    /// Samples all sources uniformly. Prefer
3379    /// [`Self::next_triplet_batch_with_weights_for_split`] with an explicit
3380    /// per-source weight map to honor a data mixture.
3381    #[deprecated(
3382        note = "samples all sources uniformly; call next_triplet_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3383    )]
3384    pub fn next_triplet_batch_for_split(
3385        &self,
3386        split: SplitLabel,
3387    ) -> Result<TripletBatch, SamplerError> {
3388        self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3389    }
3390
3391    /// Shared retry loop for all three batch types: load state, bump step,
3392    /// retry on exhaustion up to EXHAUSTION_RETRY_LIMIT times, and return
3393    /// the batch or propagate the final error.
3394    fn next_batch_with_retry<B>(
3395        &self,
3396        split: SplitLabel,
3397        weights: &HashMap<SourceId, f32>,
3398        inner_fn: impl Fn(
3399            &mut TripletSamplerInner<S>,
3400            SplitLabel,
3401            Option<&HashMap<SourceId, f32>>,
3402        ) -> Result<B, SamplerError>,
3403    ) -> Result<B, SamplerError> {
3404        let mut inner = self.inner.lock().unwrap();
3405        inner.ensure_split_allowed(split)?;
3406        // Validate weight keys up front, on every path: a typo'd source id
3407        // must fail loudly, never silently sample as if unweighted.
3408        if !weights.is_empty() {
3409            inner.ingestion.validate_weights(weights)?;
3410        }
3411        // Load persisted state BEFORE bumping step so the restored step
3412        // counter is reflected before we increment it for this call.
3413        inner.ensure_ingestion_cursors_loaded()?;
3414        // Bump epoch step once per public call, regardless of success/exhaustion.
3415        inner.ingestion.increment_epoch_step();
3416        for attempt in 0..=EXHAUSTION_RETRY_LIMIT {
3417            match inner_fn(&mut inner, split, Some(weights)) {
3418                Ok(batch) => return Ok(batch),
3419                Err(SamplerError::Exhausted(_)) if attempt < EXHAUSTION_RETRY_LIMIT => {
3420                    inner.force_ingest_refresh_with_weights_for_split(split, weights)?;
3421                }
3422                Err(err) => return Err(err),
3423            }
3424        }
3425        // The for-loop always returns via the match arms above.
3426        unreachable!()
3427    }
3428
3429    /// Return a weighted pair batch for `split` using per-source weights.
3430    ///
3431    /// Source choice is deficit round-robin over the map: deterministic and
3432    /// exact over each weight cycle (0.75/0.25 yields exactly 3:1 anchors).
3433    /// Unlisted sources get quantum 1.0; non-finite/negative quanta count as
3434    /// 0. Empty or all-equal maps keep legacy uniform cycling bit-identically.
3435    /// Weights also steer ingestion; see `ingest_with_weights_fallback`.
3436    pub fn next_pair_batch_with_weights_for_split(
3437        &self,
3438        split: SplitLabel,
3439        weights: &HashMap<SourceId, f32>,
3440    ) -> Result<SampleBatch, SamplerError> {
3441        self.next_batch_with_retry(split, weights, |inner, split, weights| {
3442            inner.next_pair_batch_inner_with_weights(split, weights)
3443        })
3444    }
3445
3446    /// Return a weighted text batch for `split` using per-source weights.
3447    ///
3448    /// Source choice is deficit round-robin over the map: deterministic and
3449    /// exact over each weight cycle. Unlisted sources get quantum 1.0;
3450    /// non-finite/negative quanta count as 0. Empty or all-equal maps keep
3451    /// legacy uniform cycling bit-identically.
3452    pub fn next_text_batch_with_weights_for_split(
3453        &self,
3454        split: SplitLabel,
3455        weights: &HashMap<SourceId, f32>,
3456    ) -> Result<TextBatch, SamplerError> {
3457        self.next_batch_with_retry(split, weights, |inner, split, weights| {
3458            inner.next_text_batch_inner_with_weights(split, weights)
3459        })
3460    }
3461
3462    /// Return a weighted triplet batch for `split` using per-source weights.
3463    ///
3464    /// Source choice is deficit round-robin over the map: deterministic and
3465    /// exact over each weight cycle. Unlisted sources get quantum 1.0;
3466    /// non-finite/negative quanta count as 0. Empty or all-equal maps keep
3467    /// legacy uniform cycling bit-identically.
3468    pub fn next_triplet_batch_with_weights_for_split(
3469        &self,
3470        split: SplitLabel,
3471        weights: &HashMap<SourceId, f32>,
3472    ) -> Result<TripletBatch, SamplerError> {
3473        self.next_batch_with_retry(split, weights, |inner, split, weights| {
3474            inner.next_triplet_batch_inner_with_weights(split, weights)
3475        })
3476    }
3477
3478    /// Spawn a background prefetcher for triplet batches.
3479    ///
3480    /// Samples all sources uniformly. Prefer
3481    /// [`Self::prefetch_triplet_batches_with_weights`] with an explicit
3482    /// per-source weight map to honor a data mixture.
3483    #[deprecated(
3484        note = "samples all sources uniformly; call prefetch_triplet_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3485    )]
3486    pub fn prefetch_triplet_batches(
3487        self: Arc<Self>,
3488        split: SplitLabel,
3489        capacity: usize,
3490    ) -> BatchPrefetcher<TripletBatch> {
3491        BatchPrefetcher::new(capacity, move || {
3492            self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3493        })
3494    }
3495
3496    /// Spawn a background prefetcher for weighted triplet batches.
3497    pub fn prefetch_triplet_batches_with_weights(
3498        self: Arc<Self>,
3499        split: SplitLabel,
3500        capacity: usize,
3501        weights: HashMap<SourceId, f32>,
3502    ) -> BatchPrefetcher<TripletBatch> {
3503        BatchPrefetcher::new(capacity, move || {
3504            self.next_triplet_batch_with_weights_for_split(split, &weights)
3505        })
3506    }
3507
3508    /// Spawn a background prefetcher for pair batches.
3509    ///
3510    /// Samples all sources uniformly. Prefer
3511    /// [`Self::prefetch_pair_batches_with_weights`] with an explicit
3512    /// per-source weight map to honor a data mixture.
3513    #[deprecated(
3514        note = "samples all sources uniformly; call prefetch_pair_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3515    )]
3516    pub fn prefetch_pair_batches(
3517        self: Arc<Self>,
3518        split: SplitLabel,
3519        capacity: usize,
3520    ) -> BatchPrefetcher<SampleBatch> {
3521        BatchPrefetcher::new(capacity, move || {
3522            self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3523        })
3524    }
3525
3526    /// Spawn a background prefetcher for weighted pair batches.
3527    pub fn prefetch_pair_batches_with_weights(
3528        self: Arc<Self>,
3529        split: SplitLabel,
3530        capacity: usize,
3531        weights: HashMap<SourceId, f32>,
3532    ) -> BatchPrefetcher<SampleBatch> {
3533        BatchPrefetcher::new(capacity, move || {
3534            self.next_pair_batch_with_weights_for_split(split, &weights)
3535        })
3536    }
3537
3538    /// Spawn a background prefetcher for text batches.
3539    ///
3540    /// Samples all sources uniformly. Prefer
3541    /// [`Self::prefetch_text_batches_with_weights`] with an explicit
3542    /// per-source weight map to honor a data mixture.
3543    #[deprecated(
3544        note = "samples all sources uniformly; call prefetch_text_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3545    )]
3546    pub fn prefetch_text_batches(
3547        self: Arc<Self>,
3548        split: SplitLabel,
3549        capacity: usize,
3550    ) -> BatchPrefetcher<TextBatch> {
3551        BatchPrefetcher::new(capacity, move || {
3552            self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3553        })
3554    }
3555
3556    /// Spawn a background prefetcher for weighted text batches.
3557    pub fn prefetch_text_batches_with_weights(
3558        self: Arc<Self>,
3559        split: SplitLabel,
3560        capacity: usize,
3561        weights: HashMap<SourceId, f32>,
3562    ) -> BatchPrefetcher<TextBatch> {
3563        BatchPrefetcher::new(capacity, move || {
3564            self.next_text_batch_with_weights_for_split(split, &weights)
3565        })
3566    }
3567
3568    /// Return the currently active text recipes.
3569    pub fn text_recipes(&self) -> Vec<TextRecipe> {
3570        let inner = self.inner.lock().unwrap();
3571        inner.text_recipes().to_vec()
3572    }
3573
3574    /// Register a data source for ingestion and sampling.
3575    ///
3576    /// Returns an error if the source's `id()` matches the reserved `__*__`
3577    /// pattern used for internal synthetic/metadata source identifiers.
3578    pub fn register_source(
3579        &self,
3580        source: Box<dyn DataSource + 'static>,
3581    ) -> Result<(), SamplerError> {
3582        let mut inner = self.inner.lock().unwrap();
3583        inner.register_source(source)
3584    }
3585
3586    /// Force sampler epoch to `epoch` (advanced deterministic replay control).
3587    pub fn set_epoch(&self, epoch: u64) -> Result<(), SamplerError> {
3588        let mut inner = self.inner.lock().unwrap();
3589        inner.set_epoch(epoch)
3590    }
3591
3592    /// Persist sampler and split runtime state for restart-resume.
3593    ///
3594    /// When `save_to` is `Some(path)`, current persisted runtime state is also
3595    /// mirrored to `path` when supported by the split-store backend.
3596    pub fn save_sampler_state(&self, save_to: Option<&Path>) -> Result<(), SamplerError> {
3597        let mut inner = self.inner.lock().unwrap();
3598        inner.save_sampler_state(save_to)
3599    }
3600
3601    /// Return cumulative `(fallback_count, selection_count)` from the BM25
3602    /// backend since it was created.  Only available when both `bm25-mining`
3603    /// and `extended-metrics` features are active.
3604    ///
3605    /// `selection_count` counts calls where the pool was non-empty.
3606    /// `fallback_count` counts the subset where BM25 produced no candidates
3607    /// intersecting the pool and random selection was used instead.
3608    #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3609    pub fn bm25_fallback_stats(&self) -> (u64, u64) {
3610        let inner = self.inner.lock().unwrap();
3611        inner.bm25_fallback_stats()
3612    }
3613}
3614
3615impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> Sampler for TripletSampler<S> {
3616    fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
3617        self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3618    }
3619
3620    fn next_pair_batch_with_weights(
3621        &self,
3622        split: SplitLabel,
3623        weights: &HashMap<SourceId, f32>,
3624    ) -> Result<SampleBatch, SamplerError> {
3625        self.next_pair_batch_with_weights_for_split(split, weights)
3626    }
3627
3628    fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3629        self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3630    }
3631
3632    fn next_text_batch_with_weights(
3633        &self,
3634        split: SplitLabel,
3635        weights: &HashMap<SourceId, f32>,
3636    ) -> Result<TextBatch, SamplerError> {
3637        self.next_text_batch_with_weights_for_split(split, weights)
3638    }
3639
3640    fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
3641        self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3642    }
3643
3644    fn next_triplet_batch_with_weights(
3645        &self,
3646        split: SplitLabel,
3647        weights: &HashMap<SourceId, f32>,
3648    ) -> Result<TripletBatch, SamplerError> {
3649        self.next_triplet_batch_with_weights_for_split(split, weights)
3650    }
3651
3652    fn save_sampler_state(&self, save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
3653        TripletSampler::save_sampler_state(self, save_to)
3654    }
3655}
3656
3657fn roles_match(target: &SectionRole, candidate: &SectionRole) -> bool {
3658    target == candidate
3659}
3660
3661fn role_cursor_key(record_key: &RecordKey, role: &SectionRole) -> (RecordKey, String) {
3662    (record_key.clone(), role_label(role))
3663}
3664
3665fn role_label(role: &SectionRole) -> String {
3666    match role {
3667        SectionRole::Anchor => ROLE_LABEL_ANCHOR.into(),
3668        SectionRole::Context => ROLE_LABEL_CONTEXT.into(),
3669    }
3670}
3671
3672fn taxonomy_value(record: &DataRecord, field: MetadataKey) -> Option<&str> {
3673    record.taxonomy.iter().find_map(|entry| field.strip(entry))
3674}
3675
3676fn strategy_reason(strategy: &NegativeStrategy) -> &'static str {
3677    match strategy {
3678        NegativeStrategy::WrongPublicationDate => NEG_REASON_WRONG_DATE,
3679        NegativeStrategy::WrongArticle => NEG_REASON_WRONG_ARTICLE,
3680        NegativeStrategy::QuestionAnswerMismatch => NEG_REASON_WRONG_QA,
3681        NegativeStrategy::SameRecord => NEG_REASON_SAME_RECORD,
3682    }
3683}
3684
3685/// Dedup key for text batch generation: (record_id, text).
3686///
3687/// Unlike `triplet_chunk_key` which compares structural identity
3688/// (record_id, section_idx, window_index), this key compares actual
3689/// text content. This prevents text-columns mode sources (where Anchor
3690/// and Context sections carry identical text) from producing duplicate
3691/// samples from the same record.
3692#[cfg(test)]
3693fn text_dedup_key(chunk: &RecordChunk) -> (String, String) {
3694    (chunk.record_id.clone(), chunk.text.clone())
3695}
3696
3697fn triplet_chunk_key(chunk: &RecordChunk) -> String {
3698    match &chunk.view {
3699        ChunkView::Window { index, .. } => {
3700            format!("{}|{}|w|{}", chunk.record_id, chunk.section_idx, index)
3701        }
3702        ChunkView::SummaryFallback { strategy, .. } => {
3703            format!("{}|{}|s|{}", chunk.record_id, chunk.section_idx, strategy)
3704        }
3705    }
3706}
3707
3708/// Bounded negative-redraw attempts per anchor-positive pair.
3709///
3710/// When a candidate negative collides with an already-accepted negative in
3711/// the current batch (or fails the intra-triplet text guard), only the
3712/// negative is redrawn — the anchor/positive pair is kept. After this many
3713/// failures the slot is skipped; the outer `attempts` budget then moves on.
3714const NEG_DEDUP_RETRIES: usize = 4;
3715
3716/// Pure `u64` text hash for per-batch negative dedup (Option A: negative-vs-negative).
3717///
3718/// InfoNCE-style objectives compare text content, so the key must be text-only:
3719/// no `RecordKey` coupling (same text from different records/windows must still
3720/// collide) and no `String` allocation in the hot loop.
3721#[inline]
3722fn negative_text_hash(chunk: &RecordChunk) -> u64 {
3723    stable_hash_str(0, &chunk.text)
3724}
3725
3726fn pad_with_reuse<T: Clone>(items: &mut Vec<T>, target: usize) {
3727    if items.is_empty() || items.len() >= target {
3728        return;
3729    }
3730    let seed = items.clone();
3731    let base_len = seed.len();
3732    for idx in 0..(target - items.len()) {
3733        items.push(seed[idx % base_len].clone());
3734    }
3735}
3736
3737#[cfg(test)]
3738mod tests;