From 24a97dbe6961c6b0b8eb492cd7f5c3a48338097d Mon Sep 17 00:00:00 2001
From: pascal <pascal.seitz@gmail.com>
Date: Sun, 31 May 2026 12:20:07 +0200
Subject: [PATCH] Split buffered refill from scorer removal

---
 src/query/union/buffered_union.rs | 64 ++++++++++++++-----------------
 1 file changed, 28 insertions(+), 36 deletions(-)

diff --git a/src/query/union/buffered_union.rs b/src/query/union/buffered_union.rs
index 565e73af1..77505c196 100644
--- a/src/query/union/buffered_union.rs
+++ b/src/query/union/buffered_union.rs
@@ -11,23 +11,6 @@ use crate::{DocId, Score};
 const HORIZON_NUM_TINYBITSETS: usize = HORIZON as usize / 64;
 const HORIZON: u32 = 64u32 * 64u32;
 
-// `drain_filter` is not stable yet.
-// This function is similar except that it does is not unstable, and
-// it does not keep the original vector ordering.
-//
-// Elements are dropped and not yielded.
-fn unordered_drain_filter<T, P>(v: &mut Vec<T>, mut predicate: P)
-where P: FnMut(&mut T) -> bool {
-    let mut i = 0;
-    while i < v.len() {
-        if predicate(&mut v[i]) {
-            v.swap_remove(i);
-        } else {
-            i += 1;
-        }
-    }
-}
-
 /// Creates a `DocSet` that iterate through the union of two or more `DocSet`s.
 pub struct BufferedUnionScorer<TScorer, TScoreCombiner = DoNothingCombiner> {
     /// Active scorers (already filtered of `TERMINATED`).
@@ -55,28 +38,37 @@ pub struct BufferedUnionScorer<TScorer, TScoreCombiner = DoNothingCombiner> {
     num_docs: u32,
 }
 
+fn refill_scorer<TScorer: Scorer, TScoreCombiner: ScoreCombiner>(
+    scorer: &mut TScorer,
+    bitsets: &mut [TinySet; HORIZON_NUM_TINYBITSETS],
+    score_combiner: &mut [TScoreCombiner; HORIZON as usize],
+    min_doc: DocId,
+    horizon: DocId,
+) {
+    loop {
+        let doc = scorer.doc();
+        if doc >= horizon {
+            break;
+        }
+        // add this document
+        let delta = doc - min_doc;
+        bitsets[(delta / 64) as usize].insert_mut(delta % 64u32);
+        score_combiner[delta as usize].update(scorer);
+        scorer.advance();
+    }
+}
+
 fn refill<TScorer: Scorer, TScoreCombiner: ScoreCombiner>(
     scorers: &mut Vec<TScorer>,
     bitsets: &mut [TinySet; HORIZON_NUM_TINYBITSETS],
     score_combiner: &mut [TScoreCombiner; HORIZON as usize],
     min_doc: DocId,
 ) {
-    unordered_drain_filter(scorers, |scorer| {
-        let horizon = min_doc + HORIZON;
-        loop {
-            let doc = scorer.doc();
-            if doc >= horizon {
-                break;
-            }
-            // add this document
-            let delta = doc - min_doc;
-            bitsets[(delta / 64) as usize].insert_mut(delta % 64u32);
-            score_combiner[delta as usize].update(scorer);
-            scorer.advance();
-        }
-        // remove the docset if it has been entirely consumed.
-        scorer.doc() == TERMINATED
-    });
+    let horizon = min_doc + HORIZON;
+    for scorer in scorers.iter_mut() {
+        refill_scorer(scorer, bitsets, score_combiner, min_doc, horizon);
+    }
+    scorers.retain(|scorer| scorer.doc() != TERMINATED);
 }
 
 impl<TScorer: Scorer, TScoreCombiner: ScoreCombiner> BufferedUnionScorer<TScorer, TScoreCombiner> {
@@ -247,12 +239,12 @@ where
 
             // The target is outside of the buffered horizon.
             // advance all docsets to a doc >= to the target.
-            unordered_drain_filter(&mut self.docsets, |docset| {
+            for docset in &mut self.docsets {
                 if docset.doc() < target {
                     docset.seek(target);
                 }
-                docset.doc() == TERMINATED
-            });
+            }
+            self.docsets.retain(|docset| docset.doc() != TERMINATED);
 
             // at this point all of the docsets
             // are positioned on a doc >= to the target.