Skip to main content

hermes_core/segment/builder/
graph_bisection.rs

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