Skip to main content

hermes_core/segment/builder/
graph_bisection.rs

1//! Level-synchronized Graph Bisection (BP) for BMP document ordering.
2//!
3//! Based on Dhulipala et al. (KDD 2016) and Mackenzie et al. — the same
4//! algorithm used in Lucene and PISA for document reordering.
5//!
6//! Directly optimizes log-gap cost: docs sharing dimensions end up in the
7//! same BMP blocks, producing tight upper bounds and effective pruning.
8//!
9//! Scheduling follows Mackenzie et al.'s level-synchronized iterative variant:
10//! all partitions at one depth finish before the next depth begins.
11//!
12//! Memory is budgeted: CSR terms plus roughly 32 bytes/document of graph
13//! scratch, with lazily initialized direct-index term-degree arrays.
14
15#[cfg(feature = "native")]
16use rayon::prelude::*;
17const TERM_DEGREE_VALUE_BYTES: usize = std::mem::size_of::<[u32; 2]>();
18const CANDIDATE_ENTRY_BYTES: usize = std::mem::size_of::<(usize, u32)>();
19// Radix selection scans the gain array four times and parallel degree updates
20// require private vocabulary arrays. Quickselect wins decisively below this
21// scale; above it, bounded memory and worker utilization dominate.
22const PARALLEL_BP_MIN_ENTITIES: usize = 1_048_576;
23const MIN_RELATIVE_OBJECTIVE_IMPROVEMENT: f64 = 1e-6;
24const MIN_OBJECTIVE_ITERATIONS: usize = 4;
25const OBJECTIVE_STALL_ITERATIONS: usize = 2;
26
27fn term_degree_bytes(num_terms: usize) -> usize {
28    let bitmap_words = num_terms.div_ceil(64);
29    num_terms
30        .saturating_mul(TERM_DEGREE_VALUE_BYTES)
31        .saturating_add(bitmap_words.saturating_mul(std::mem::size_of::<u64>()))
32        // A reusable lane remembers only bitmap words touched by its previous
33        // partition, so reset is proportional to active terms rather than the
34        // complete vocabulary.
35        .saturating_add(bitmap_words.saturating_mul(std::mem::size_of::<u32>()))
36}
37
38/// Count a dense vocabulary without a cache-line-contended atomic increment
39/// for every posting. Workers own private tables, and the number of tables is
40/// capped by the caller's remaining memory budget. A single plain table is
41/// used when the budget cannot afford useful parallelism.
42#[cfg(feature = "native")]
43trait FrequencyParallelSafe: Sync {}
44#[cfg(feature = "native")]
45impl<T: Sync + ?Sized> FrequencyParallelSafe for T {}
46
47#[cfg(not(feature = "native"))]
48trait FrequencyParallelSafe {}
49#[cfg(not(feature = "native"))]
50impl<T: ?Sized> FrequencyParallelSafe for T {}
51
52fn count_frequencies_bounded<T: FrequencyParallelSafe>(
53    items: &[T],
54    num_terms: usize,
55    available_bytes: usize,
56    count_item: impl Fn(&T, &mut [u32]) + FrequencyParallelSafe,
57) -> Option<Vec<u32>> {
58    if num_terms == 0 {
59        return Some(Vec::new());
60    }
61    let table_bytes = num_terms
62        .checked_mul(std::mem::size_of::<u32>())?
63        .checked_add(std::mem::size_of::<Vec<u32>>())?;
64    let affordable_tables = available_bytes.checked_div(table_bytes)?;
65    if affordable_tables == 0 {
66        return None;
67    }
68
69    #[cfg(feature = "native")]
70    {
71        let lanes = affordable_tables
72            .min(rayon::current_num_threads().max(1))
73            .min(items.len().max(1));
74        if lanes > 1 {
75            let chunk_len = items.len().div_ceil(lanes);
76            return items
77                .par_chunks(chunk_len)
78                .map(|chunk| {
79                    let mut counts = vec![0u32; num_terms];
80                    for item in chunk {
81                        count_item(item, &mut counts);
82                    }
83                    counts
84                })
85                .reduce_with(|mut left, right| {
86                    for (total, count) in left.iter_mut().zip(right) {
87                        *total = total.saturating_add(count);
88                    }
89                    left
90                });
91        }
92    }
93
94    let mut counts = vec![0u32; num_terms];
95    for item in items {
96        count_item(item, &mut counts);
97    }
98    Some(counts)
99}
100
101/// Retain the lowest-frequency eligible dimensions while the frequency table
102/// is live. Both record- and block-level builders use the same bounded policy.
103fn select_frequency_candidates(
104    frequencies: &[u32],
105    min_frequency: usize,
106    max_frequency: usize,
107    candidate_budget_bytes: usize,
108) -> (Vec<(u32, usize)>, bool) {
109    let eligible_count = frequencies
110        .iter()
111        .filter(|&&frequency| {
112            let frequency = frequency as usize;
113            frequency >= min_frequency && frequency <= max_frequency
114        })
115        .count();
116    let capacity = candidate_budget_bytes
117        .checked_div(CANDIDATE_ENTRY_BYTES)
118        .unwrap_or(0)
119        .min(eligible_count);
120    let mut candidates = std::collections::BinaryHeap::with_capacity(capacity);
121    for (term_id, &frequency) in frequencies.iter().enumerate() {
122        let frequency = frequency as usize;
123        if frequency < min_frequency || frequency > max_frequency {
124            continue;
125        }
126        let candidate = (frequency, term_id as u32);
127        if candidates.len() < capacity {
128            candidates.push(candidate);
129        } else if capacity > 0 && candidate < *candidates.peek().unwrap() {
130            candidates.pop();
131            candidates.push(candidate);
132        }
133    }
134    let selected = candidates
135        .into_vec()
136        .into_iter()
137        .map(|(frequency, term_id)| (term_id, frequency))
138        .collect();
139    (selected, capacity < eligible_count)
140}
141
142struct CandidateFit {
143    estimated_bytes: usize,
144    retained_postings: usize,
145    dropped: usize,
146}
147
148/// Apply the shared forward-index memory model, preferring low-frequency
149/// dimensions because they add the least CSR storage and the strongest
150/// clustering signal.
151fn fit_candidates_to_budget(
152    candidates: &mut Vec<(u32, usize)>,
153    fixed_bytes: usize,
154    memory_budget_bytes: usize,
155) -> CandidateFit {
156    let total_postings = candidates.iter().fold(0usize, |total, (_, frequency)| {
157        total.saturating_add(*frequency)
158    });
159    let estimated_bytes = total_postings
160        .saturating_mul(std::mem::size_of::<u32>())
161        .saturating_add(fixed_bytes)
162        .saturating_add(candidates.len().saturating_mul(CANDIDATE_ENTRY_BYTES))
163        .saturating_add(term_degree_bytes(candidates.len()));
164    if estimated_bytes <= memory_budget_bytes || candidates.is_empty() {
165        return CandidateFit {
166            estimated_bytes,
167            retained_postings: total_postings,
168            dropped: 0,
169        };
170    }
171
172    candidates.sort_by_key(|&(_, frequency)| frequency);
173    let mut used_bytes = fixed_bytes;
174    let mut retained_postings = 0usize;
175    let mut keep = 0usize;
176    for &(_, frequency) in candidates.iter() {
177        let term_bytes = frequency
178            .saturating_mul(std::mem::size_of::<u32>())
179            .saturating_add(TERM_DEGREE_VALUE_BYTES + 1)
180            .saturating_add(CANDIDATE_ENTRY_BYTES);
181        if term_bytes > memory_budget_bytes.saturating_sub(used_bytes) {
182            break;
183        }
184        used_bytes = used_bytes.saturating_add(term_bytes);
185        retained_postings = retained_postings.saturating_add(frequency);
186        keep += 1;
187    }
188    let dropped = candidates.len() - keep;
189    candidates.truncate(keep);
190    // BinaryHeap::into_vec retains its original capacity. Release dropped
191    // candidates before allocating the remap, CSR, and graph scratch.
192    candidates.shrink_to_fit();
193    CandidateFit {
194        estimated_bytes,
195        retained_postings,
196        dropped,
197    }
198}
199
200fn parallel_bisect_lanes(
201    memory_budget_bytes: usize,
202    non_degree_bytes: usize,
203    num_terms: usize,
204) -> usize {
205    let per_node = term_degree_bytes(num_terms).max(1);
206    let affordable_nodes = memory_budget_bytes
207        .saturating_sub(non_degree_bytes)
208        .checked_div(per_node)
209        .unwrap_or(0)
210        .max(1);
211    #[cfg(feature = "native")]
212    let worker_limit = rayon::current_num_threads().max(1);
213    #[cfg(not(feature = "native"))]
214    let worker_limit = 1usize;
215    affordable_nodes.min(worker_limit)
216}
217
218/// Per-partition left/right term degrees with direct compact-term indexing.
219///
220/// Recursive BP used to zero two `num_terms`-long vectors at every node. At
221/// 100k vocabulary terms and hundreds of thousands of fine partitions, that
222/// turns into a large amount of memory traffic unrelated to actual postings.
223/// A one-bit initialization map lets us retain array-speed lookups while only
224/// touching degree slots present in the current partition.
225struct TermDegrees {
226    values: Vec<std::mem::MaybeUninit<[u32; 2]>>,
227    initialized: Vec<u64>,
228    touched_words: Vec<u32>,
229}
230
231impl TermDegrees {
232    fn new(num_terms: usize) -> Self {
233        let bitmap_words = num_terms.div_ceil(64);
234        let mut values = Vec::with_capacity(num_terms);
235        values.resize_with(num_terms, std::mem::MaybeUninit::uninit);
236        Self {
237            values,
238            initialized: vec![0; bitmap_words],
239            touched_words: Vec::with_capacity(bitmap_words),
240        }
241    }
242
243    /// Reuse this vocabulary-sized lane for another partition without
244    /// clearing the complete initialization bitmap.
245    fn reset(&mut self) {
246        for word in self.touched_words.drain(..) {
247            self.initialized[word as usize] = 0;
248        }
249    }
250
251    fn sort_touched_words(&mut self) {
252        self.touched_words.sort_unstable();
253    }
254
255    #[inline]
256    fn entry_mut(&mut self, term: usize) -> &mut [u32; 2] {
257        let word = term / 64;
258        let mask = 1u64 << (term % 64);
259        if self.initialized[word] & mask == 0 {
260            if self.initialized[word] == 0 {
261                self.touched_words.push(word as u32);
262            }
263            self.values[term].write([0, 0]);
264            self.initialized[word] |= mask;
265        }
266        // SAFETY: the bit above is set only after writing this exact slot.
267        unsafe { self.values[term].assume_init_mut() }
268    }
269
270    #[inline]
271    fn get(&self, term: usize) -> [u32; 2] {
272        let word = term / 64;
273        let mask = 1u64 << (term % 64);
274        if self.initialized[word] & mask == 0 {
275            return [0, 0];
276        }
277        // SAFETY: an initialized bit is published only after the slot write;
278        // scoring reads degrees after construction, with no concurrent writes.
279        unsafe { *self.values[term].assume_init_ref() }
280    }
281
282    fn merge_from(&mut self, other: &Self) {
283        for &word_idx in &other.touched_words {
284            let word_idx = word_idx as usize;
285            let mut pending = other.initialized[word_idx];
286            while pending != 0 {
287                let bit = pending.trailing_zeros() as usize;
288                let term = word_idx * 64 + bit;
289                // SAFETY: `pending` is derived from the initialized bitmap.
290                let [left, right] = unsafe { *other.values[term].assume_init_ref() };
291                let entry = self.entry_mut(term);
292                entry[0] += left;
293                entry[1] += right;
294                pending &= pending - 1;
295            }
296        }
297    }
298
299    /// Apply directional movement counts accumulated in a reusable lane.
300    ///
301    /// Each entry is `[right_to_left, left_to_right]`. Keeping two unsigned
302    /// counters lets partition workers reuse the exact same storage as degree
303    /// construction instead of allocating a separate signed delta table.
304    fn apply_moves_to(&self, degrees: &mut Self) {
305        for &word_idx in &self.touched_words {
306            let word_idx = word_idx as usize;
307            let mut pending = self.initialized[word_idx];
308            while pending != 0 {
309                let bit = pending.trailing_zeros() as usize;
310                let term = word_idx * 64 + bit;
311                // SAFETY: `pending` is derived from the initialized bitmap.
312                let [right_to_left, left_to_right] =
313                    unsafe { *self.values[term].assume_init_ref() };
314                debug_assert!(
315                    degrees.initialized[word_idx] & (1u64 << bit) != 0,
316                    "a moved term must already exist in the partition degrees"
317                );
318                let degree = degrees.entry_mut(term);
319                let new_left =
320                    i64::from(degree[0]) + i64::from(right_to_left) - i64::from(left_to_right);
321                let new_right =
322                    i64::from(degree[1]) + i64::from(left_to_right) - i64::from(right_to_left);
323                debug_assert!(new_left >= 0 && new_right >= 0);
324                debug_assert!(new_left <= i64::from(u32::MAX));
325                debug_assert!(new_right <= i64::from(u32::MAX));
326                degree[0] = new_left as u32;
327                degree[1] = new_right as u32;
328                pending &= pending - 1;
329            }
330        }
331    }
332
333    /// Exact bisection objective optimized by the BP gain approximation.
334    ///
335    /// This returns the negative assignment-dependent BiMLogA bisection cost
336    /// from Dhulipala et al., so a larger value means lower cost. Keeping the
337    /// partition-size term matters for odd-sized partitions, whose halves
338    /// differ by one entity.
339    fn bisection_objective(&self, left_size: usize, right_size: usize, log_table: &[f32]) -> f64 {
340        let mut objective = 0.0f64;
341        let side_log = [
342            fast_log2_lookup(left_size, log_table) as f64,
343            fast_log2_lookup(right_size, log_table) as f64,
344        ];
345        for &word_idx in &self.touched_words {
346            let word_idx = word_idx as usize;
347            let mut pending = self.initialized[word_idx];
348            while pending != 0 {
349                let bit = pending.trailing_zeros() as usize;
350                let term = word_idx * 64 + bit;
351                // SAFETY: `pending` is derived from the initialized bitmap.
352                let [left, right] = unsafe { *self.values[term].assume_init_ref() };
353                for (side, count) in [left, right].into_iter().enumerate() {
354                    if count > 0 {
355                        objective += count as f64
356                            * (fast_log2_lookup(count as usize + 1, log_table) as f64
357                                - side_log[side]);
358                    }
359                }
360                pending &= pending - 1;
361            }
362        }
363        objective
364    }
365}
366
367// ── Forward index (CSR) ──────────────────────────────────────────────────
368
369/// Forward index in CSR format: doc `d`'s terms are `terms[offsets[d]..offsets[d+1]]`.
370///
371/// Term IDs are remapped to compact range `0..num_terms` for flat-array degree tracking.
372pub(crate) struct ForwardIndex {
373    terms: Vec<u32>,
374    /// u64, not u32: a 58M-doc / ~85-dims-per-doc reorder pass carries ~4.9B
375    /// postings — u32 prefix sums wrapped and the CSR carving panicked
376    /// (prod 2026-07-14, "mid > len"). The old 8 GB memory budget masked it
377    /// by dropping dims below the u32 limit.
378    offsets: Vec<u64>,
379    pub num_terms: usize,
380    /// Exact number of simultaneous vocabulary-sized degree-array lanes
381    /// allowed by the memory budget. Recursion divides this allowance between
382    /// children without rounding non-power-of-two worker pools down.
383    parallel_bisect_lanes: usize,
384    /// True when the configured memory limit forced graph signal to be
385    /// discarded. Callers must not report the resulting order as fully
386    /// converged: a later pass with a larger budget may still improve it.
387    budget_limited: bool,
388}
389
390/// Build CSR offsets (prefix sums) from per-entity counts. u64 output — the
391/// sum of counts legitimately exceeds u32::MAX on large reorder passes.
392fn build_csr_offsets(counts: &[u32]) -> Vec<u64> {
393    let mut offsets = Vec::with_capacity(counts.len() + 1);
394    offsets.push(0u64);
395    for &c in counts {
396        offsets.push(offsets.last().unwrap() + c as u64);
397    }
398    offsets
399}
400
401impl ForwardIndex {
402    #[inline]
403    pub fn num_docs(&self) -> usize {
404        if self.offsets.is_empty() {
405            0
406        } else {
407            self.offsets.len() - 1
408        }
409    }
410
411    #[inline]
412    fn doc_terms(&self, doc: usize) -> &[u32] {
413        let start = self.offsets[doc] as usize;
414        let end = self.offsets[doc + 1] as usize;
415        &self.terms[start..end]
416    }
417
418    /// Total postings in the forward index.
419    pub fn total_postings(&self) -> u64 {
420        self.offsets.last().copied().unwrap_or(0)
421    }
422
423    #[inline]
424    pub fn budget_limited(&self) -> bool {
425        self.budget_limited
426    }
427}
428
429/// Build virtual→real and real→virtual vid maps from a BMP doc map.
430///
431/// A virtual slot is real iff its doc-map entry is not the `u32::MAX` padding
432/// sentinel. Realness must come from the doc map itself: block-copy merged
433/// segments carry each source's tail padding as *interior* padding, so
434/// `vid < num_real_docs` does NOT identify real docs there.
435///
436/// Returns `(virtual_to_real, real_to_virtual)` where `virtual_to_real[vid]`
437/// is the dense real index or `u32::MAX` for padding.
438pub(crate) fn build_vid_maps(
439    bmp: &crate::segment::reader::bmp::BmpIndex,
440) -> crate::Result<(Vec<u32>, Vec<u32>)> {
441    let ids = bmp.doc_map_ids_slice();
442    let num_virtual = bmp.num_virtual_docs as usize;
443    let expected_real = bmp.num_real_docs() as usize;
444    let mut virtual_to_real = vec![u32::MAX; num_virtual];
445    let mut real_to_virtual = Vec::with_capacity(expected_real);
446    debug_assert_eq!(ids.len(), virtual_to_real.len() * 4);
447    bmp.visit_real_slots_for_rewrite(|vid| {
448        virtual_to_real[vid] = real_to_virtual.len() as u32;
449        real_to_virtual.push(vid as u32);
450    })?;
451    Ok((virtual_to_real, real_to_virtual))
452}
453
454/// One (source, block) unit of forward-index construction. Because
455/// [`build_vid_maps`] assigns real ids in ascending vid order, a block's real
456/// docs form the contiguous per-source range
457/// `real_start..real_start + real_len` — blocks can be processed in parallel
458/// with disjoint output slices.
459struct BlockJob {
460    src: u32,
461    block_id: u32,
462    /// Per-source real index of the block's first real doc.
463    real_start: u32,
464    /// Number of real (non-padding) docs in the block.
465    real_len: u32,
466}
467
468/// Enumerate jobs in (source, block) order — cumulative `real_len` tiles the
469/// global real-id space `0..total_docs` exactly.
470fn build_block_jobs(
471    bmps: &[&crate::segment::reader::bmp::BmpIndex],
472    vid_maps: &[(Vec<u32>, Vec<u32>)],
473) -> Vec<BlockJob> {
474    let total_blocks: usize = bmps.iter().map(|b| b.num_blocks as usize).sum();
475    let mut jobs = Vec::with_capacity(total_blocks);
476    for (src, (bmp, (v2r, _))) in bmps.iter().zip(vid_maps).enumerate() {
477        let block_size = bmp.bmp_block_size as usize;
478        let mut real_cursor = 0u32;
479        for block_id in 0..bmp.num_blocks as usize {
480            let vid_start = block_id * block_size;
481            let vid_end = ((block_id + 1) * block_size).min(v2r.len());
482            let real_len = v2r[vid_start..vid_end]
483                .iter()
484                .filter(|&&r| r != u32::MAX)
485                .count() as u32;
486            jobs.push(BlockJob {
487                src: src as u32,
488                block_id: block_id as u32,
489                real_start: real_cursor,
490                real_len,
491            });
492            real_cursor += real_len;
493        }
494    }
495    jobs
496}
497
498/// Build forward index from BmpIndex sources (single or multi-source).
499///
500/// Documents are identified by dense *real* indices assigned sequentially
501/// across sources: source 0 gets 0..n0, source 1 gets n0..n0+n1, etc., where
502/// each n is the source's real (non-padding) doc count derived from its doc
503/// map via [`build_vid_maps`].
504///
505/// Filters dims with doc_freq outside `[min_doc_freq, max_doc_freq]`.
506/// If the estimated forward index memory exceeds `memory_budget_bytes`, the
507/// highest-frequency dims are dropped to stay within budget. This prevents OOM
508/// for huge segments at the cost of slightly reduced reorder quality.
509///
510/// Remaps term IDs to compact range for flat-array degree tracking.
511#[cfg(test)]
512pub(crate) fn build_forward_index_from_bmps(
513    bmps: &[&crate::segment::reader::bmp::BmpIndex],
514    min_doc_freq: usize,
515    max_doc_freq: usize,
516    memory_budget_bytes: usize,
517) -> crate::Result<ForwardIndex> {
518    let vid_maps: Vec<(Vec<u32>, Vec<u32>)> = bmps
519        .iter()
520        .map(|bmp| build_vid_maps(bmp))
521        .collect::<crate::Result<_>>()?;
522    Ok(build_forward_index_from_bmps_with_maps(
523        bmps,
524        &vid_maps,
525        min_doc_freq,
526        max_doc_freq,
527        memory_budget_bytes,
528    ))
529}
530
531/// Variant for reorder callers that already need the virtual/real maps during
532/// output encoding. Reusing them avoids a second full document-map scan and a
533/// duplicate real-to-virtual allocation on very large segments.
534pub(crate) fn build_forward_index_from_bmps_with_maps(
535    bmps: &[&crate::segment::reader::bmp::BmpIndex],
536    vid_maps: &[(Vec<u32>, Vec<u32>)],
537    min_doc_freq: usize,
538    max_doc_freq: usize,
539    memory_budget_bytes: usize,
540) -> ForwardIndex {
541    debug_assert_eq!(bmps.len(), vid_maps.len());
542    let total_docs: usize = vid_maps.iter().map(|(_, r2v)| r2v.len()).sum();
543
544    if total_docs == 0 {
545        return ForwardIndex {
546            terms: Vec::new(),
547            offsets: Vec::new(),
548            num_terms: 0,
549            parallel_bisect_lanes: 1,
550            budget_limited: false,
551        };
552    }
553
554    // Job list: one entry per (source, block). Real ids are assigned in
555    // ascending vid order (see build_vid_maps), so each block owns a
556    // contiguous real-id range — every phase below can process blocks in
557    // parallel, writing disjoint slices.
558    let jobs = build_block_jobs(bmps, vid_maps);
559
560    // Phase 1: count document frequencies in bounded worker-local tables.
561    // Shared atomics made popular dimensions a cache-coherence bottleneck;
562    // private dense tables avoid that while retaining an exact memory cap.
563    let max_dims = bmps
564        .iter()
565        .map(|bmp| bmp.dims() as usize)
566        .max()
567        .unwrap_or(0);
568    let jobs_bytes = jobs
569        .len()
570        .saturating_mul(std::mem::size_of::<BlockJob>().saturating_add(40));
571    let frequency_bytes = max_dims.saturating_mul(std::mem::size_of::<u32>());
572    if frequency_bytes > memory_budget_bytes.saturating_sub(jobs_bytes) {
573        log::warn!(
574            "[reorder] memory budget {} cannot hold the {} dimension-frequency table; using identity order",
575            crate::format_bytes(memory_budget_bytes as u64),
576            crate::format_bytes(frequency_bytes as u64),
577        );
578        return ForwardIndex {
579            terms: Vec::new(),
580            offsets: Vec::new(),
581            num_terms: 0,
582            parallel_bisect_lanes: 1,
583            budget_limited: true,
584        };
585    }
586    let Some(dim_df) = count_frequencies_bounded(
587        &jobs,
588        max_dims,
589        memory_budget_bytes.saturating_sub(jobs_bytes),
590        |job, counts| {
591            let bmp = bmps[job.src as usize];
592            let (v2r, _) = &vid_maps[job.src as usize];
593            let block_size = bmp.bmp_block_size as usize;
594            for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
595                let mut frequency = 0u32;
596                for posting in postings {
597                    let vid = job.block_id as usize * block_size + posting.local_slot as usize;
598                    if v2r.get(vid).is_some_and(|&real| real != u32::MAX) && posting.impact > 0 {
599                        frequency = frequency.saturating_add(1);
600                    }
601                }
602                if frequency > 0
603                    && let Some(total) = counts.get_mut(dim_id as usize)
604                {
605                    *total = total.saturating_add(frequency);
606                }
607            }
608        },
609    ) else {
610        log::warn!(
611            "[reorder] memory budget {} cannot hold a bounded dimension-frequency table; using identity order",
612            crate::format_bytes(memory_budget_bytes as u64),
613        );
614        return ForwardIndex {
615            terms: Vec::new(),
616            offsets: Vec::new(),
617            num_terms: 0,
618            parallel_bisect_lanes: 1,
619            budget_limited: true,
620        };
621    };
622
623    // Retain the lowest-frequency candidates in a bounded heap while the
624    // frequency table is live. This makes candidate discovery itself obey the
625    // configured limit even for extremely large vocabularies.
626    let (mut eligible, mut budget_limited) = select_frequency_candidates(
627        &dim_df,
628        min_doc_freq,
629        max_doc_freq,
630        memory_budget_bytes
631            .saturating_sub(jobs_bytes)
632            .saturating_sub(frequency_bytes),
633    );
634    drop(dim_df);
635
636    // Memory budget: estimate forward index + bisection scratch.
637    // Includes jobs/slice descriptors, dense remap, all per-document scratch,
638    // and at least one exact TermDegrees allocation.
639    let entity_scratch_bytes = total_docs.saturating_mul(32);
640    let remap_bytes = max_dims.saturating_mul(4);
641    let fixed_bytes = entity_scratch_bytes
642        .saturating_add(remap_bytes)
643        .saturating_add(jobs_bytes);
644    let fit = fit_candidates_to_budget(&mut eligible, fixed_bytes, memory_budget_bytes);
645    if fit.dropped > 0 {
646        budget_limited = true;
647        log::warn!(
648            "[reorder] memory budget {}: estimated {}, dropped {} highest-df dims, keeping {} ({} postings)",
649            crate::format_bytes(memory_budget_bytes as u64),
650            crate::format_bytes(fit.estimated_bytes as u64),
651            fit.dropped,
652            eligible.len(),
653            fit.retained_postings,
654        );
655    }
656
657    if eligible.is_empty() {
658        // The caller emits an identity permutation when there is no graph
659        // signal. Avoid allocating per-document counts and u64 offsets only
660        // to discover that the terms array is empty, especially when the
661        // configured budget is below the fixed document scratch cost.
662        return ForwardIndex {
663            terms: Vec::new(),
664            offsets: Vec::new(),
665            num_terms: 0,
666            parallel_bisect_lanes: 1,
667            budget_limited,
668        };
669    }
670
671    let mut term_remap = vec![u32::MAX; max_dims];
672    for (compact_id, &(dim_id, _)) in eligible.iter().enumerate() {
673        term_remap[dim_id as usize] = compact_id as u32;
674    }
675    let num_active_terms = eligible.len();
676    // Jobs, candidate metadata, and the dense remap are construction scratch
677    // and are gone before graph recursion. Admit lanes against graph-resident
678    // CSR/entity scratch so non-power-of-two pools are not needlessly rounded
679    // down under a tight budget.
680    let non_degree_bytes = entity_scratch_bytes.saturating_add(
681        fit.retained_postings
682            .saturating_mul(std::mem::size_of::<u32>()),
683    );
684    let parallel_bisect_lanes =
685        parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_active_terms);
686    drop(eligible);
687
688    // Phase 2: count terms per doc (filtered) — per-block disjoint slices
689    let mut counts = vec![0u32; total_docs];
690    let fill_block_counts = |job: &BlockJob, out: &mut [u32]| {
691        let bmp = bmps[job.src as usize];
692        let (v2r, _) = &vid_maps[job.src as usize];
693        let block_size = bmp.bmp_block_size as usize;
694        for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
695            if term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX) == u32::MAX {
696                continue;
697            }
698            for p in postings {
699                let vid = job.block_id as usize * block_size + p.local_slot as usize;
700                let Some(&real) = v2r.get(vid) else {
701                    continue;
702                };
703                if real != u32::MAX && p.impact > 0 {
704                    out[(real - job.real_start) as usize] += 1;
705                }
706            }
707        }
708    };
709    {
710        let mut slices: Vec<(&BlockJob, &mut [u32])> = Vec::with_capacity(jobs.len());
711        let mut rest: &mut [u32] = &mut counts;
712        for job in &jobs {
713            let (head, tail) = rest.split_at_mut(job.real_len as usize);
714            slices.push((job, head));
715            rest = tail;
716        }
717        #[cfg(feature = "native")]
718        slices
719            .into_par_iter()
720            .for_each(|(job, out)| fill_block_counts(job, out));
721        #[cfg(not(feature = "native"))]
722        for (job, out) in slices {
723            fill_block_counts(job, out);
724        }
725    }
726
727    // Phase 3: build CSR offsets (u64 — sums exceed u32::MAX at scale)
728    let offsets = build_csr_offsets(&counts);
729    let total = *offsets.last().unwrap() as usize;
730    drop(counts);
731
732    // Phase 4: fill terms (compact IDs) — each block writes the contiguous
733    // terms range covering its real docs; per-doc write cursors are local.
734    let mut terms = vec![0u32; total];
735    let fill_block_terms = |job: &BlockJob, global_real_start: usize, out: &mut [u32]| {
736        let bmp = bmps[job.src as usize];
737        let (v2r, _) = &vid_maps[job.src as usize];
738        let block_size = bmp.bmp_block_size as usize;
739        // local_slot is u8, so a block never holds more than 256 real docs
740        assert!(job.real_len as usize <= 256, "BMP block exceeds 256 docs");
741        let mut cursor = [0u32; 256];
742        let base = offsets[global_real_start] as usize;
743        for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
744            let compact = term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX);
745            if compact == u32::MAX {
746                continue;
747            }
748            for p in postings {
749                let vid = job.block_id as usize * block_size + p.local_slot as usize;
750                let Some(&real) = v2r.get(vid) else {
751                    continue;
752                };
753                if real != u32::MAX && p.impact > 0 {
754                    let local = (real - job.real_start) as usize;
755                    let pos =
756                        offsets[global_real_start + local] as usize - base + cursor[local] as usize;
757                    out[pos] = compact;
758                    cursor[local] += 1;
759                }
760            }
761        }
762    };
763    {
764        let mut slices: Vec<(&BlockJob, usize, &mut [u32])> = Vec::with_capacity(jobs.len());
765        let mut rest: &mut [u32] = &mut terms;
766        let mut global_real = 0usize;
767        for job in &jobs {
768            let len =
769                (offsets[global_real + job.real_len as usize] - offsets[global_real]) as usize;
770            let (head, tail) = rest.split_at_mut(len);
771            slices.push((job, global_real, head));
772            rest = tail;
773            global_real += job.real_len as usize;
774        }
775        #[cfg(feature = "native")]
776        slices
777            .into_par_iter()
778            .for_each(|(job, g, out)| fill_block_terms(job, g, out));
779        #[cfg(not(feature = "native"))]
780        for (job, g, out) in slices {
781            fill_block_terms(job, g, out);
782        }
783    }
784
785    ForwardIndex {
786        terms,
787        offsets,
788        num_terms: num_active_terms,
789        parallel_bisect_lanes,
790        budget_limited,
791    }
792}
793
794/// Build a forward index over BLOCKS (one entity per block, its terms = the
795/// block's header dim list). Used by block-level reorder: BP over blocks is
796/// ~block_size× cheaper than over records and only needs to decide superblock
797/// assignment. Blocks are numbered globally across sources in source order.
798///
799/// Dims appearing in fewer than 2 blocks carry no clustering signal and are
800/// dropped; the memory budget applies as in the record-level builder.
801pub(crate) fn build_forward_index_from_blocks(
802    bmps: &[&crate::segment::reader::bmp::BmpIndex],
803    memory_budget_bytes: usize,
804) -> ForwardIndex {
805    let total_blocks: usize = bmps.iter().map(|b| b.num_blocks as usize).sum();
806    if total_blocks == 0 {
807        return ForwardIndex {
808            terms: Vec::new(),
809            offsets: Vec::new(),
810            num_terms: 0,
811            parallel_bisect_lanes: 1,
812            budget_limited: false,
813        };
814    }
815
816    // (source, block) pairs in global block order — the parallel unit.
817    let blocks: Vec<(u32, u32)> = bmps
818        .iter()
819        .enumerate()
820        .flat_map(|(src, bmp)| (0..bmp.num_blocks).map(move |b| (src as u32, b)))
821        .collect();
822
823    // Phase 1: bounded worker-local frequency tables. This is the same policy
824    // as record-level BP and avoids hot atomic increments on common terms.
825    let max_dims = bmps
826        .iter()
827        .map(|bmp| bmp.dims() as usize)
828        .max()
829        .unwrap_or(0);
830    let blocks_bytes = blocks
831        .len()
832        .saturating_mul(std::mem::size_of::<(u32, u32)>().saturating_add(32));
833    let frequency_bytes = max_dims.saturating_mul(std::mem::size_of::<u32>());
834    if frequency_bytes > memory_budget_bytes.saturating_sub(blocks_bytes) {
835        log::warn!(
836            "[reorder] block-level frequency table exceeds memory budget; using identity order"
837        );
838        return ForwardIndex {
839            terms: Vec::new(),
840            offsets: Vec::new(),
841            num_terms: 0,
842            parallel_bisect_lanes: 1,
843            budget_limited: true,
844        };
845    }
846    let Some(dim_bf) = count_frequencies_bounded(
847        &blocks,
848        max_dims,
849        memory_budget_bytes.saturating_sub(blocks_bytes),
850        |&(src, block_id), counts| {
851            for (dim_id, _, _) in bmps[src as usize].iter_block_terms(block_id) {
852                if let Some(count) = counts.get_mut(dim_id as usize) {
853                    *count = count.saturating_add(1);
854                }
855            }
856        },
857    ) else {
858        log::warn!(
859            "[reorder] block-level frequency table cannot fit its bounded allocation; using identity order"
860        );
861        return ForwardIndex {
862            terms: Vec::new(),
863            offsets: Vec::new(),
864            num_terms: 0,
865            parallel_bisect_lanes: 1,
866            budget_limited: true,
867        };
868    };
869
870    let max_bf = (total_blocks as f64 * 0.9) as usize;
871    let (mut eligible, mut budget_limited) = select_frequency_candidates(
872        &dim_bf,
873        2,
874        max_bf.max(2),
875        memory_budget_bytes
876            .saturating_sub(blocks_bytes)
877            .saturating_sub(frequency_bytes),
878    );
879    drop(dim_bf);
880
881    let entity_scratch_bytes = total_blocks.saturating_mul(32);
882    let remap_bytes = max_dims.saturating_mul(4);
883    let fixed_bytes = entity_scratch_bytes
884        .saturating_add(remap_bytes)
885        .saturating_add(blocks_bytes);
886    let fit = fit_candidates_to_budget(&mut eligible, fixed_bytes, memory_budget_bytes);
887    if fit.dropped > 0 {
888        budget_limited = true;
889        log::warn!(
890            "[reorder] block-level fwd index over budget — dropped {} highest-bf dims",
891            fit.dropped,
892        );
893    }
894
895    if eligible.is_empty() {
896        return ForwardIndex {
897            terms: Vec::new(),
898            offsets: Vec::new(),
899            num_terms: 0,
900            parallel_bisect_lanes: 1,
901            budget_limited,
902        };
903    }
904
905    let mut term_remap = vec![u32::MAX; max_dims];
906    for (compact, &(dim_id, _)) in eligible.iter().enumerate() {
907        term_remap[dim_id as usize] = compact as u32;
908    }
909    let num_terms = eligible.len();
910    let non_degree_bytes = entity_scratch_bytes.saturating_add(
911        fit.retained_postings
912            .saturating_mul(std::mem::size_of::<u32>()),
913    );
914    let parallel_bisect_lanes =
915        parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_terms);
916    drop(eligible);
917
918    // Phase 2+3: counts and CSR fill — one entity per block, so each block
919    // maps to a single count cell and a contiguous terms range.
920    let count_remapped = |&(src, block_id): &(u32, u32)| -> u32 {
921        bmps[src as usize]
922            .iter_block_terms(block_id)
923            .filter(|(dim_id, _, _)| {
924                term_remap
925                    .get(*dim_id as usize)
926                    .copied()
927                    .unwrap_or(u32::MAX)
928                    != u32::MAX
929            })
930            .count() as u32
931    };
932    #[cfg(feature = "native")]
933    let counts: Vec<u32> = blocks.par_iter().map(count_remapped).collect();
934    #[cfg(not(feature = "native"))]
935    let counts: Vec<u32> = blocks.iter().map(count_remapped).collect();
936
937    let offsets = build_csr_offsets(&counts);
938    let total = *offsets.last().unwrap() as usize;
939    drop(counts);
940
941    let mut terms = vec![0u32; total];
942    let fill_block = |&(src, block_id): &(u32, u32), out: &mut [u32]| {
943        let mut n = 0usize;
944        for (dim_id, _, _) in bmps[src as usize].iter_block_terms(block_id) {
945            let compact = term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX);
946            if compact != u32::MAX {
947                out[n] = compact;
948                n += 1;
949            }
950        }
951    };
952    {
953        let mut slices: Vec<(&(u32, u32), &mut [u32])> = Vec::with_capacity(blocks.len());
954        let mut rest: &mut [u32] = &mut terms;
955        for (gb, b) in blocks.iter().enumerate() {
956            let len = (offsets[gb + 1] - offsets[gb]) as usize;
957            let (head, tail) = rest.split_at_mut(len);
958            slices.push((b, head));
959            rest = tail;
960        }
961        #[cfg(feature = "native")]
962        slices
963            .into_par_iter()
964            .for_each(|(b, out)| fill_block(b, out));
965        #[cfg(not(feature = "native"))]
966        for (b, out) in slices {
967            fill_block(b, out);
968        }
969    }
970
971    ForwardIndex {
972        terms,
973        offsets,
974        num_terms,
975        parallel_bisect_lanes,
976        budget_limited,
977    }
978}
979
980// ── Recursive Graph Bisection ────────────────────────────────────────────
981
982/// CPU/depth budget for a BP pass. BP is an anytime algorithm: stopping at
983/// any depth or deadline still yields a valid permutation, and because the
984/// output layout becomes the next pass's input order, repeated budgeted
985/// passes warm-start and deepen (top levels converge in ~0 swaps, the budget
986/// flows to deeper levels).
987#[derive(Clone, Copy, Debug, Default)]
988pub struct BpBudget {
989    /// Stop recursion at partitions of at most this many docs instead of
990    /// descending to block granularity. `None` = full depth. Capping at
991    /// superblock granularity (superblock_size × block_size docs) keeps most
992    /// of the superblock-pruning win at ~⅓ less depth.
993    pub min_partition_docs: Option<usize>,
994    /// Wall-clock cap for the whole BP computation. The pass ends cleanly at
995    /// the deadline with whatever depth it reached (`converged = false`).
996    /// Ignored on wasm (no monotonic clock).
997    pub time_budget: Option<std::time::Duration>,
998}
999
1000impl BpBudget {
1001    /// Unbudgeted: full depth, no deadline.
1002    pub fn full() -> Self {
1003        Self::default()
1004    }
1005}
1006
1007/// Build one partition's term degrees in preallocated lanes.
1008///
1009/// At coarse levels each worker scans a contiguous document range into a
1010/// private lane, then the lanes are reduced into `workspaces[0]`. Recursive
1011/// siblings receive disjoint workspace slices, so the complete pass performs
1012/// exactly the vocabulary-sized allocations admitted before BP starts.
1013fn build_term_degrees(
1014    docs: &[u32],
1015    mid: usize,
1016    fwd: &ForwardIndex,
1017    workspaces: &mut [TermDegrees],
1018    cancellation: Option<&std::sync::atomic::AtomicBool>,
1019) {
1020    debug_assert!(!workspaces.is_empty());
1021    let build_range = |degrees: &mut TermDegrees, start: usize, chunk: &[u32]| {
1022        degrees.reset();
1023        for (offset, &doc) in chunk.iter().enumerate() {
1024            if offset.is_multiple_of(1024)
1025                && cancellation
1026                    .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
1027            {
1028                break;
1029            }
1030            let side = usize::from(start + offset >= mid);
1031            for &term in fwd.doc_terms(doc as usize) {
1032                degrees.entry_mut(term as usize)[side] += 1;
1033            }
1034        }
1035    };
1036
1037    #[cfg(feature = "native")]
1038    {
1039        let workers = workspaces
1040            .len()
1041            .min(docs.len().div_ceil(PARALLEL_BP_MIN_ENTITIES).max(1));
1042        if workers > 1 {
1043            let chunk_len = docs.len().div_ceil(workers);
1044            workspaces[..workers]
1045                .par_iter_mut()
1046                .enumerate()
1047                .for_each(|(worker, degrees)| {
1048                    let start = worker * chunk_len;
1049                    let end = (start + chunk_len).min(docs.len());
1050                    build_range(degrees, start, &docs[start..end]);
1051                });
1052            let (degrees, partials) = workspaces[..workers]
1053                .split_first_mut()
1054                .expect("at least one BP degree workspace");
1055            for partial in partials {
1056                degrees.merge_from(partial);
1057            }
1058            return;
1059        }
1060        build_range(&mut workspaces[0], 0, docs);
1061    }
1062
1063    #[cfg(not(feature = "native"))]
1064    build_range(&mut workspaces[0], 0, docs);
1065}
1066
1067/// Convert `f32::total_cmp` ordering into an unsigned radix key.
1068///
1069/// BP gains are finite, but preserving the complete IEEE total order makes
1070/// the parallel selector exactly equivalent to the former comparator even if
1071/// malformed input ever produces a signed zero, infinity, or NaN.
1072#[inline]
1073fn gain_order_key(gain: f32) -> u32 {
1074    let bits = gain.to_bits();
1075    if bits & 0x8000_0000 != 0 {
1076        !bits
1077    } else {
1078        bits ^ 0x8000_0000
1079    }
1080}
1081
1082/// Find the gain key of the last entity admitted to the left half and the
1083/// number of strictly lower keys. Four byte-histogram passes replace the old
1084/// serial `select_nth_unstable_by` over an 8-byte index per entity.
1085fn select_gain_threshold(
1086    gains: &[f32],
1087    left_count: usize,
1088    allow_inner_parallelism: bool,
1089) -> (u32, usize) {
1090    #[cfg(not(feature = "native"))]
1091    let _ = allow_inner_parallelism;
1092
1093    debug_assert!(left_count > 0 && left_count <= gains.len());
1094    let mut rank_within_prefix = left_count - 1;
1095    let mut strictly_lower = 0usize;
1096    let mut prefix = 0u32;
1097    let mut prefix_mask = 0u32;
1098
1099    for shift in [24u32, 16, 8, 0] {
1100        let histogram = {
1101            #[cfg(feature = "native")]
1102            {
1103                if allow_inner_parallelism && gains.len() >= PARALLEL_BP_MIN_ENTITIES {
1104                    gains
1105                        .par_iter()
1106                        .fold(
1107                            || Box::new([0usize; 256]),
1108                            |mut counts, &gain| {
1109                                let key = gain_order_key(gain);
1110                                if key & prefix_mask == prefix {
1111                                    counts[((key >> shift) & 0xff) as usize] += 1;
1112                                }
1113                                counts
1114                            },
1115                        )
1116                        .reduce(
1117                            || Box::new([0usize; 256]),
1118                            |mut left, right| {
1119                                for (dst, &count) in left.iter_mut().zip(right.iter()) {
1120                                    *dst += count;
1121                                }
1122                                left
1123                            },
1124                        )
1125                } else {
1126                    let mut counts = [0usize; 256];
1127                    for &gain in gains {
1128                        let key = gain_order_key(gain);
1129                        if key & prefix_mask == prefix {
1130                            counts[((key >> shift) & 0xff) as usize] += 1;
1131                        }
1132                    }
1133                    Box::new(counts)
1134                }
1135            }
1136            #[cfg(not(feature = "native"))]
1137            {
1138                let mut counts = [0usize; 256];
1139                for &gain in gains {
1140                    let key = gain_order_key(gain);
1141                    if key & prefix_mask == prefix {
1142                        counts[((key >> shift) & 0xff) as usize] += 1;
1143                    }
1144                }
1145                counts
1146            }
1147        };
1148
1149        let mut before_bucket = 0usize;
1150        let mut selected_bucket = None;
1151        for (bucket, count) in histogram.iter().copied().enumerate() {
1152            if rank_within_prefix < before_bucket + count {
1153                selected_bucket = Some(bucket as u32);
1154                rank_within_prefix -= before_bucket;
1155                strictly_lower += before_bucket;
1156                break;
1157            }
1158            before_bucket += count;
1159        }
1160        let selected_bucket = selected_bucket.expect("BP radix selection lost the requested rank");
1161        prefix |= selected_bucket << shift;
1162        prefix_mask |= 0xffu32 << shift;
1163    }
1164
1165    (prefix, strictly_lower)
1166}
1167
1168enum PartitionDegreeUpdate {
1169    /// Parallel workers accumulated directional movement counts in the first
1170    /// reusable movement workspace.
1171    Moves,
1172    /// Small partitions use the former unstable selection order. Keep its
1173    /// reusable rank map long enough to update only records that crossed the
1174    /// cut.
1175    Ranked,
1176    /// A large partition with only one affordable degree lane still uses the
1177    /// bounded-memory radix selector, then applies deltas serially.
1178    Threshold {
1179        threshold_key: u32,
1180        ties_left: usize,
1181    },
1182}
1183
1184struct PartitionOutcome {
1185    swap_count: usize,
1186    degree_update: PartitionDegreeUpdate,
1187}
1188
1189#[derive(Clone, Copy)]
1190struct PartitionChunk {
1191    start: usize,
1192    end: usize,
1193    strictly_lower: usize,
1194    equal: usize,
1195    ties_left: usize,
1196}
1197
1198#[inline]
1199fn select_left(key: u32, threshold_key: u32, equal_seen: &mut usize, ties_left: usize) -> bool {
1200    if key < threshold_key {
1201        true
1202    } else if key == threshold_key {
1203        let selected = *equal_seen < ties_left;
1204        *equal_seen += 1;
1205        selected
1206    } else {
1207        false
1208    }
1209}
1210
1211/// Exact, deterministic parallel partition by `(gain.total_cmp(), old_index)`.
1212///
1213/// Output is stable within each half. Parallel workers also accumulate
1214/// per-term degree deltas for moved entities, eliminating the former serial
1215/// postings update without rescanning every posting after each iteration.
1216/// Small partitions retain the old unstable selector in their preallocated
1217/// rank slice: it is faster than four radix scans, and its within-half
1218/// permutation supplies the graph algorithm's established symmetry breaking.
1219#[allow(clippy::too_many_arguments)]
1220fn partition_by_gain(
1221    docs: &[u32],
1222    gains: &[f32],
1223    mid: usize,
1224    fwd: &ForwardIndex,
1225    movement_workspaces: &mut [TermDegrees],
1226    output: &mut [u32],
1227    ranked_scratch: &mut [usize],
1228    allow_inner_parallelism: bool,
1229) -> PartitionOutcome {
1230    #[cfg(not(feature = "native"))]
1231    let _ = (fwd, &movement_workspaces);
1232
1233    #[cfg(feature = "native")]
1234    if allow_inner_parallelism
1235        && !movement_workspaces.is_empty()
1236        && docs.len() >= PARALLEL_BP_MIN_ENTITIES
1237    {
1238        let (threshold_key, strictly_lower) =
1239            select_gain_threshold(gains, mid, allow_inner_parallelism);
1240        let ties_left = mid - strictly_lower;
1241
1242        // `degrees` remains live while these deltas are built. Reserving one
1243        // lane for it keeps total vocabulary arrays within the admitted set.
1244        let chunk_count = movement_workspaces
1245            .len()
1246            .min(docs.len().div_ceil(PARALLEL_BP_MIN_ENTITIES).max(1));
1247        let chunk_len = docs.len().div_ceil(chunk_count);
1248        let mut chunks: Vec<PartitionChunk> = gains
1249            .par_chunks(chunk_len)
1250            .enumerate()
1251            .map(|(chunk_id, chunk)| {
1252                let mut lower = 0usize;
1253                let mut equal = 0usize;
1254                for &gain in chunk {
1255                    match gain_order_key(gain).cmp(&threshold_key) {
1256                        std::cmp::Ordering::Less => lower += 1,
1257                        std::cmp::Ordering::Equal => equal += 1,
1258                        std::cmp::Ordering::Greater => {}
1259                    }
1260                }
1261                let start = chunk_id * chunk_len;
1262                PartitionChunk {
1263                    start,
1264                    end: start + chunk.len(),
1265                    strictly_lower: lower,
1266                    equal,
1267                    ties_left: 0,
1268                }
1269            })
1270            .collect();
1271
1272        let mut remaining_ties = ties_left;
1273        for chunk in &mut chunks {
1274            chunk.ties_left = remaining_ties.min(chunk.equal);
1275            remaining_ties -= chunk.ties_left;
1276        }
1277        debug_assert_eq!(remaining_ties, 0);
1278
1279        let (mut left_rest, mut right_rest) = output.split_at_mut(mid);
1280        let mut jobs = Vec::with_capacity(chunks.len());
1281        for chunk in chunks {
1282            let left_len = chunk.strictly_lower + chunk.ties_left;
1283            let right_len = chunk.end - chunk.start - left_len;
1284            let (left_out, next_left) = left_rest.split_at_mut(left_len);
1285            let (right_out, next_right) = right_rest.split_at_mut(right_len);
1286            jobs.push((
1287                chunk.start,
1288                &docs[chunk.start..chunk.end],
1289                &gains[chunk.start..chunk.end],
1290                chunk.ties_left,
1291                left_out,
1292                right_out,
1293            ));
1294            left_rest = next_left;
1295            right_rest = next_right;
1296        }
1297        debug_assert!(left_rest.is_empty() && right_rest.is_empty());
1298
1299        let swap_count = movement_workspaces[..chunk_count]
1300            .par_iter_mut()
1301            .zip(jobs.into_par_iter())
1302            .map(
1303                |(moves, (start, docs, gains, ties_for_chunk, left_out, right_out))| {
1304                    moves.reset();
1305                    let mut equal_seen = 0usize;
1306                    let mut left_cursor = 0usize;
1307                    let mut right_cursor = 0usize;
1308                    let mut swaps = 0usize;
1309
1310                    for (offset, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1311                        let key = gain_order_key(gain);
1312                        let now_left =
1313                            select_left(key, threshold_key, &mut equal_seen, ties_for_chunk);
1314                        if now_left {
1315                            left_out[left_cursor] = doc;
1316                            left_cursor += 1;
1317                        } else {
1318                            right_out[right_cursor] = doc;
1319                            right_cursor += 1;
1320                        }
1321
1322                        let was_left = start + offset < mid;
1323                        if was_left != now_left {
1324                            swaps += 1;
1325                            // [right→left, left→right]
1326                            let direction = usize::from(was_left);
1327                            for &term in fwd.doc_terms(doc as usize) {
1328                                moves.entry_mut(term as usize)[direction] += 1;
1329                            }
1330                        }
1331                    }
1332                    debug_assert_eq!(left_cursor, left_out.len());
1333                    debug_assert_eq!(right_cursor, right_out.len());
1334                    swaps
1335                },
1336            )
1337            .sum();
1338
1339        let (moves, partials) = movement_workspaces[..chunk_count]
1340            .split_first_mut()
1341            .expect("parallel BP partition must use at least one movement workspace");
1342        for partial in partials {
1343            moves.merge_from(partial);
1344        }
1345
1346        return PartitionOutcome {
1347            swap_count,
1348            degree_update: PartitionDegreeUpdate::Moves,
1349        };
1350    }
1351
1352    if docs.len() < PARALLEL_BP_MIN_ENTITIES {
1353        debug_assert_eq!(ranked_scratch.len(), docs.len());
1354        for (index, rank) in ranked_scratch.iter_mut().enumerate() {
1355            *rank = index;
1356        }
1357        ranked_scratch.select_nth_unstable_by(mid, |&left, &right| {
1358            gains[left]
1359                .total_cmp(&gains[right])
1360                .then_with(|| left.cmp(&right))
1361        });
1362
1363        let mut swaps = 0usize;
1364        for (rank, &old_index) in ranked_scratch.iter().enumerate() {
1365            output[rank] = docs[old_index];
1366            swaps += usize::from((old_index < mid) != (rank < mid));
1367        }
1368        return PartitionOutcome {
1369            swap_count: swaps,
1370            degree_update: PartitionDegreeUpdate::Ranked,
1371        };
1372    }
1373
1374    let (threshold_key, strictly_lower) =
1375        select_gain_threshold(gains, mid, allow_inner_parallelism);
1376    let ties_left = mid - strictly_lower;
1377    let mut equal_seen = 0usize;
1378    let mut left_cursor = 0usize;
1379    let mut right_cursor = mid;
1380    let mut swaps = 0usize;
1381    for (idx, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1382        let key = gain_order_key(gain);
1383        let now_left = select_left(key, threshold_key, &mut equal_seen, ties_left);
1384        if now_left {
1385            output[left_cursor] = doc;
1386            left_cursor += 1;
1387        } else {
1388            output[right_cursor] = doc;
1389            right_cursor += 1;
1390        }
1391        swaps += usize::from((idx < mid) != now_left);
1392    }
1393    debug_assert_eq!(left_cursor, mid);
1394    debug_assert_eq!(right_cursor, docs.len());
1395
1396    PartitionOutcome {
1397        swap_count: swaps,
1398        degree_update: PartitionDegreeUpdate::Threshold {
1399            threshold_key,
1400            ties_left,
1401        },
1402    }
1403}
1404
1405/// Apply degree changes for a bounded-memory serial radix partition.
1406fn update_degrees_for_threshold_partition(
1407    docs: &[u32],
1408    gains: &[f32],
1409    mid: usize,
1410    threshold_key: u32,
1411    ties_left: usize,
1412    fwd: &ForwardIndex,
1413    degrees: &mut TermDegrees,
1414) {
1415    let mut equal_seen = 0usize;
1416    for (idx, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1417        let key = gain_order_key(gain);
1418        let now_left = select_left(key, threshold_key, &mut equal_seen, ties_left);
1419        let was_left = idx < mid;
1420        if was_left == now_left {
1421            continue;
1422        }
1423        let left_delta = if was_left { -1i64 } else { 1i64 };
1424        for &term in fwd.doc_terms(doc as usize) {
1425            let degree = degrees.entry_mut(term as usize);
1426            let new_left = degree[0] as i64 + left_delta;
1427            let new_right = degree[1] as i64 - left_delta;
1428            debug_assert!(new_left >= 0 && new_right >= 0);
1429            degree[0] = new_left as u32;
1430            degree[1] = new_right as u32;
1431        }
1432    }
1433}
1434
1435/// Apply degree changes using the exact unstable rank order selected for a
1436/// small partition.
1437fn update_degrees_for_ranked_partition(
1438    docs: &[u32],
1439    ranked: &[usize],
1440    mid: usize,
1441    fwd: &ForwardIndex,
1442    degrees: &mut TermDegrees,
1443) {
1444    for (rank, &old_index) in ranked.iter().enumerate() {
1445        let was_left = old_index < mid;
1446        let now_left = rank < mid;
1447        if was_left == now_left {
1448            continue;
1449        }
1450        let left_delta = if was_left { -1i64 } else { 1i64 };
1451        for &term in fwd.doc_terms(docs[old_index] as usize) {
1452            let degree = degrees.entry_mut(term as usize);
1453            let new_left = degree[0] as i64 + left_delta;
1454            let new_right = degree[1] as i64 - left_delta;
1455            debug_assert!(new_left >= 0 && new_right >= 0);
1456            degree[0] = new_left as u32;
1457            degree[1] = new_right as u32;
1458        }
1459    }
1460}
1461
1462#[derive(Clone, Copy)]
1463pub(crate) struct BpProgressLabel<'a> {
1464    pub index: &'a str,
1465    pub field: &'a str,
1466    pub entity_kind: &'static str,
1467}
1468
1469#[cfg(test)]
1470impl BpProgressLabel<'static> {
1471    fn anonymous() -> Self {
1472        Self {
1473            index: "unknown",
1474            field: "unknown",
1475            entity_kind: "entities",
1476        }
1477    }
1478}
1479
1480#[cfg(feature = "native")]
1481struct BpProgress<'a> {
1482    label: BpProgressLabel<'a>,
1483    start: std::time::Instant,
1484    total_entities: usize,
1485    total_postings: u64,
1486    expected_depth: usize,
1487    next_log_ms: std::sync::atomic::AtomicU64,
1488    active_partitions: std::sync::atomic::AtomicU64,
1489    partitions_started: std::sync::atomic::AtomicU64,
1490    partitions_completed: std::sync::atomic::AtomicU64,
1491    iterations: std::sync::atomic::AtomicU64,
1492    entity_passes: std::sync::atomic::AtomicU64,
1493    swaps: std::sync::atomic::AtomicU64,
1494    deepest_level: std::sync::atomic::AtomicU64,
1495    objective_stops: std::sync::atomic::AtomicU64,
1496    last_objective_delta_bits: std::sync::atomic::AtomicU64,
1497    last_relative_delta_bits: std::sync::atomic::AtomicU64,
1498    active_metric_released: std::sync::atomic::AtomicBool,
1499}
1500
1501#[cfg(feature = "native")]
1502impl<'a> BpProgress<'a> {
1503    fn new(
1504        label: BpProgressLabel<'a>,
1505        total_entities: usize,
1506        total_postings: u64,
1507        expected_depth: usize,
1508        degree_lanes: usize,
1509    ) -> Self {
1510        log::info!(
1511            "[reorder][bp] started: index={} field={} entity_kind={} scheduler=level_synchronized degree_lanes={} entities={} postings={} expected_depth={} objective_stall_threshold={:.1e}x{} min_objective_iterations={}",
1512            label.index,
1513            label.field,
1514            label.entity_kind,
1515            degree_lanes,
1516            total_entities,
1517            total_postings,
1518            expected_depth,
1519            MIN_RELATIVE_OBJECTIVE_IMPROVEMENT,
1520            OBJECTIVE_STALL_ITERATIONS,
1521            MIN_OBJECTIVE_ITERATIONS,
1522        );
1523        crate::observe::reorder_bp_started(label.index, label.field, label.entity_kind);
1524        Self {
1525            label,
1526            start: std::time::Instant::now(),
1527            total_entities,
1528            total_postings,
1529            expected_depth,
1530            next_log_ms: std::sync::atomic::AtomicU64::new(30_000),
1531            active_partitions: std::sync::atomic::AtomicU64::new(0),
1532            partitions_started: std::sync::atomic::AtomicU64::new(0),
1533            partitions_completed: std::sync::atomic::AtomicU64::new(0),
1534            iterations: std::sync::atomic::AtomicU64::new(0),
1535            entity_passes: std::sync::atomic::AtomicU64::new(0),
1536            swaps: std::sync::atomic::AtomicU64::new(0),
1537            deepest_level: std::sync::atomic::AtomicU64::new(0),
1538            objective_stops: std::sync::atomic::AtomicU64::new(0),
1539            last_objective_delta_bits: std::sync::atomic::AtomicU64::new(0f64.to_bits()),
1540            last_relative_delta_bits: std::sync::atomic::AtomicU64::new(0f64.to_bits()),
1541            active_metric_released: std::sync::atomic::AtomicBool::new(false),
1542        }
1543    }
1544
1545    fn partition_started(&self, level: usize) {
1546        self.active_partitions
1547            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1548        self.partitions_started
1549            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1550        self.deepest_level
1551            .fetch_max(level as u64, std::sync::atomic::Ordering::Relaxed);
1552    }
1553
1554    fn partition_finished(&self) {
1555        self.partitions_completed
1556            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1557        self.active_partitions
1558            .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
1559    }
1560
1561    fn iteration(&self, entities: usize, swaps: usize, objective_delta: f64, relative_delta: f64) {
1562        self.iterations
1563            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1564        self.entity_passes
1565            .fetch_add(entities as u64, std::sync::atomic::Ordering::Relaxed);
1566        self.swaps
1567            .fetch_add(swaps as u64, std::sync::atomic::Ordering::Relaxed);
1568        self.last_objective_delta_bits.store(
1569            objective_delta.to_bits(),
1570            std::sync::atomic::Ordering::Relaxed,
1571        );
1572        self.last_relative_delta_bits.store(
1573            relative_delta.to_bits(),
1574            std::sync::atomic::Ordering::Relaxed,
1575        );
1576        self.maybe_log();
1577    }
1578
1579    fn objective_stop(&self) {
1580        self.objective_stops
1581            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1582    }
1583
1584    fn maybe_log(&self) {
1585        let elapsed_ms = self.start.elapsed().as_millis().min(u64::MAX as u128) as u64;
1586        let next = self.next_log_ms.load(std::sync::atomic::Ordering::Relaxed);
1587        if elapsed_ms < next
1588            || self
1589                .next_log_ms
1590                .compare_exchange(
1591                    next,
1592                    elapsed_ms.saturating_add(30_000),
1593                    std::sync::atomic::Ordering::Relaxed,
1594                    std::sync::atomic::Ordering::Relaxed,
1595                )
1596                .is_err()
1597        {
1598            return;
1599        }
1600
1601        let active = self
1602            .active_partitions
1603            .load(std::sync::atomic::Ordering::Relaxed);
1604        let started = self
1605            .partitions_started
1606            .load(std::sync::atomic::Ordering::Relaxed);
1607        let completed = self
1608            .partitions_completed
1609            .load(std::sync::atomic::Ordering::Relaxed);
1610        let iterations = self.iterations.load(std::sync::atomic::Ordering::Relaxed);
1611        let entity_passes = self
1612            .entity_passes
1613            .load(std::sync::atomic::Ordering::Relaxed);
1614        let swaps = self.swaps.load(std::sync::atomic::Ordering::Relaxed);
1615        let deepest = self
1616            .deepest_level
1617            .load(std::sync::atomic::Ordering::Relaxed);
1618        let objective_delta = f64::from_bits(
1619            self.last_objective_delta_bits
1620                .load(std::sync::atomic::Ordering::Relaxed),
1621        );
1622        let relative_delta = f64::from_bits(
1623            self.last_relative_delta_bits
1624                .load(std::sync::atomic::Ordering::Relaxed),
1625        );
1626        log::info!(
1627            "[reorder][bp] progress: index={} field={} entity_kind={} elapsed={:.1}s depth={}/{} partitions={}/{} active={} iterations={} entity_passes={} swaps={} last_objective_delta={:.3} relative={:.3e}",
1628            self.label.index,
1629            self.label.field,
1630            self.label.entity_kind,
1631            self.start.elapsed().as_secs_f64(),
1632            deepest,
1633            self.expected_depth,
1634            completed,
1635            started,
1636            active,
1637            iterations,
1638            entity_passes,
1639            swaps,
1640            objective_delta,
1641            relative_delta,
1642        );
1643    }
1644
1645    fn finish(
1646        &self,
1647        converged: bool,
1648        memory_limited: bool,
1649        deadline_exhausted: bool,
1650        cancelled: bool,
1651    ) {
1652        let elapsed = self.start.elapsed().as_secs_f64();
1653        let partitions = self
1654            .partitions_completed
1655            .load(std::sync::atomic::Ordering::Relaxed);
1656        let iterations = self.iterations.load(std::sync::atomic::Ordering::Relaxed);
1657        let entity_passes = self
1658            .entity_passes
1659            .load(std::sync::atomic::Ordering::Relaxed);
1660        let swaps = self.swaps.load(std::sync::atomic::Ordering::Relaxed);
1661        let deepest = self
1662            .deepest_level
1663            .load(std::sync::atomic::Ordering::Relaxed);
1664        let objective_stops = self
1665            .objective_stops
1666            .load(std::sync::atomic::Ordering::Relaxed);
1667        let stop_reason = if cancelled {
1668            "shutdown"
1669        } else if memory_limited {
1670            "memory_budget"
1671        } else if deadline_exhausted {
1672            "time_budget"
1673        } else if objective_stops > 0 {
1674            "objective"
1675        } else {
1676            "complete"
1677        };
1678        log::info!(
1679            "[reorder][bp] completed: index={} field={} entity_kind={} entities={} postings={} elapsed={:.1}s depth={}/{} partitions={} iterations={} entity_passes={} swaps={} objective_stops={} converged={} stop_reason={}",
1680            self.label.index,
1681            self.label.field,
1682            self.label.entity_kind,
1683            self.total_entities,
1684            self.total_postings,
1685            elapsed,
1686            deepest,
1687            self.expected_depth,
1688            partitions,
1689            iterations,
1690            entity_passes,
1691            swaps,
1692            objective_stops,
1693            converged,
1694            stop_reason,
1695        );
1696        crate::observe::reorder_bp_pass(
1697            self.label.index,
1698            self.label.field,
1699            self.label.entity_kind,
1700            stop_reason,
1701            elapsed,
1702            self.total_entities,
1703            self.total_postings,
1704            partitions,
1705            iterations,
1706            entity_passes,
1707            swaps,
1708            converged,
1709        );
1710        self.release_active_metric();
1711    }
1712
1713    fn release_active_metric(&self) {
1714        if !self
1715            .active_metric_released
1716            .swap(true, std::sync::atomic::Ordering::AcqRel)
1717        {
1718            crate::observe::reorder_bp_finished(
1719                self.label.index,
1720                self.label.field,
1721                self.label.entity_kind,
1722            );
1723        }
1724    }
1725}
1726
1727#[cfg(feature = "native")]
1728impl Drop for BpProgress<'_> {
1729    fn drop(&mut self) {
1730        self.release_active_metric();
1731    }
1732}
1733
1734#[cfg(not(feature = "native"))]
1735struct BpProgress<'a>(std::marker::PhantomData<&'a ()>);
1736
1737#[cfg(not(feature = "native"))]
1738impl BpProgress<'_> {
1739    fn new(_: BpProgressLabel<'_>, _: usize, _: u64, _: usize, _: usize) -> Self {
1740        Self(std::marker::PhantomData)
1741    }
1742    fn partition_started(&self, _: usize) {}
1743    fn partition_finished(&self) {}
1744    fn iteration(&self, _: usize, _: usize, _: f64, _: f64) {}
1745    fn objective_stop(&self) {}
1746    fn finish(&self, _: bool, _: bool, _: bool, _: bool) {}
1747}
1748
1749/// Level-synchronized graph bisection. Returns `(perm, converged)` where
1750/// `perm[new_pos] = old_index`. Convergence is false when the wall-clock or
1751/// memory budget prevents the requested work from finishing; a configured
1752/// depth cap is a chosen target and reports converged.
1753///
1754/// `min_partition_size` should be the configured BMP block size.
1755/// `max_iters` controls convergence (20 is standard).
1756///
1757/// Term IDs in the forward index must be compact (0..num_terms) so we can
1758/// use flat arrays for O(1) degree lookups instead of hash maps.
1759#[cfg(test)]
1760pub(crate) fn graph_bisection(
1761    fwd: &ForwardIndex,
1762    min_partition_size: usize,
1763    max_iters: usize,
1764    budget: BpBudget,
1765) -> (Vec<u32>, bool) {
1766    graph_bisection_with_progress(
1767        fwd,
1768        min_partition_size,
1769        max_iters,
1770        budget,
1771        None,
1772        BpProgressLabel::anonymous(),
1773    )
1774}
1775
1776pub(crate) fn graph_bisection_with_progress(
1777    fwd: &ForwardIndex,
1778    min_partition_size: usize,
1779    max_iters: usize,
1780    budget: BpBudget,
1781    cancellation: Option<&std::sync::atomic::AtomicBool>,
1782    progress_label: BpProgressLabel<'_>,
1783) -> (Vec<u32>, bool) {
1784    let n = fwd.num_docs();
1785    if n == 0 {
1786        return (Vec::new(), !fwd.budget_limited);
1787    }
1788
1789    let effective_min_partition = budget
1790        .min_partition_docs
1791        .unwrap_or(0)
1792        .max(min_partition_size)
1793        // A singleton cannot be bisected. The public callers use a positive
1794        // BMP block size, but enforcing the structural minimum here also
1795        // prevents an accidental zero configuration from walking empty tree
1796        // levels until the partition counter overflows.
1797        .max(1);
1798
1799    let mut docs: Vec<u32> = (0..n as u32).collect();
1800    let depth = if effective_min_partition > 0 {
1801        ((n as f64) / (effective_min_partition as f64))
1802            .log2()
1803            .ceil() as usize
1804    } else {
1805        0
1806    };
1807    #[cfg(feature = "native")]
1808    let degree_lanes = fwd.parallel_bisect_lanes.max(1);
1809    #[cfg(not(feature = "native"))]
1810    let degree_lanes = 1usize;
1811    let log_table = build_log_table(4096);
1812    let progress = BpProgress::new(progress_label, n, fwd.total_postings(), depth, degree_lanes);
1813
1814    log::debug!(
1815        "BP graph_bisection: n={}, min_partition={}, max_iters={}, depth=~{}, time_budget={:?}",
1816        n,
1817        effective_min_partition,
1818        max_iters,
1819        depth,
1820        budget.time_budget,
1821    );
1822
1823    #[cfg(feature = "native")]
1824    let deadline = budget.time_budget.map(|duration| {
1825        let now = std::time::Instant::now();
1826        now.checked_add(duration).unwrap_or(now)
1827    });
1828    #[cfg(not(feature = "native"))]
1829    let deadline: Option<()> = None;
1830
1831    let exhausted = std::sync::atomic::AtomicBool::new(false);
1832    let context = BisectContext {
1833        fwd,
1834        min_partition_size: effective_min_partition,
1835        max_iters,
1836        log_table: &log_table,
1837        #[cfg(feature = "native")]
1838        deadline,
1839        #[cfg(not(feature = "native"))]
1840        deadline,
1841        exhausted: &exhausted,
1842        cancellation,
1843        progress: &progress,
1844    };
1845    let immediately_exhausted = {
1846        #[cfg(feature = "native")]
1847        {
1848            deadline.is_some_and(|deadline| std::time::Instant::now() >= deadline)
1849        }
1850        #[cfg(not(feature = "native"))]
1851        {
1852            false
1853        }
1854    };
1855    let initially_cancelled =
1856        cancellation.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire));
1857    if immediately_exhausted || initially_cancelled {
1858        exhausted.store(true, std::sync::atomic::Ordering::Relaxed);
1859    } else if n > effective_min_partition {
1860        // Entity scratch and vocabulary-sized degree lanes are allocated
1861        // exactly once. At each depth, workers dynamically claim disjoint
1862        // partitions and reuse their lane for the whole level.
1863        let mut gains = vec![0.0f32; n];
1864        let mut partitioned = vec![0u32; n];
1865        let mut ranked = vec![0usize; n];
1866        let mut degree_workspaces: Vec<TermDegrees> = (0..degree_lanes)
1867            .map(|_| TermDegrees::new(fwd.num_terms))
1868            .collect();
1869        bisect_level_synchronized(
1870            &mut docs,
1871            &mut gains,
1872            &mut partitioned,
1873            &mut ranked,
1874            &mut degree_workspaces,
1875            &context,
1876        );
1877    }
1878
1879    let stopped_early = exhausted.load(std::sync::atomic::Ordering::Relaxed);
1880    let cancelled =
1881        cancellation.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire));
1882    let deadline_exhausted = stopped_early && !cancelled;
1883    let converged = !fwd.budget_limited && !stopped_early && !cancelled;
1884    progress.finish(converged, fwd.budget_limited, deadline_exhausted, cancelled);
1885    if !converged {
1886        log::info!(
1887            "BP graph_bisection: pass incomplete at n={} (time={:?}, memory_limited={}, cancelled={}) — emitting partial (still valid) permutation",
1888            n,
1889            budget.time_budget,
1890            fwd.budget_limited,
1891            cancelled,
1892        );
1893    }
1894    (docs, converged)
1895}
1896
1897/// Context shared by all partitions in a level-synchronized BP pass.
1898///
1899/// Uses flat `Vec<u32>` degree arrays indexed by compact term_id for cache-friendly
1900/// O(1) lookups (vs FxHashMap which has poor cache locality at scale).
1901///
1902struct BisectContext<'a> {
1903    fwd: &'a ForwardIndex,
1904    min_partition_size: usize,
1905    max_iters: usize,
1906    log_table: &'a [f32],
1907    #[cfg(feature = "native")]
1908    deadline: Option<std::time::Instant>,
1909    #[cfg(not(feature = "native"))]
1910    deadline: Option<()>,
1911    exhausted: &'a std::sync::atomic::AtomicBool,
1912    cancellation: Option<&'a std::sync::atomic::AtomicBool>,
1913    progress: &'a BpProgress<'a>,
1914}
1915
1916/// Return the range of partition `partition_id` at a fixed recursion `level`.
1917///
1918/// Replaying the path bits produces exactly the same floor-left/ceil-right
1919/// boundaries as recursive `split_at_mut(len / 2)`, including for non-power
1920/// of-two collection sizes.
1921fn partition_range(total: usize, level: usize, partition_id: usize) -> std::ops::Range<usize> {
1922    debug_assert!(level < usize::BITS as usize);
1923    debug_assert!(partition_id < (1usize << level));
1924    let mut start = 0usize;
1925    let mut len = total;
1926    for bit in (0..level).rev() {
1927        let left_len = len / 2;
1928        if partition_id & (1usize << bit) == 0 {
1929            len = left_len;
1930        } else {
1931            start += left_len;
1932            len -= left_len;
1933        }
1934    }
1935    start..start + len
1936}
1937
1938/// Collect all active partitions when they fit in `limit` lanes. Returning
1939/// `None` means there are more active partitions than lanes, so the caller
1940/// should use dynamic claiming instead of assigning multiple lanes per node.
1941fn small_active_partition_set(
1942    total: usize,
1943    level: usize,
1944    min_partition_size: usize,
1945    limit: usize,
1946) -> Option<Vec<usize>> {
1947    let partition_count = 1usize.checked_shl(level as u32)?;
1948    let mut active = Vec::with_capacity(partition_count.min(limit));
1949    for partition_id in 0..partition_count {
1950        if partition_range(total, level, partition_id).len() <= min_partition_size {
1951            continue;
1952        }
1953        active.push(partition_id);
1954        if active.len() > limit {
1955            return None;
1956        }
1957    }
1958    Some(active)
1959}
1960
1961/// Mutable pass buffers shared by native level workers.
1962///
1963/// The scheduler hands each claimed partition ID to exactly one worker, and
1964/// `partition_range` yields non-overlapping ranges at a fixed level. The level
1965/// barrier completes before any buffer is accessed again or the next level is
1966/// started. Those invariants make concurrent slice construction sound.
1967#[cfg(feature = "native")]
1968#[derive(Clone, Copy)]
1969struct LevelBuffers {
1970    docs: *mut u32,
1971    gains: *mut f32,
1972    partitioned: *mut u32,
1973    ranked: *mut usize,
1974    len: usize,
1975}
1976
1977#[cfg(feature = "native")]
1978unsafe impl Send for LevelBuffers {}
1979#[cfg(feature = "native")]
1980unsafe impl Sync for LevelBuffers {}
1981
1982#[cfg(feature = "native")]
1983impl LevelBuffers {
1984    fn new(
1985        docs: &mut [u32],
1986        gains: &mut [f32],
1987        partitioned: &mut [u32],
1988        ranked: &mut [usize],
1989    ) -> Self {
1990        debug_assert_eq!(docs.len(), gains.len());
1991        debug_assert_eq!(docs.len(), partitioned.len());
1992        debug_assert_eq!(docs.len(), ranked.len());
1993        Self {
1994            docs: docs.as_mut_ptr(),
1995            gains: gains.as_mut_ptr(),
1996            partitioned: partitioned.as_mut_ptr(),
1997            ranked: ranked.as_mut_ptr(),
1998            len: docs.len(),
1999        }
2000    }
2001
2002    /// # Safety
2003    ///
2004    /// During one level, every requested range must be in bounds and disjoint
2005    /// from every range concurrently handed out from this `LevelBuffers`.
2006    unsafe fn with_slices<R>(
2007        &self,
2008        range: std::ops::Range<usize>,
2009        use_slices: impl for<'a> FnOnce(
2010            &'a mut [u32],
2011            &'a mut [f32],
2012            &'a mut [u32],
2013            &'a mut [usize],
2014        ) -> R,
2015    ) -> R {
2016        debug_assert!(range.start <= range.end && range.end <= self.len);
2017        let len = range.len();
2018        // SAFETY: upheld by the caller as documented above. All four backing
2019        // allocations remain live and immovable until the level barrier.
2020        unsafe {
2021            use_slices(
2022                std::slice::from_raw_parts_mut(self.docs.add(range.start), len),
2023                std::slice::from_raw_parts_mut(self.gains.add(range.start), len),
2024                std::slice::from_raw_parts_mut(self.partitioned.add(range.start), len),
2025                std::slice::from_raw_parts_mut(self.ranked.add(range.start), len),
2026            )
2027        }
2028    }
2029}
2030
2031/// Process the recursion tree breadth-first. Same-depth partitions have
2032/// similar sizes and dynamically claim bounded degree lanes, avoiding the
2033/// inter-level contention and permanently imbalanced descendant ownership of
2034/// recursive fork-join BP.
2035fn bisect_level_synchronized(
2036    docs: &mut [u32],
2037    gains: &mut [f32],
2038    partitioned: &mut [u32],
2039    ranked_scratch: &mut [usize],
2040    degree_workspaces: &mut [TermDegrees],
2041    context: &BisectContext<'_>,
2042) {
2043    debug_assert_eq!(docs.len(), gains.len());
2044    debug_assert_eq!(docs.len(), partitioned.len());
2045    debug_assert_eq!(docs.len(), ranked_scratch.len());
2046    debug_assert!(!degree_workspaces.is_empty());
2047
2048    #[cfg(feature = "native")]
2049    {
2050        let buffers = LevelBuffers::new(docs, gains, partitioned, ranked_scratch);
2051        let mut level = 0usize;
2052        while let Some(partition_count) = 1usize.checked_shl(level as u32) {
2053            let small_set = small_active_partition_set(
2054                docs.len(),
2055                level,
2056                context.min_partition_size,
2057                degree_workspaces.len(),
2058            );
2059
2060            match small_set {
2061                Some(active) if active.is_empty() => break,
2062                Some(active) => {
2063                    // With fewer partitions than degree lanes, give each node
2064                    // a lane group so its frequency/gain work can still use
2065                    // the whole pool during BP's startup levels.
2066                    let active_count = active.len();
2067                    let allow_inner_parallelism =
2068                        active_count < rayon::current_num_threads().max(1);
2069                    let mut rest: &mut [TermDegrees] = &mut *degree_workspaces;
2070                    let mut groups = Vec::with_capacity(active_count);
2071                    for remaining_groups in (1..=active_count).rev() {
2072                        let group_len = rest.len().div_ceil(remaining_groups);
2073                        let (group, tail) = rest.split_at_mut(group_len);
2074                        groups.push(group);
2075                        rest = tail;
2076                    }
2077                    debug_assert!(rest.is_empty());
2078                    groups.into_par_iter().zip(active.into_par_iter()).for_each(
2079                        |(workspaces, partition_id)| {
2080                            let range = partition_range(docs.len(), level, partition_id);
2081                            // SAFETY: active IDs are unique at one fixed level,
2082                            // hence their ranges are disjoint. `for_each` is
2083                            // the barrier before buffers are reused.
2084                            unsafe {
2085                                buffers.with_slices(
2086                                    range,
2087                                    |part_docs, part_gains, part_output, part_ranked| {
2088                                        bisect_partition(
2089                                            part_docs,
2090                                            part_gains,
2091                                            part_output,
2092                                            part_ranked,
2093                                            workspaces,
2094                                            level,
2095                                            allow_inner_parallelism,
2096                                            context,
2097                                        );
2098                                    },
2099                                );
2100                            }
2101                        },
2102                    );
2103                }
2104                None => {
2105                    // More partitions than degree lanes: one long-lived worker
2106                    // per admitted lane repeatedly claims the next partition.
2107                    // This preserves the memory bound and lets short/deep work
2108                    // steal naturally instead of pinning a whole subtree to a
2109                    // lane for the remainder of the pass.
2110                    let next = std::sync::atomic::AtomicUsize::new(0);
2111                    let allow_inner_parallelism =
2112                        degree_workspaces.len() < rayon::current_num_threads().max(1);
2113                    degree_workspaces.par_iter_mut().for_each(|workspace| {
2114                        loop {
2115                            let partition_id =
2116                                next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2117                            if partition_id >= partition_count {
2118                                break;
2119                            }
2120                            let range = partition_range(docs.len(), level, partition_id);
2121                            if range.len() <= context.min_partition_size {
2122                                continue;
2123                            }
2124                            // SAFETY: `fetch_add` assigns every partition ID
2125                            // once, and fixed-level partition ranges are
2126                            // pairwise disjoint.
2127                            unsafe {
2128                                buffers.with_slices(
2129                                    range,
2130                                    |part_docs, part_gains, part_output, part_ranked| {
2131                                        bisect_partition(
2132                                            part_docs,
2133                                            part_gains,
2134                                            part_output,
2135                                            part_ranked,
2136                                            std::slice::from_mut(workspace),
2137                                            level,
2138                                            allow_inner_parallelism,
2139                                            context,
2140                                        );
2141                                    },
2142                                );
2143                            }
2144                        }
2145                    });
2146                }
2147            }
2148
2149            if context.exhausted.load(std::sync::atomic::Ordering::Relaxed) {
2150                break;
2151            }
2152            level += 1;
2153        }
2154    }
2155
2156    #[cfg(not(feature = "native"))]
2157    {
2158        let mut level = 0usize;
2159        while let Some(partition_count) = 1usize.checked_shl(level as u32) {
2160            let mut active = false;
2161            for partition_id in 0..partition_count {
2162                let range = partition_range(docs.len(), level, partition_id);
2163                if range.len() <= context.min_partition_size {
2164                    continue;
2165                }
2166                active = true;
2167                bisect_partition(
2168                    &mut docs[range.clone()],
2169                    &mut gains[range.clone()],
2170                    &mut partitioned[range.clone()],
2171                    &mut ranked_scratch[range],
2172                    degree_workspaces,
2173                    level,
2174                    false,
2175                    context,
2176                );
2177            }
2178            if !active || context.exhausted.load(std::sync::atomic::Ordering::Relaxed) {
2179                break;
2180            }
2181            level += 1;
2182        }
2183    }
2184}
2185
2186/// Bisection of one document slice, without scheduling its children.
2187///
2188/// Uses flat `Vec<u32>` degree arrays indexed by compact term_id for
2189/// cache-friendly O(1) lookups. Adaptive iteration count reduces work at top
2190/// levels where coarse splits converge faster and dominate total runtime.
2191#[allow(clippy::too_many_arguments)]
2192fn bisect_partition(
2193    docs: &mut [u32],
2194    gains: &mut [f32],
2195    partitioned: &mut [u32],
2196    ranked_scratch: &mut [usize],
2197    degree_workspaces: &mut [TermDegrees],
2198    level: usize,
2199    allow_inner_parallelism: bool,
2200    context: &BisectContext<'_>,
2201) {
2202    let n = docs.len();
2203    debug_assert_eq!(gains.len(), n);
2204    debug_assert_eq!(partitioned.len(), n);
2205    debug_assert_eq!(ranked_scratch.len(), n);
2206    debug_assert!(!degree_workspaces.is_empty());
2207    if n <= context.min_partition_size {
2208        return;
2209    }
2210    // Anytime cutoff: leave this subtree in its current (valid) order.
2211    if context.exhausted.load(std::sync::atomic::Ordering::Relaxed)
2212        || context
2213            .cancellation
2214            .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2215    {
2216        context
2217            .exhausted
2218            .store(true, std::sync::atomic::Ordering::Relaxed);
2219        return;
2220    }
2221    #[cfg(feature = "native")]
2222    if let Some(dl) = context.deadline
2223        && std::time::Instant::now() >= dl
2224    {
2225        context
2226            .exhausted
2227            .store(true, std::sync::atomic::Ordering::Relaxed);
2228        return;
2229    }
2230    #[cfg(not(feature = "native"))]
2231    let _ = context.deadline;
2232
2233    context.progress.partition_started(level);
2234    let mid = n / 2;
2235
2236    // Adaptive iteration count: large partitions converge faster with
2237    // coarse splits, so fewer refinement passes suffice. The fine-grained
2238    // clustering is handled by deeper recursion levels with full iterations.
2239    let effective_iters = if n > 100_000 {
2240        context.max_iters.min(12)
2241    } else {
2242        context.max_iters
2243    };
2244
2245    // Compact term IDs permit direct indexing. Slots are initialized lazily so
2246    // deep partitions touch only their active terms. Coarse partitions use
2247    // multiple preallocated lanes; fine levels dynamically reuse one lane per
2248    // active partition.
2249    build_term_degrees(
2250        docs,
2251        mid,
2252        context.fwd,
2253        degree_workspaces,
2254        context.cancellation,
2255    );
2256    if context
2257        .cancellation
2258        .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2259    {
2260        context
2261            .exhausted
2262            .store(true, std::sync::atomic::Ordering::Relaxed);
2263        context.progress.partition_finished();
2264        return;
2265    }
2266    // A full-vocabulary objective scan pays off only on coarse partitions,
2267    // where one avoided refinement skips millions of postings. At fine levels
2268    // it cost more than the work it could save, so retain the cheap gain
2269    // cooling condition there.
2270    let track_objective = n >= PARALLEL_BP_MIN_ENTITIES;
2271    if track_objective {
2272        // The old bitmap scan accumulated terms in ascending order. Sorting
2273        // once preserves that exact floating-point order while subsequent
2274        // objective evaluations visit only words active in this partition.
2275        degree_workspaces[0].sort_touched_words();
2276    }
2277    let mut previous_objective = if track_objective {
2278        degree_workspaces[0].bisection_objective(mid, n - mid, context.log_table)
2279    } else {
2280        0.0
2281    };
2282    let mut best_objective = previous_objective;
2283    let mut objective_stalls = 0usize;
2284
2285    for iter in 0..effective_iters {
2286        // Anytime cutoff between refinement passes: keep the current split.
2287        if context
2288            .cancellation
2289            .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2290        {
2291            context
2292                .exhausted
2293                .store(true, std::sync::atomic::Ordering::Relaxed);
2294            break;
2295        }
2296        #[cfg(feature = "native")]
2297        if let Some(dl) = context.deadline
2298            && std::time::Instant::now() >= dl
2299        {
2300            context
2301                .exhausted
2302                .store(true, std::sync::atomic::Ordering::Relaxed);
2303            break;
2304        }
2305        // Compute gain for each document (approx_1 from Dhulipala et al.)
2306        // Parallelized for large partitions where per-doc work dominates.
2307        compute_gains(
2308            docs,
2309            context.fwd,
2310            mid,
2311            &degree_workspaces[0],
2312            context.log_table,
2313            gains,
2314            allow_inner_parallelism,
2315        );
2316        if context
2317            .cancellation
2318            .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2319        {
2320            context
2321                .exhausted
2322                .store(true, std::sync::atomic::Ordering::Relaxed);
2323            break;
2324        }
2325
2326        // Exact median selection and stable partition. The old path allocated
2327        // `Vec<usize>` (8 bytes/entity) and selected/applied it serially; the
2328        // radix path is O(4n), parallel, and accumulates moved-term deltas in
2329        // the same bounded worker lanes used by degree construction.
2330        let partition = partition_by_gain(
2331            docs,
2332            gains,
2333            mid,
2334            context.fwd,
2335            &mut degree_workspaces[1..],
2336            partitioned,
2337            ranked_scratch,
2338            allow_inner_parallelism,
2339        );
2340
2341        if partition.swap_count == 0 {
2342            context.progress.iteration(n, 0, 0.0, 0.0);
2343            // Preserve the selector's within-half order for the next level.
2344            // Small partitions intentionally retain quickselect's established
2345            // symmetry-breaking permutation.
2346            docs.copy_from_slice(partitioned);
2347            break;
2348        }
2349
2350        match &partition.degree_update {
2351            PartitionDegreeUpdate::Moves => {
2352                let (degrees, movement_workspaces) = degree_workspaces
2353                    .split_first_mut()
2354                    .expect("BP always has one degree workspace");
2355                movement_workspaces
2356                    .first()
2357                    .expect("parallel BP movement update requires a spare workspace")
2358                    .apply_moves_to(degrees);
2359            }
2360            PartitionDegreeUpdate::Ranked => update_degrees_for_ranked_partition(
2361                docs,
2362                ranked_scratch,
2363                mid,
2364                context.fwd,
2365                &mut degree_workspaces[0],
2366            ),
2367            PartitionDegreeUpdate::Threshold {
2368                threshold_key,
2369                ties_left,
2370            } => update_degrees_for_threshold_partition(
2371                docs,
2372                gains,
2373                mid,
2374                *threshold_key,
2375                *ties_left,
2376                context.fwd,
2377                &mut degree_workspaces[0],
2378            ),
2379        }
2380
2381        let (new_objective, objective_improvement, relative_improvement) = if track_objective {
2382            let new_objective =
2383                degree_workspaces[0].bisection_objective(mid, n - mid, context.log_table);
2384            let objective_improvement = new_objective - previous_objective;
2385            let relative_improvement = objective_improvement / previous_objective.abs().max(1.0);
2386            (new_objective, objective_improvement, relative_improvement)
2387        } else {
2388            (0.0, 0.0, 0.0)
2389        };
2390        context.progress.iteration(
2391            n,
2392            partition.swap_count,
2393            objective_improvement,
2394            relative_improvement,
2395        );
2396
2397        // Keep the existing approximate-BP semantics: one median refinement
2398        // can temporarily reduce the exact objective before the reciprocal
2399        // move settles. Rejecting that first step made valid clustered inputs
2400        // no-op. Instead, accept refinements and stop only after a warm-up plus
2401        // consecutive iterations that fail to improve the best exact
2402        // objective by a meaningful relative amount.
2403        docs.copy_from_slice(partitioned);
2404        if track_objective {
2405            previous_objective = new_objective;
2406            let relative_best_improvement =
2407                (new_objective - best_objective) / best_objective.abs().max(1.0);
2408            if relative_best_improvement >= MIN_RELATIVE_OBJECTIVE_IMPROVEMENT {
2409                best_objective = new_objective;
2410                objective_stalls = 0;
2411            } else if iter + 1 >= MIN_OBJECTIVE_ITERATIONS {
2412                objective_stalls += 1;
2413            }
2414            if objective_stalls >= OBJECTIVE_STALL_ITERATIONS {
2415                context.progress.objective_stop();
2416                break;
2417            }
2418        }
2419
2420        // Early termination: if < 0.5% of docs swapped, partition is stable
2421        if iter > 2 && partition.swap_count < n / 200 {
2422            break;
2423        }
2424
2425        // Fine partitions retain the previous cheap cooling rule. Computing
2426        // an exact vocabulary objective here costs more than the posting work
2427        // it can avoid.
2428        if !track_objective && iter > 5 {
2429            let max_abs_gain = gains
2430                .iter()
2431                .copied()
2432                .fold(0.0f32, |max_gain, gain| max_gain.max(gain.abs()));
2433            if max_abs_gain < 0.001 {
2434                break;
2435            }
2436        }
2437    }
2438
2439    context.progress.partition_finished();
2440}
2441
2442/// Compute gains for all documents, parallelized via rayon for large partitions.
2443///
2444/// Each doc's gain is independent: iterate its terms, accumulate the log-gap
2445/// cost delta of moving it to the other side. Read-only access to degree arrays
2446/// makes this embarrassingly parallel.
2447#[inline(never)]
2448fn compute_gains(
2449    docs: &[u32],
2450    fwd: &ForwardIndex,
2451    mid: usize,
2452    degrees: &TermDegrees,
2453    log_table: &[f32],
2454    gains: &mut [f32],
2455    allow_inner_parallelism: bool,
2456) {
2457    #[cfg(not(feature = "native"))]
2458    let _ = allow_inner_parallelism;
2459
2460    // Single coherent key: HIGH = belongs in the RIGHT half.
2461    // Left docs get +approx_one(from=left, to=right) — a misplaced left doc
2462    // (terms concentrated right) scores high. Right docs get
2463    // -approx_one(from=right, to=left) — a misplaced right doc scores low.
2464    // This matches the reference two-sided formulation (compute_gains_left /
2465    // compute_gains_right with negation); ranking both halves by raw
2466    // "move gain" instead made both sides' misplaced docs rank identically,
2467    // so the partition step could never exchange them.
2468    let gain_for_doc = |i: usize| -> f32 {
2469        let doc = docs[i] as usize;
2470        let in_left = i < mid;
2471        let mut g = 0.0f32;
2472        for &term in fwd.doc_terms(doc) {
2473            let [left, right] = degrees.get(term as usize);
2474            let (from, to) = if in_left {
2475                (left, right)
2476            } else {
2477                (right, left)
2478            };
2479            let move_gain = fast_log2_lookup(to as usize + 2, log_table)
2480                - fast_log2_lookup(from as usize, log_table)
2481                - std::f32::consts::LOG2_E / (1.0 + to as f32);
2482            g += if in_left { move_gain } else { -move_gain };
2483        }
2484        g
2485    };
2486
2487    #[cfg(feature = "native")]
2488    {
2489        if allow_inner_parallelism && docs.len() > 4096 {
2490            gains
2491                .par_iter_mut()
2492                .enumerate()
2493                .for_each(|(i, gain)| *gain = gain_for_doc(i));
2494        } else {
2495            for (i, gain) in gains.iter_mut().enumerate().take(docs.len()) {
2496                *gain = gain_for_doc(i);
2497            }
2498        }
2499    }
2500    #[cfg(not(feature = "native"))]
2501    {
2502        for (i, gain) in gains.iter_mut().enumerate().take(docs.len()) {
2503            *gain = gain_for_doc(i);
2504        }
2505    }
2506}
2507
2508// ── Helpers ──────────────────────────────────────────────────────────────
2509
2510/// Build precomputed log2 table for values 0..size.
2511fn build_log_table(size: usize) -> Vec<f32> {
2512    let mut table = vec![0.0f32; size];
2513    // log2(0) is undefined; use a large negative value
2514    table[0] = -10.0;
2515    for (i, entry) in table.iter_mut().enumerate().skip(1) {
2516        *entry = (i as f32).log2();
2517    }
2518    table
2519}
2520
2521/// Fast log2 with precomputed table lookup.
2522#[inline]
2523fn fast_log2_lookup(val: usize, table: &[f32]) -> f32 {
2524    if val < table.len() {
2525        table[val]
2526    } else {
2527        (val as f32).log2()
2528    }
2529}
2530
2531#[cfg(test)]
2532mod tests {
2533    use super::*;
2534
2535    #[test]
2536    fn bounded_frequency_count_and_candidate_selection_are_exact() {
2537        let items: Vec<u32> = (0..10_000).collect();
2538        let frequencies = count_frequencies_bounded(&items, 7, 16 * 1024, |item, counts| {
2539            let term = *item as usize % counts.len();
2540            counts[term] += 1;
2541        })
2542        .unwrap();
2543        assert_eq!(frequencies.iter().sum::<u32>(), items.len() as u32);
2544        for (term, &frequency) in frequencies.iter().enumerate() {
2545            let expected = (term..items.len()).step_by(frequencies.len()).count() as u32;
2546            assert_eq!(frequency, expected);
2547        }
2548
2549        let (selected, limited) =
2550            select_frequency_candidates(&[8, 3, 0, 5, 3], 1, 10, 2 * CANDIDATE_ENTRY_BYTES);
2551        assert!(limited);
2552        let mut selected = selected;
2553        selected.sort_unstable();
2554        assert_eq!(selected, vec![(1, 3), (4, 3)]);
2555    }
2556
2557    #[test]
2558    fn bounded_frequency_count_rejects_an_undersized_budget() {
2559        assert!(
2560            count_frequencies_bounded(&[0u32], 32, 32, |_, _| {}).is_none(),
2561            "one complete dense frequency table must fit before counting"
2562        );
2563    }
2564
2565    #[test]
2566    fn lazy_term_degrees_initialize_only_on_first_write() {
2567        let mut degrees = TermDegrees::new(130);
2568        let values_ptr = degrees.values.as_ptr();
2569        let bitmap_ptr = degrees.initialized.as_ptr();
2570        assert_eq!(degrees.get(65), [0, 0]);
2571        degrees.entry_mut(65)[0] += 3;
2572        degrees.entry_mut(65)[1] += 2;
2573        assert_eq!(degrees.get(65), [3, 2]);
2574        assert_eq!(degrees.get(64), [0, 0]);
2575        assert_eq!(
2576            degrees
2577                .initialized
2578                .iter()
2579                .map(|w| w.count_ones())
2580                .sum::<u32>(),
2581            1
2582        );
2583
2584        degrees.reset();
2585        assert_eq!(degrees.values.as_ptr(), values_ptr);
2586        assert_eq!(degrees.initialized.as_ptr(), bitmap_ptr);
2587        assert_eq!(degrees.get(65), [0, 0]);
2588        assert!(degrees.touched_words.is_empty());
2589        degrees.entry_mut(129)[1] = 7;
2590        assert_eq!(degrees.get(129), [0, 7]);
2591        assert_eq!(degrees.get(65), [0, 0]);
2592    }
2593
2594    #[test]
2595    fn sparse_objective_keeps_the_original_ascending_term_order() {
2596        let mut degrees = TermDegrees::new(130);
2597        *degrees.entry_mut(129) = [5, 2];
2598        *degrees.entry_mut(1) = [3, 7];
2599        *degrees.entry_mut(65) = [11, 13];
2600        assert_eq!(degrees.touched_words, vec![2, 0, 1]);
2601        degrees.sort_touched_words();
2602        assert_eq!(degrees.touched_words, vec![0, 1, 2]);
2603
2604        let log_table = build_log_table(4096);
2605        let actual = degrees.bisection_objective(31, 32, &log_table);
2606        let side_log = [
2607            fast_log2_lookup(31, &log_table) as f64,
2608            fast_log2_lookup(32, &log_table) as f64,
2609        ];
2610        let mut original_bitmap_scan = 0.0f64;
2611        for (word_idx, &initialized) in degrees.initialized.iter().enumerate() {
2612            let mut pending = initialized;
2613            while pending != 0 {
2614                let bit = pending.trailing_zeros() as usize;
2615                let term = word_idx * 64 + bit;
2616                let [left, right] = degrees.get(term);
2617                for (side, count) in [left, right].into_iter().enumerate() {
2618                    if count > 0 {
2619                        original_bitmap_scan += count as f64
2620                            * (fast_log2_lookup(count as usize + 1, &log_table) as f64
2621                                - side_log[side]);
2622                    }
2623                }
2624                pending &= pending - 1;
2625            }
2626        }
2627        assert_eq!(actual.to_bits(), original_bitmap_scan.to_bits());
2628    }
2629
2630    #[test]
2631    fn gain_radix_key_matches_total_cmp() {
2632        let values = [
2633            f32::from_bits(0xffc0_0001),
2634            f32::NEG_INFINITY,
2635            -42.0,
2636            -0.0,
2637            0.0,
2638            42.0,
2639            f32::INFINITY,
2640            f32::from_bits(0x7fc0_0001),
2641        ];
2642        let mut by_cmp = values;
2643        by_cmp.sort_by(f32::total_cmp);
2644        let mut by_key = values;
2645        by_key.sort_by_key(|value| gain_order_key(*value));
2646        assert_eq!(
2647            by_cmp.map(f32::to_bits),
2648            by_key.map(f32::to_bits),
2649            "radix selection must preserve the former total_cmp order"
2650        );
2651    }
2652
2653    #[test]
2654    fn radix_threshold_matches_exact_rank_with_ties() {
2655        let gains = [3.0, -1.0, 7.0, -1.0, 0.0, -0.0, 3.0, 9.0, 3.0, 2.0, 2.0];
2656        let mut sorted: Vec<(u32, usize)> = gains
2657            .iter()
2658            .enumerate()
2659            .map(|(idx, &gain)| (gain_order_key(gain), idx))
2660            .collect();
2661        sorted.sort_unstable();
2662
2663        for left_count in 1..=gains.len() {
2664            let (threshold, lower) = select_gain_threshold(&gains, left_count, true);
2665            assert_eq!(threshold, sorted[left_count - 1].0);
2666            assert_eq!(lower, sorted.partition_point(|&(key, _)| key < threshold),);
2667        }
2668    }
2669
2670    #[cfg(feature = "native")]
2671    #[test]
2672    fn parallel_partition_matches_exact_selection_and_degree_rebuild() {
2673        const N: usize = PARALLEL_BP_MIN_ENTITIES + 1;
2674        const TERMS: usize = 101;
2675        let mut terms = Vec::with_capacity(N * 3);
2676        let mut offsets = Vec::with_capacity(N + 1);
2677        offsets.push(0);
2678        for doc in 0..N {
2679            terms.extend_from_slice(&[
2680                (doc % TERMS) as u32,
2681                ((doc / 7) % TERMS) as u32,
2682                ((doc * 13) % TERMS) as u32,
2683            ]);
2684            offsets.push(terms.len() as u64);
2685        }
2686        let fwd = ForwardIndex {
2687            terms,
2688            offsets,
2689            num_terms: TERMS,
2690            parallel_bisect_lanes: 4,
2691            budget_limited: false,
2692        };
2693        let docs: Vec<u32> = (0..N as u32)
2694            .map(|idx| ((idx as usize * 7_919) % N) as u32)
2695            .collect();
2696        let gains: Vec<f32> = docs
2697            .iter()
2698            .map(|&doc| ((doc as usize * 37) % 257) as f32 - 128.0)
2699            .collect();
2700        let mid = N / 2;
2701
2702        let mut ranked: Vec<usize> = (0..N).collect();
2703        ranked.sort_unstable_by(|&left, &right| {
2704            gains[left]
2705                .total_cmp(&gains[right])
2706                .then_with(|| left.cmp(&right))
2707        });
2708        let mut selected_left = vec![false; N];
2709        for &idx in &ranked[..mid] {
2710            selected_left[idx] = true;
2711        }
2712        let expected: Vec<u32> = docs
2713            .iter()
2714            .enumerate()
2715            .filter(|(idx, _)| selected_left[*idx])
2716            .chain(
2717                docs.iter()
2718                    .enumerate()
2719                    .filter(|(idx, _)| !selected_left[*idx]),
2720            )
2721            .map(|(_, &doc)| doc)
2722            .collect();
2723
2724        let mut output = vec![0; N];
2725        let mut ranked_scratch = Vec::new();
2726        let mut movement_workspaces: Vec<_> = (0..3).map(|_| TermDegrees::new(TERMS)).collect();
2727        let outcome = partition_by_gain(
2728            &docs,
2729            &gains,
2730            mid,
2731            &fwd,
2732            &mut movement_workspaces,
2733            &mut output,
2734            &mut ranked_scratch,
2735            true,
2736        );
2737        assert_eq!(output, expected);
2738        assert!(
2739            matches!(&outcome.degree_update, PartitionDegreeUpdate::Moves),
2740            "test must exercise parallel movement counts"
2741        );
2742
2743        let mut updated_workspaces: Vec<_> = (0..4).map(|_| TermDegrees::new(TERMS)).collect();
2744        build_term_degrees(&docs, mid, &fwd, &mut updated_workspaces, None);
2745        movement_workspaces[0].apply_moves_to(&mut updated_workspaces[0]);
2746
2747        let mut rebuilt_workspaces: Vec<_> = (0..4).map(|_| TermDegrees::new(TERMS)).collect();
2748        build_term_degrees(&output, mid, &fwd, &mut rebuilt_workspaces, None);
2749        for term in 0..TERMS {
2750            assert_eq!(
2751                updated_workspaces[0].get(term),
2752                rebuilt_workspaces[0].get(term),
2753                "parallel movement-count mismatch for term {term}"
2754            );
2755        }
2756    }
2757
2758    /// Regression: CSR offsets were u32 and wrapped past 4.29B postings —
2759    /// a 58M-doc / ~85-dims-per-doc prod reorder pass (~4.9B postings)
2760    /// panicked with "mid > len" in the terms carving. The old 8 GB memory
2761    /// budget masked the overflow by dropping dims; raising the budget
2762    /// exposed it. Offsets must be u64.
2763    #[test]
2764    fn test_csr_offsets_do_not_wrap_past_u32() {
2765        let counts = [1_500_000_000u32; 3]; // 4.5B total > u32::MAX
2766        let offsets = build_csr_offsets(&counts);
2767        assert_eq!(
2768            offsets,
2769            vec![0, 1_500_000_000, 3_000_000_000, 4_500_000_000]
2770        );
2771        assert!(*offsets.last().unwrap() > u32::MAX as u64);
2772    }
2773
2774    /// Build a simple forward index from (doc_id, terms) pairs.
2775    fn make_fwd(docs: &[&[u32]], num_terms: usize) -> ForwardIndex {
2776        let mut terms = Vec::new();
2777        let mut offsets = vec![0u64];
2778        for doc_terms in docs {
2779            terms.extend_from_slice(doc_terms);
2780            offsets.push(terms.len() as u64);
2781        }
2782        ForwardIndex {
2783            terms,
2784            offsets,
2785            num_terms,
2786            parallel_bisect_lanes: 1,
2787            budget_limited: false,
2788        }
2789    }
2790
2791    #[test]
2792    fn level_partition_ranges_match_recursive_halving() {
2793        for total in 1..=129 {
2794            let mut expected: Vec<std::ops::Range<usize>> = std::iter::once(0..total).collect();
2795            for level in 0..8 {
2796                let actual: Vec<_> = (0..(1usize << level))
2797                    .map(|partition_id| partition_range(total, level, partition_id))
2798                    .collect();
2799                assert_eq!(actual, expected, "total={total}, level={level}");
2800
2801                expected = expected
2802                    .into_iter()
2803                    .flat_map(|range| {
2804                        let mid = range.start + range.len() / 2;
2805                        [range.start..mid, mid..range.end]
2806                    })
2807                    .collect();
2808            }
2809        }
2810    }
2811
2812    #[cfg(feature = "native")]
2813    #[test]
2814    fn level_synchronized_scheduler_is_thread_count_deterministic() {
2815        let docs: Vec<Vec<u32>> = (0..1_027)
2816            .map(|doc| {
2817                vec![
2818                    (doc % 31) as u32,
2819                    ((doc / 5) % 53) as u32,
2820                    ((doc * 29 + 3) % 97) as u32,
2821                ]
2822            })
2823            .collect();
2824        let doc_refs: Vec<_> = docs.iter().map(Vec::as_slice).collect();
2825        let mut fwd = make_fwd(&doc_refs, 97);
2826        fwd.parallel_bisect_lanes = 8;
2827        let one_thread = rayon::ThreadPoolBuilder::new()
2828            .num_threads(1)
2829            .build()
2830            .unwrap();
2831        let eight_threads = rayon::ThreadPoolBuilder::new()
2832            .num_threads(8)
2833            .build()
2834            .unwrap();
2835
2836        let one = one_thread.install(|| graph_bisection(&fwd, 8, 12, BpBudget::full()).0);
2837        let eight = eight_threads.install(|| graph_bisection(&fwd, 8, 12, BpBudget::full()).0);
2838
2839        assert_eq!(eight, one);
2840    }
2841
2842    /// Scheduler benchmark on a posting-skewed graph. Run manually:
2843    /// `cargo test -p hermes-core --release --features native \
2844    ///    bench_level_synchronized_scheduler -- --ignored --nocapture`
2845    #[cfg(feature = "native")]
2846    #[test]
2847    #[ignore]
2848    fn bench_level_synchronized_scheduler() {
2849        const DOCS: usize = 500_000;
2850        const TERMS: usize = 32_768;
2851        const ROUNDS: usize = 3;
2852
2853        let mut terms = Vec::with_capacity(DOCS * 12);
2854        let mut offsets = Vec::with_capacity(DOCS + 1);
2855        offsets.push(0);
2856        for doc in 0..DOCS {
2857            // Make posting work deliberately uneven between neighboring
2858            // neighboring tree regions. Dynamic same-level claiming should
2859            // absorb this skew instead of pinning it to one lane.
2860            let term_count = 4 + ((doc.wrapping_mul(2_654_435_761) >> 12) % 17);
2861            for term in 0..term_count {
2862                terms.push(
2863                    ((doc / 64)
2864                        .wrapping_mul(131)
2865                        .wrapping_add(term.wrapping_mul(7_919))
2866                        % TERMS) as u32,
2867                );
2868            }
2869            offsets.push(terms.len() as u64);
2870        }
2871        let fwd = ForwardIndex {
2872            terms,
2873            offsets,
2874            num_terms: TERMS,
2875            parallel_bisect_lanes: 8,
2876            budget_limited: false,
2877        };
2878        let pool = rayon::ThreadPoolBuilder::new()
2879            .num_threads(8)
2880            .build()
2881            .unwrap();
2882        let mut level_times = Vec::with_capacity(ROUNDS);
2883        let mut expected_output = None;
2884
2885        pool.install(|| {
2886            for _ in 0..ROUNDS {
2887                let started = std::time::Instant::now();
2888                let output = graph_bisection(&fwd, 32, 12, BpBudget::full()).0;
2889                level_times.push(started.elapsed());
2890                if let Some(expected) = &expected_output {
2891                    assert_eq!(&output, expected);
2892                } else {
2893                    expected_output = Some(output);
2894                }
2895            }
2896        });
2897
2898        level_times.sort_unstable();
2899        let level = level_times[ROUNDS / 2];
2900        println!(
2901            "BP level-synchronized scheduler median: {:.3}s",
2902            level.as_secs_f64(),
2903        );
2904    }
2905
2906    #[test]
2907    fn test_bp_empty() {
2908        let fwd = ForwardIndex {
2909            terms: Vec::new(),
2910            offsets: Vec::new(),
2911            num_terms: 0,
2912            parallel_bisect_lanes: 1,
2913            budget_limited: false,
2914        };
2915        let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2916        assert!(perm.is_empty());
2917    }
2918
2919    #[test]
2920    fn test_bp_small() {
2921        // 4 docs, min_partition_size=4 → no bisection, identity
2922        let fwd = make_fwd(&[&[0, 1], &[0, 2], &[1, 3], &[2, 3]], 4);
2923        let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2924        assert_eq!(perm.len(), 4);
2925        // All docs present
2926        let mut sorted = perm.clone();
2927        sorted.sort();
2928        assert_eq!(sorted, vec![0, 1, 2, 3]);
2929    }
2930
2931    #[test]
2932    fn test_bp_clusters() {
2933        // 8 docs in 2 clear clusters:
2934        // Cluster A (docs 0-3): share terms 0, 1
2935        // Cluster B (docs 4-7): share terms 2, 3
2936        let fwd = make_fwd(
2937            &[
2938                &[0, 1],
2939                &[0, 1],
2940                &[0, 1],
2941                &[0, 1],
2942                &[2, 3],
2943                &[2, 3],
2944                &[2, 3],
2945                &[2, 3],
2946            ],
2947            4,
2948        );
2949        let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2950        assert_eq!(perm.len(), 8);
2951
2952        // After bisection, docs from same cluster should be in same half
2953        let left: Vec<u32> = perm[..4].to_vec();
2954
2955        // Either all of cluster A is in left and B in right, or vice versa
2956        let a_in_left = left.iter().filter(|&&d| d < 4).count();
2957        let b_in_left = left.iter().filter(|&&d| d >= 4).count();
2958        assert!(
2959            (a_in_left == 4 && b_in_left == 0) || (a_in_left == 0 && b_in_left == 4),
2960            "Clusters should be separated: a_left={}, b_left={}",
2961            a_in_left,
2962            b_in_left,
2963        );
2964    }
2965
2966    #[test]
2967    fn test_bp_permutation_valid() {
2968        // 16 docs with mixed terms: terms range from 0..4 and 10..18
2969        let docs: Vec<Vec<u32>> = (0..16).map(|i| vec![i / 4, 10 + i / 2]).collect();
2970        let doc_refs: Vec<&[u32]> = docs.iter().map(|v| v.as_slice()).collect();
2971        let fwd = make_fwd(&doc_refs, 18); // max term = 10 + 15/2 = 17, so need 18
2972        let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2973
2974        assert_eq!(perm.len(), 16);
2975        // Must be a valid permutation
2976        let mut sorted = perm.clone();
2977        sorted.sort();
2978        let expected: Vec<u32> = (0..16).collect();
2979        assert_eq!(sorted, expected);
2980    }
2981
2982    /// Depth-capped BP: with min_partition_docs above the cluster size, only
2983    /// the top-level split happens — clusters still separate (coarse
2984    /// clustering), the permutation stays valid, and the pass converges
2985    /// (a depth cap is a chosen target, not an interruption).
2986    #[test]
2987    fn test_bp_depth_cap_separates_clusters_and_converges() {
2988        // Mostly-separated clusters with one misplaced doc per half — the
2989        // top-level swap pass must exchange docs 3 and 4.
2990        let fwd = make_fwd(
2991            &[
2992                &[0, 1],
2993                &[0, 1],
2994                &[0, 1],
2995                &[2, 3],
2996                &[0, 1],
2997                &[2, 3],
2998                &[2, 3],
2999                &[2, 3],
3000            ],
3001            4,
3002        );
3003        let budget = BpBudget {
3004            min_partition_docs: Some(4),
3005            time_budget: None,
3006        };
3007        let (perm, converged) = graph_bisection(&fwd, 2, 20, budget);
3008        assert!(converged, "depth cap must report converged");
3009        assert_eq!(perm.len(), 8);
3010        let mut sorted = perm.clone();
3011        sorted.sort();
3012        assert_eq!(
3013            sorted,
3014            (0..8).collect::<Vec<u32>>(),
3015            "must stay a valid permutation"
3016        );
3017        // Top-level split separates the clusters (docs {0,1,2,4} share terms
3018        // 0/1; docs {3,5,6,7} share terms 2/3)
3019        let cluster_a = [0u32, 1, 2, 4];
3020        let a_in_left = perm[..4].iter().filter(|d| cluster_a.contains(d)).count();
3021        assert!(
3022            a_in_left == 4 || a_in_left == 0,
3023            "clusters should separate at the top level: {:?}",
3024            perm
3025        );
3026    }
3027
3028    /// Zero wall-clock budget: the pass ends immediately, reports
3029    /// converged=false, and still emits a valid (identity) permutation.
3030    #[test]
3031    fn test_bp_zero_time_budget_emits_valid_partial_permutation() {
3032        let docs: Vec<Vec<u32>> = (0..64).map(|i| vec![i % 4]).collect();
3033        let doc_refs: Vec<&[u32]> = docs.iter().map(|v| v.as_slice()).collect();
3034        let fwd = make_fwd(&doc_refs, 4);
3035        let budget = BpBudget {
3036            min_partition_docs: None,
3037            time_budget: Some(std::time::Duration::ZERO),
3038        };
3039        let (perm, converged) = graph_bisection(&fwd, 4, 20, budget);
3040        assert!(!converged, "zero budget must report unconverged");
3041        assert_eq!(perm.len(), 64);
3042        let mut sorted = perm.clone();
3043        sorted.sort();
3044        assert_eq!(sorted, (0..64).collect::<Vec<u32>>());
3045    }
3046
3047    #[test]
3048    fn test_bp_shutdown_cancellation_emits_valid_partial_permutation() {
3049        let docs: Vec<Vec<u32>> = (0..64).map(|i| vec![i % 4]).collect();
3050        let doc_refs: Vec<&[u32]> = docs.iter().map(Vec::as_slice).collect();
3051        let fwd = make_fwd(&doc_refs, 4);
3052        let cancellation = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
3053        let budget = BpBudget {
3054            min_partition_docs: None,
3055            time_budget: None,
3056        };
3057
3058        let (perm, converged) = graph_bisection_with_progress(
3059            &fwd,
3060            4,
3061            20,
3062            budget,
3063            Some(cancellation.as_ref()),
3064            BpProgressLabel::anonymous(),
3065        );
3066
3067        assert!(!converged, "cancelled BP must report unconverged");
3068        assert_eq!(perm, (0..64).collect::<Vec<u32>>());
3069    }
3070
3071    #[test]
3072    fn test_memory_limited_graph_never_reports_converged() {
3073        let mut fwd = make_fwd(&[&[0], &[0], &[1], &[1]], 2);
3074        fwd.budget_limited = true;
3075
3076        let (perm, converged) = graph_bisection(&fwd, 2, 20, BpBudget::full());
3077
3078        assert!(!converged);
3079        let mut sorted = perm;
3080        sorted.sort_unstable();
3081        assert_eq!(sorted, vec![0, 1, 2, 3]);
3082    }
3083
3084    #[test]
3085    fn test_fast_log2() {
3086        let table = build_log_table(4096);
3087        assert!((table[1] - 0.0).abs() < 0.001);
3088        assert!((table[2] - 1.0).abs() < 0.001);
3089        assert!((table[4] - 2.0).abs() < 0.001);
3090        assert!((table[1024] - 10.0).abs() < 0.001);
3091        // Fallback for values beyond table
3092        let val = fast_log2_lookup(8192, &table);
3093        assert!((val - 13.0).abs() < 0.001);
3094    }
3095}