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