positional.rs 25.5 KB
Newer Older
Yan Ru Pei's avatar
Yan Ru Pei committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Positional HashMap-based KV cache index with nested structure.
//!
//! This module provides a `PositionalIndexer` that uses nested HashMaps
//! keyed by position for better cache locality and enables jump/binary-search
//! optimizations in find_matches.
//!
//! # Structure
//!
//! - `index`: position -> local_hash -> seq_hash -> workers
//!   The main lookup structure. Position-first nesting enables O(1) position access.
//! - `worker_blocks`: worker -> seq_hash -> (position, local_hash)
//!   Per-worker reverse lookup for efficient remove operations.
//!
//! # Threading
//!
//! `PositionalIndexer` implements `SyncIndexer`, meaning all its methods are
//! synchronous and thread-safe (via `DashMap` and `RwLock`). To get the full
//! `KvIndexerInterface` with sticky event routing and worker threads, wrap it
//! in a `ThreadPoolIndexer`.
use dashmap::DashMap;
24
use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
25
use std::sync::Arc;
26
use std::sync::atomic::{AtomicUsize, Ordering};
Yan Ru Pei's avatar
Yan Ru Pei committed
27

28
use super::{EventKind, KvIndexerMetrics, SyncIndexer, WorkerTask};
29
use crate::active_set::reconcile_active_workers;
Yan Ru Pei's avatar
Yan Ru Pei committed
30
use crate::protocols::{
31
32
33
    DpRank, ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheEventError,
    KvCacheStoreData, KvCacheStoredBlockData, LocalBlockHash, OverlapScores, RouterEvent, WorkerId,
    WorkerWithDpRank,
Yan Ru Pei's avatar
Yan Ru Pei committed
34
35
36
37
38
39
40
41
42
};

/// Entry for the innermost level of the index.
///
/// Optimizes for the common case where there's only one sequence hash
/// at a given (position, local_hash) pair, avoiding HashMap allocation.
#[derive(Debug, Clone)]
enum SeqEntry {
    /// Single seq_hash -> workers mapping (common case, no HashMap allocation)
43
    Single(ExternalSequenceBlockHash, FxHashSet<WorkerWithDpRank>),
Yan Ru Pei's avatar
Yan Ru Pei committed
44
    /// Multiple seq_hash -> workers mappings (rare case, different prefixes)
45
    Multi(FxHashMap<ExternalSequenceBlockHash, FxHashSet<WorkerWithDpRank>>),
Yan Ru Pei's avatar
Yan Ru Pei committed
46
47
48
49
50
}

impl SeqEntry {
    /// Create a new entry with a single worker.
    fn new(seq_hash: ExternalSequenceBlockHash, worker: WorkerWithDpRank) -> Self {
51
        let mut workers = FxHashSet::default();
Yan Ru Pei's avatar
Yan Ru Pei committed
52
53
54
55
56
57
58
59
60
61
62
63
        workers.insert(worker);
        Self::Single(seq_hash, workers)
    }

    /// Insert a worker for a given seq_hash, upgrading to Multi if needed.
    fn insert(&mut self, seq_hash: ExternalSequenceBlockHash, worker: WorkerWithDpRank) {
        match self {
            Self::Single(existing_hash, workers) if *existing_hash == seq_hash => {
                workers.insert(worker);
            }
            Self::Single(existing_hash, existing_workers) => {
                // Upgrade to Multi
64
                let mut map = FxHashMap::with_capacity_and_hasher(2, FxBuildHasher);
Yan Ru Pei's avatar
Yan Ru Pei committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
                map.insert(*existing_hash, std::mem::take(existing_workers));
                map.entry(seq_hash).or_default().insert(worker);
                *self = Self::Multi(map);
            }
            Self::Multi(map) => {
                map.entry(seq_hash).or_default().insert(worker);
            }
        }
    }

    /// Remove a worker from a given seq_hash.
    /// Returns true if the entry is now completely empty and should be removed.
    fn remove(&mut self, seq_hash: ExternalSequenceBlockHash, worker: WorkerWithDpRank) -> bool {
        match self {
            Self::Single(existing_hash, workers) if *existing_hash == seq_hash => {
                workers.remove(&worker);
                workers.is_empty()
            }
            Self::Single(_, _) => false, // Different hash, nothing to remove
            Self::Multi(map) => {
                if let Some(workers) = map.get_mut(&seq_hash) {
                    workers.remove(&worker);
                    if workers.is_empty() {
                        map.remove(&seq_hash);
                    }
                }
                map.is_empty()
            }
        }
    }

    /// Get workers for a specific seq_hash.
97
    fn get(&self, seq_hash: ExternalSequenceBlockHash) -> Option<&FxHashSet<WorkerWithDpRank>> {
Yan Ru Pei's avatar
Yan Ru Pei committed
98
99
100
101
102
103
104
105
        match self {
            Self::Single(existing_hash, workers) if *existing_hash == seq_hash => Some(workers),
            Self::Single(_, _) => None,
            Self::Multi(map) => map.get(&seq_hash),
        }
    }
}

106
pub type LevelIndex = FxHashMap<ExternalSequenceBlockHash, (usize, LocalBlockHash)>;
Yan Ru Pei's avatar
Yan Ru Pei committed
107
108
109
110
111
112

/// Positional HashMap-based KV cache index.
///
/// Implements [`SyncIndexer`] for use with [`ThreadPoolIndexer`](crate::indexer::ThreadPoolIndexer).
/// All methods are synchronous and thread-safe.
pub struct PositionalIndexer {
113
    index: DashMap<(usize, LocalBlockHash), SeqEntry, FxBuildHasher>,
114
115

    tree_sizes: DashMap<WorkerWithDpRank, AtomicUsize, FxBuildHasher>,
Yan Ru Pei's avatar
Yan Ru Pei committed
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

    jump_size: usize,
}

impl PositionalIndexer {
    /// Create a new PositionalIndexer.
    ///
    /// # Arguments
    /// * `jump_size` - Jump size for find_matches optimization (e.g., 32).
    ///   The algorithm jumps by this many positions at a time, only scanning
    ///   intermediate positions when workers drain (stop matching).
    pub fn new(jump_size: usize) -> Self {
        assert!(jump_size > 0, "jump_size must be greater than 0");

        Self {
131
            index: DashMap::with_hasher(FxBuildHasher),
132
            tree_sizes: DashMap::with_hasher(FxBuildHasher),
Yan Ru Pei's avatar
Yan Ru Pei committed
133
134
135
136
137
138
139
140
141
142
            jump_size,
        }
    }
}

// ============================================================================
// SyncIndexer implementation
// ============================================================================

impl SyncIndexer for PositionalIndexer {
143
144
145
146
147
    fn worker(
        &self,
        event_receiver: flume::Receiver<WorkerTask>,
        metrics: Option<Arc<KvIndexerMetrics>>,
    ) -> anyhow::Result<()> {
148
        let mut worker_blocks = FxHashMap::default();
149
        let counters = metrics.as_ref().map(|m| m.prebind());
150
151
152
153

        while let Ok(task) = event_receiver.recv() {
            match task {
                WorkerTask::Event(event) => {
154
155
156
157
158
159
160
                    let kind = EventKind::of(&event.event.data);
                    let result = self.apply_event(&mut worker_blocks, event);
                    if result.is_err() {
                        tracing::warn!("Failed to apply event: {:?}", result.as_ref().err());
                    }
                    if let Some(ref c) = counters {
                        c.inc(kind, result);
Yan Ru Pei's avatar
Yan Ru Pei committed
161
                    }
162
163
164
165
                }
                WorkerTask::RemoveWorker(worker_id) => {
                    self.remove_or_clear_worker_blocks_impl(&mut worker_blocks, worker_id, false);
                }
166
167
168
                WorkerTask::RemoveWorkerDpRank(worker_id, dp_rank) => {
                    self.remove_worker_dp_rank_impl(&mut worker_blocks, worker_id, dp_rank);
                }
169
170
171
                WorkerTask::CleanupStaleChildren => {
                    self.run_cleanup_task();
                }
172
173
174
175
176
177
178
179
180
                WorkerTask::DumpEvents(sender) => {
                    let events = self.dump_events(&worker_blocks);
                    if let Err(e) = sender.send(Ok(events)) {
                        tracing::warn!("Failed to send events: {:?}", e);
                    }
                }
                WorkerTask::Terminate => {
                    break;
                }
Yan Ru Pei's avatar
Yan Ru Pei committed
181
182
183
            }
        }

184
185
186
187
188
189
        tracing::debug!("PositionalIndexer worker thread shutting down");
        Ok(())
    }

    fn find_matches(&self, sequence: &[LocalBlockHash], early_exit: bool) -> OverlapScores {
        self.jump_search_matches(sequence, early_exit)
Yan Ru Pei's avatar
Yan Ru Pei committed
190
191
192
193
194
195
196
197
198
199
    }
}

// ============================================================================
// Event processing (write operations)
// ============================================================================

impl PositionalIndexer {
    /// Process an event using the provided index and worker_blocks.
    /// This is called from worker threads.
200
201
202
    pub fn apply_event(
        &self,
        worker_blocks: &mut FxHashMap<WorkerWithDpRank, LevelIndex>,
Yan Ru Pei's avatar
Yan Ru Pei committed
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
        event: RouterEvent,
    ) -> Result<(), KvCacheEventError> {
        let (worker_id, kv_event) = (event.worker_id, event.event);
        let (id, op) = (kv_event.event_id, kv_event.data);

        let worker = WorkerWithDpRank::new(worker_id, kv_event.dp_rank);

        tracing::trace!(
            id,
            "PositionalIndexer::apply_event_impl: operation: {:?}",
            op
        );

        match op {
            KvCacheEventData::Stored(store_data) => {
218
                self.store_blocks_impl(worker_blocks, worker, store_data, id)?;
Yan Ru Pei's avatar
Yan Ru Pei committed
219
220
221
222

                Ok(())
            }
            KvCacheEventData::Removed(remove_data) => {
223
                self.remove_blocks_impl(worker_blocks, worker, &remove_data.block_hashes, id)?;
Yan Ru Pei's avatar
Yan Ru Pei committed
224
225
226
                Ok(())
            }
            KvCacheEventData::Cleared => {
227
                self.clear_worker_blocks_impl(worker_blocks, worker_id);
Yan Ru Pei's avatar
Yan Ru Pei committed
228
229
230
231
232
233
                Ok(())
            }
        }
    }

    fn store_blocks_impl(
234
235
        &self,
        worker_blocks: &mut FxHashMap<WorkerWithDpRank, LevelIndex>,
Yan Ru Pei's avatar
Yan Ru Pei committed
236
237
238
239
        worker: WorkerWithDpRank,
        store_data: KvCacheStoreData,
        event_id: u64,
    ) -> Result<(), KvCacheEventError> {
240
241
242
243
244
        let KvCacheStoreData {
            parent_hash,
            start_position,
            blocks,
        } = store_data;
245
        let worker_map = worker_blocks.entry(worker).or_default();
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        let start_pos = match start_position {
            Some(start_position) => start_position as usize,
            None => match parent_hash {
                Some(parent_hash) => {
                    let Some(entry) = worker_map.get(&parent_hash) else {
                        tracing::warn!(
                            worker_id = worker.worker_id.to_string(),
                            dp_rank = worker.dp_rank,
                            event_id,
                            parent_hash = ?parent_hash,
                        );
                        return Err(KvCacheEventError::ParentBlockNotFound);
                    };

                    entry.0 + 1 // parent position + 1
                }
                None => 0, // Start from position 0
            },
Yan Ru Pei's avatar
Yan Ru Pei committed
264
265
        };

266
        let worker_blocks_entry = worker_blocks.entry(worker).or_default();
Yan Ru Pei's avatar
Yan Ru Pei committed
267

268
        let num_stored_blocks = blocks.len();
Yan Ru Pei's avatar
Yan Ru Pei committed
269

270
        for (i, block_data) in blocks.into_iter().enumerate() {
Yan Ru Pei's avatar
Yan Ru Pei committed
271
272
273
274
            let position = start_pos + i;
            let local_hash = block_data.tokens_hash;
            let seq_hash = block_data.block_hash;

275
            self.index
Yan Ru Pei's avatar
Yan Ru Pei committed
276
277
278
279
280
                .entry((position, local_hash))
                .and_modify(|entry| entry.insert(seq_hash, worker))
                .or_insert_with(|| SeqEntry::new(seq_hash, worker));

            // Insert into worker_blocks: worker -> seq_hash -> (position, local_hash)
281
282
283
284
285
286
287
288
289
290
291
            worker_blocks_entry.insert(seq_hash, (position, local_hash));
        }

        match self.tree_sizes.get(&worker) {
            Some(size) => {
                size.fetch_add(num_stored_blocks, Ordering::Relaxed);
            }
            None => {
                self.tree_sizes
                    .insert(worker, AtomicUsize::new(num_stored_blocks));
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
292
293
294
295
296
297
        }

        Ok(())
    }

    fn remove_blocks_impl(
298
299
        &self,
        worker_blocks: &mut FxHashMap<WorkerWithDpRank, LevelIndex>,
Yan Ru Pei's avatar
Yan Ru Pei committed
300
301
302
303
        worker: WorkerWithDpRank,
        seq_hashes: &Vec<ExternalSequenceBlockHash>,
        event_id: u64,
    ) -> Result<(), KvCacheEventError> {
304
        let worker_map = worker_blocks.get_mut(&worker).ok_or_else(|| {
Yan Ru Pei's avatar
Yan Ru Pei committed
305
306
307
308
309
310
311
312
313
314
            tracing::warn!(
                worker_id = worker.worker_id.to_string(),
                dp_rank = worker.dp_rank,
                event_id,
                block_hashes = ?seq_hashes,
                "Failed to find worker blocks to remove"
            );
            KvCacheEventError::BlockNotFound
        })?;

315
        let mut num_removed_blocks = 0;
Yan Ru Pei's avatar
Yan Ru Pei committed
316
317
318
319
320
321
322
323
324
325

        for seq_hash in seq_hashes {
            let Some((position, local_hash)) = worker_map.remove(seq_hash) else {
                tracing::warn!(
                    worker_id = worker.worker_id.to_string(),
                    dp_rank = worker.dp_rank,
                    event_id,
                    block_hash = ?seq_hash,
                    "Failed to find block to remove; skipping remove operation"
                );
326
327
328
329
330

                if let Some(size) = self.tree_sizes.get(&worker) {
                    size.fetch_sub(num_removed_blocks, Ordering::Relaxed);
                }

Yan Ru Pei's avatar
Yan Ru Pei committed
331
332
333
                return Err(KvCacheEventError::BlockNotFound);
            };

334
            if let Some(mut entry) = self.index.get_mut(&(position, local_hash)) {
Yan Ru Pei's avatar
Yan Ru Pei committed
335
336
                let _ = entry.remove(*seq_hash, worker);
            }
337
338
339
340
341
342

            num_removed_blocks += 1;
        }

        if let Some(size) = self.tree_sizes.get(&worker) {
            size.fetch_sub(num_removed_blocks, Ordering::Relaxed);
Yan Ru Pei's avatar
Yan Ru Pei committed
343
344
345
346
347
348
349
350
        }

        Ok(())
    }

    /// Clear all blocks for a specific worker_id (all dp_ranks), but keep worker tracked.
    /// Static version for use in worker threads.
    fn clear_worker_blocks_impl(
351
352
        &self,
        worker_blocks: &mut FxHashMap<WorkerWithDpRank, LevelIndex>,
Yan Ru Pei's avatar
Yan Ru Pei committed
353
354
        worker_id: WorkerId,
    ) {
355
        self.remove_or_clear_worker_blocks_impl(worker_blocks, worker_id, true);
Yan Ru Pei's avatar
Yan Ru Pei committed
356
357
    }

358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
    fn remove_worker_dp_rank_impl(
        &self,
        worker_blocks: &mut FxHashMap<WorkerWithDpRank, LevelIndex>,
        worker_id: WorkerId,
        dp_rank: DpRank,
    ) {
        let key = WorkerWithDpRank { worker_id, dp_rank };
        if let Some(worker_map) = worker_blocks.remove(&key) {
            for (seq_hash, (position, local_hash)) in worker_map.iter() {
                if let Some(mut entry) = self.index.get_mut(&(*position, *local_hash)) {
                    let _ = entry.remove(*seq_hash, key);
                }
            }
            self.tree_sizes.remove(&key);
        }
    }

Yan Ru Pei's avatar
Yan Ru Pei committed
375
376
377
378
    /// Helper function to remove or clear blocks for a worker.
    /// If `keep_worker` is true, the worker remains tracked with empty blocks.
    /// If `keep_worker` is false, the worker is completely removed.
    fn remove_or_clear_worker_blocks_impl(
379
380
        &self,
        worker_blocks: &mut FxHashMap<WorkerWithDpRank, LevelIndex>,
Yan Ru Pei's avatar
Yan Ru Pei committed
381
382
383
384
        worker_id: WorkerId,
        keep_worker: bool,
    ) {
        let workers: Vec<WorkerWithDpRank> = worker_blocks
385
386
387
            .iter()
            .filter(|entry| entry.0.worker_id == worker_id)
            .map(|entry| *entry.0)
Yan Ru Pei's avatar
Yan Ru Pei committed
388
389
390
            .collect();

        for worker in workers {
391
392
393
            if let Some(worker_map) = worker_blocks.remove(&worker) {
                for (seq_hash, (position, local_hash)) in worker_map.iter() {
                    if let Some(mut entry) = self.index.get_mut(&(*position, *local_hash)) {
394
                        let _ = entry.remove(*seq_hash, worker);
Yan Ru Pei's avatar
Yan Ru Pei committed
395
396
397
398
399
                    }
                }
            }

            if keep_worker {
400
401
402
403
404
405
406
407
408
                // Re-insert worker with empty map to keep it tracked
                worker_blocks.insert(worker, FxHashMap::default());
                // Reset tree size to 0 but keep the entry so scoring remains consistent.
                if let Some(size) = self.tree_sizes.get(&worker) {
                    size.store(0, Ordering::Relaxed);
                }
            } else {
                // Fully remove the worker from tree_sizes.
                self.tree_sizes.remove(&worker);
Yan Ru Pei's avatar
Yan Ru Pei committed
409
410
411
            }
        }
    }
412
413
414
415
416
417
418
419
420

    fn dump_events(
        &self,
        worker_blocks: &FxHashMap<WorkerWithDpRank, LevelIndex>,
    ) -> Vec<RouterEvent> {
        let mut events = Vec::new();
        let mut event_id = 0u64;

        for (worker, worker_map) in worker_blocks.iter() {
421
            // Collect (position, local_hash, seq_hash) and sort by position.
422
423
424
425
426
427
428
429
430
            let mut blocks: Vec<_> = worker_map
                .iter()
                .map(|(seq_hash, (pos, local_hash))| (*pos, *local_hash, *seq_hash))
                .collect();
            blocks.sort_unstable_by_key(|(pos, _, _)| *pos);

            for (pos, local_hash, seq_hash) in blocks {
                events.push(RouterEvent {
                    worker_id: worker.worker_id,
431
                    storage_tier: crate::protocols::StorageTier::Device,
432
433
434
                    event: KvCacheEvent {
                        event_id,
                        data: KvCacheEventData::Stored(KvCacheStoreData {
435
436
                            parent_hash: None,
                            start_position: Some(pos as u32),
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
                            blocks: vec![KvCacheStoredBlockData {
                                block_hash: seq_hash,
                                tokens_hash: local_hash,
                                mm_extra_info: None,
                            }],
                        }),
                        dp_rank: worker.dp_rank,
                    },
                });
                event_id += 1;
            }
        }

        events
    }
Yan Ru Pei's avatar
Yan Ru Pei committed
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
}

// -----------------------------------------------------------------------------
// Jump-based search methods (associated functions for use in worker threads)
// -----------------------------------------------------------------------------

impl PositionalIndexer {
    /// Compute sequence hash incrementally from previous hash and current local hash.
    #[inline]
    fn compute_next_seq_hash(prev_seq_hash: u64, current_local_hash: u64) -> u64 {
        let mut bytes = [0u8; 16];

        bytes[..8].copy_from_slice(&prev_seq_hash.to_le_bytes());
        bytes[8..].copy_from_slice(&current_local_hash.to_le_bytes());

        crate::protocols::compute_hash(&bytes)
    }

    /// Ensure seq_hashes is computed up to and including target_pos.
    /// Lazily extends the seq_hashes vector as needed.
    #[inline]
    fn ensure_seq_hash_computed(
        seq_hashes: &mut Vec<ExternalSequenceBlockHash>,
        target_pos: usize,
        sequence: &[LocalBlockHash],
    ) {
        while seq_hashes.len() <= target_pos {
            let pos = seq_hashes.len();
            if pos == 0 {
                // First block's seq_hash equals its local_hash
                seq_hashes.push(ExternalSequenceBlockHash::from(sequence[0].0));
            } else {
                let prev_seq_hash = seq_hashes[pos - 1].0;
                let current_local_hash = sequence[pos].0;
                let next_hash = Self::compute_next_seq_hash(prev_seq_hash, current_local_hash);
                seq_hashes.push(ExternalSequenceBlockHash::from(next_hash));
            }
        }
    }

    /// Get workers at a position by verifying both local_hash and seq_hash match.
    ///
    /// Returns None if no workers match at this position.
    /// Always computes and verifies the seq_hash to ensure correctness when
    /// the query may have diverged from stored sequences at earlier positions.
    fn get_workers_lazy(
        &self,
        position: usize,
        local_hash: LocalBlockHash,
        seq_hashes: &mut Vec<ExternalSequenceBlockHash>,
        sequence: &[LocalBlockHash],
503
    ) -> Option<FxHashSet<WorkerWithDpRank>> {
Yan Ru Pei's avatar
Yan Ru Pei committed
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
        let entry = self.index.get(&(position, local_hash))?;

        // Always compute and verify seq_hash to handle divergent queries correctly.
        // Even if there's only one seq_hash entry, the query's seq_hash might differ
        // if the query diverged from the stored sequence at an earlier position.
        Self::ensure_seq_hash_computed(seq_hashes, position, sequence);
        let seq_hash = seq_hashes[position];
        entry.get(seq_hash).cloned()
    }

    fn count_workers_at(
        &self,
        position: usize,
        local_hash: LocalBlockHash,
        seq_hashes: &mut Vec<ExternalSequenceBlockHash>,
        sequence: &[LocalBlockHash],
    ) -> Option<usize> {
        let entry = self.index.get(&(position, local_hash))?;

        // Always compute and verify seq_hash to handle divergent queries correctly.
        // Even if there's only one seq_hash entry, the query's seq_hash might differ
        // if the query diverged from the stored sequence at an earlier position.
        Self::ensure_seq_hash_computed(seq_hashes, position, sequence);
        let seq_hash = seq_hashes[position];
        Some(
            entry
                .get(seq_hash)
                .map(|workers| workers.len())
                .unwrap_or(0),
        )
    }

    /// Scan positions sequentially, updating active set and recording drain scores.
537
    #[expect(clippy::too_many_arguments)]
Yan Ru Pei's avatar
Yan Ru Pei committed
538
539
540
541
    fn linear_scan_drain(
        &self,
        sequence: &[LocalBlockHash],
        seq_hashes: &mut Vec<ExternalSequenceBlockHash>,
542
        active: &mut FxHashSet<WorkerWithDpRank>,
Yan Ru Pei's avatar
Yan Ru Pei committed
543
544
545
546
547
        scores: &mut OverlapScores,
        lo: usize,
        hi: usize,
        early_exit: bool,
    ) {
548
549
550
        if active.is_empty() {
            return;
        }
Yan Ru Pei's avatar
Yan Ru Pei committed
551
        for pos in lo..hi {
552
553
554
555
            if active.is_empty() {
                break;
            }

Yan Ru Pei's avatar
Yan Ru Pei committed
556
            let Some(entry) = self.index.get(&(pos, sequence[pos])) else {
557
558
559
                for worker in active.drain() {
                    scores.scores.insert(worker, pos as u32);
                }
Yan Ru Pei's avatar
Yan Ru Pei committed
560
561
562
563
                break;
            };

            Self::ensure_seq_hash_computed(seq_hashes, pos, sequence);
564
            let Some(workers) = entry.get(seq_hashes[pos]) else {
565
566
567
                for worker in active.drain() {
                    scores.scores.insert(worker, pos as u32);
                }
568
569
570
                break;
            };

571
572
573
            if workers.len() != active.len() {
                reconcile_active_workers(active, workers, |worker| {
                    scores.scores.insert(worker, pos as u32);
574
575
576
577
578
                });
            }

            if early_exit && !active.is_empty() {
                break;
Yan Ru Pei's avatar
Yan Ru Pei committed
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
            }
        }
    }

    /// Jump-based search to find matches for a sequence of block hashes.
    ///
    /// # Algorithm
    ///
    /// 1. Check first position - initialize active set with matching workers
    /// 2. Initialize seq_hashes with first block's hash (seq_hash[0] = local_hash[0])
    /// 3. Loop: jump by jump_size positions
    ///    - At each jump, check if active workers still match:
    ///      - All match: Continue jumping (skip intermediate positions)
    ///      - None match: Scan range with linear_scan_drain
    ///      - Partial match: Scan range to find exact drain points
    /// 4. Record final scores for remaining active workers
    /// 5. Populate tree_sizes from worker_blocks
    ///
    /// # Arguments
    /// * `index` - The position -> local_hash -> SeqEntry index
    /// * `worker_blocks` - Per-worker reverse lookup for tree sizes
    /// * `local_hashes` - Sequence of LocalBlockHash to match
    /// * `jump_size` - Number of positions to jump at a time
    /// * `early_exit` - If true, stop after finding any match
    fn jump_search_matches(
        &self,
        local_hashes: &[LocalBlockHash],
        early_exit: bool,
    ) -> OverlapScores {
        let mut scores = OverlapScores::new();

        if local_hashes.is_empty() {
            return scores;
        }

        // Lazily computed sequence hashes
615
        let mut seq_hashes: Vec<ExternalSequenceBlockHash> = Vec::with_capacity(local_hashes.len());
Yan Ru Pei's avatar
Yan Ru Pei committed
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636

        // Check first position to initialize active set
        let Some(initial_workers) =
            self.get_workers_lazy(0, local_hashes[0], &mut seq_hashes, local_hashes)
        else {
            return scores;
        };

        let mut active = initial_workers;

        if active.is_empty() {
            return scores;
        }

        if early_exit {
            // For early exit, just record that these workers matched at least position 0
            for worker in &active {
                scores.scores.insert(*worker, 1);
            }
            // Populate tree_sizes
            for worker in scores.scores.keys() {
637
638
639
640
                if let Some(worker_tree_size) = self.tree_sizes.get(worker) {
                    scores
                        .tree_sizes
                        .insert(*worker, worker_tree_size.load(Ordering::Relaxed));
Yan Ru Pei's avatar
Yan Ru Pei committed
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
                }
            }
            return scores;
        }

        let len = local_hashes.len();
        let mut current_pos = 0;

        // Jump through positions
        while current_pos < len - 1 && !active.is_empty() {
            let next_pos = (current_pos + self.jump_size).min(len - 1);

            // Check workers at jump destination
            let num_workers_at_next = self
                .count_workers_at(
                    next_pos,
                    local_hashes[next_pos],
                    &mut seq_hashes,
                    local_hashes,
                )
                .unwrap_or(0);

            if num_workers_at_next == active.len() {
                current_pos = next_pos;
            } else {
                // No active workers match at jump destination
                // Scan the range to find where each worker drained
                self.linear_scan_drain(
                    local_hashes,
                    &mut seq_hashes,
                    &mut active,
                    &mut scores,
                    current_pos + 1,
                    next_pos + 1,
                    false,
                );
                current_pos = next_pos;
            }
        }

        // Record final scores for remaining active workers
        // They matched all positions through the end
        let final_score = len as u32;
        for worker in active {
            scores.scores.insert(worker, final_score);
        }

        for worker in scores.scores.keys() {
689
690
691
692
            if let Some(worker_tree_size) = self.tree_sizes.get(worker) {
                scores
                    .tree_sizes
                    .insert(*worker, worker_tree_size.load(Ordering::Relaxed));
Yan Ru Pei's avatar
Yan Ru Pei committed
693
694
695
696
697
698
            }
        }

        scores
    }
}