use crate::comparator::UserComparator;
use crate::loser_tree::LoserTree;
use crate::merge_source::{CoherentMergeSource, IterItem, MergeSource};
use alloc::vec::Vec;
use core::cmp::Ordering;
struct MergerEntry {
source: usize,
value: crate::InternalValue,
}
#[derive(Clone)]
struct MinCmp<C: UserComparator + Clone> {
comparator: C,
}
impl<C: UserComparator + Clone> crate::loser_tree::EntryComparator<MergerEntry> for MinCmp<C> {
#[expect(
clippy::inline_always,
reason = "called O(log cap) per replay step on the merger's hot path; \
matching the loser-tree's own #[inline(always)] on cmp_indices \
is what makes the dispatch flatten — verified in disassembly"
)]
#[inline(always)]
fn compare(&self, a: &MergerEntry, b: &MergerEntry) -> Ordering {
a.value
.key
.compare_with(&b.value.key, &self.comparator)
.then_with(|| a.source.cmp(&b.source))
}
}
#[derive(Clone)]
struct MaxCmp<C: UserComparator + Clone> {
comparator: C,
}
impl<C: UserComparator + Clone> crate::loser_tree::EntryComparator<MergerEntry> for MaxCmp<C> {
#[expect(
clippy::inline_always,
reason = "called O(log cap) per replay step on the merger's hot path"
)]
#[inline(always)]
fn compare(&self, a: &MergerEntry, b: &MergerEntry) -> Ordering {
b.value
.key
.compare_with(&a.value.key, &self.comparator)
.then_with(|| b.source.cmp(&a.source))
}
}
fn build_min_cmp<C: UserComparator + Clone>(comparator: C) -> MinCmp<C> {
MinCmp { comparator }
}
fn build_max_cmp<C: UserComparator + Clone>(comparator: C) -> MaxCmp<C> {
MaxCmp { comparator }
}
#[derive(Clone, Copy)]
enum ParkedCache {
Stale,
Empty,
Best(usize),
}
fn pick_parked<A, O, C>(
active: &LoserTree<MergerEntry, A>,
other: Option<&LoserTree<MergerEntry, O>>,
vacated: &[usize],
cache: &mut ParkedCache,
cmp: &C,
) -> Option<usize>
where
A: crate::loser_tree::EntryComparator<MergerEntry>,
O: crate::loser_tree::EntryComparator<MergerEntry>,
C: crate::loser_tree::EntryComparator<MergerEntry>,
{
let other = other?;
if vacated.is_empty() || other.is_empty() {
return None;
}
let best_slot = match *cache {
ParkedCache::Best(slot) if other.peek_slot(slot).is_some() => Some(slot),
ParkedCache::Empty => None,
ParkedCache::Stale | ParkedCache::Best(_) => {
let mut best: Option<(usize, &MergerEntry)> = None;
for &slot in vacated {
debug_assert!(
active.peek_slot(slot).is_none(),
"vacated slot holds a leaf in the active tournament"
);
let Some(parked) = other.peek_slot(slot) else {
continue;
};
if best.is_none_or(|(_, current)| cmp.compare(parked, current) == Ordering::Less) {
best = Some((slot, parked));
}
}
let computed = best.map(|(slot, _)| slot);
*cache = computed.map_or(ParkedCache::Empty, ParkedCache::Best);
computed
}
};
let slot = best_slot?;
let parked = other.peek_slot(slot)?;
let winner = active.winner_slot().and_then(|slot| active.peek_slot(slot));
match winner {
Some(w) => (cmp.compare(parked, w) == Ordering::Less).then_some(slot),
None => Some(slot),
}
}
pub struct SeekingMerger<S: MergeSource, C: UserComparator + Clone> {
sources: Vec<S>,
n_sources: usize,
comparator: C,
forward_tree: Option<LoserTree<MergerEntry, MinCmp<C>>>,
backward_tree: Option<LoserTree<MergerEntry, MaxCmp<C>>>,
pending_forward_error: Option<crate::Error>,
pending_backward_error: Option<crate::Error>,
forward_primed: bool,
backward_primed: bool,
forward_vacated: Vec<usize>,
backward_vacated: Vec<usize>,
forward_parked_best: ParkedCache,
backward_parked_best: ParkedCache,
}
impl<S: MergeSource, C: UserComparator + Clone> SeekingMerger<S, C> {
#[must_use]
pub fn new(sources: Vec<S>, comparator: C) -> Self {
let n = sources.len();
Self {
sources,
n_sources: n,
comparator,
forward_tree: None,
backward_tree: None,
pending_forward_error: None,
pending_backward_error: None,
forward_primed: false,
backward_primed: false,
forward_vacated: Vec::new(),
backward_vacated: Vec::new(),
forward_parked_best: ParkedCache::Stale,
backward_parked_best: ParkedCache::Stale,
}
}
fn initialize_forward(&mut self) {
let Self {
sources,
n_sources,
comparator,
forward_tree,
backward_tree,
pending_forward_error,
forward_vacated,
backward_vacated,
forward_parked_best,
backward_parked_best,
..
} = self;
let n = *n_sources;
let mut pull = |i: usize| -> Option<MergerEntry> {
#[expect(
clippy::indexing_slicing,
reason = "i < n_sources by construction; sources len == n_sources"
)]
match MergeSource::next(&mut sources[i]) {
Some(Ok(value)) => Some(MergerEntry { source: i, value }),
Some(Err(e)) => {
pending_forward_error.get_or_insert(e);
None
}
None => {
let migrated = backward_tree.as_mut().and_then(|bt| bt.take_slot(i));
if migrated.is_some() {
backward_vacated.push(i);
}
migrated
}
}
};
let tree = if let Some(tree) = forward_tree {
tree.refill_with(&mut pull);
tree
} else {
let mut initial: Vec<Option<MergerEntry>> = Vec::with_capacity(n);
for i in 0..n {
initial.push(pull(i));
}
let cmp = build_min_cmp(comparator.clone());
forward_tree.insert(LoserTree::build(initial, cmp))
};
forward_vacated.clear();
forward_vacated.extend((0..n).filter(|&i| tree.peek_slot(i).is_none()));
*forward_parked_best = ParkedCache::Stale;
*backward_parked_best = ParkedCache::Stale;
}
fn initialize_backward(&mut self) {
let Self {
sources,
n_sources,
comparator,
forward_tree,
backward_tree,
pending_backward_error,
forward_vacated,
backward_vacated,
forward_parked_best,
backward_parked_best,
..
} = self;
let n = *n_sources;
let mut pull = |i: usize| -> Option<MergerEntry> {
#[expect(
clippy::indexing_slicing,
reason = "i < n_sources by construction; sources len == n_sources"
)]
match MergeSource::next_back(&mut sources[i]) {
Some(Ok(value)) => Some(MergerEntry { source: i, value }),
Some(Err(e)) => {
pending_backward_error.get_or_insert(e);
None
}
None => {
let migrated = forward_tree.as_mut().and_then(|ft| ft.take_slot(i));
if migrated.is_some() {
forward_vacated.push(i);
}
migrated
}
}
};
let tree = if let Some(tree) = backward_tree {
tree.refill_with(&mut pull);
tree
} else {
let mut initial: Vec<Option<MergerEntry>> = Vec::with_capacity(n);
for i in 0..n {
initial.push(pull(i));
}
let cmp = build_max_cmp(comparator.clone());
backward_tree.insert(LoserTree::build(initial, cmp))
};
backward_vacated.clear();
backward_vacated.extend((0..n).filter(|&i| tree.peek_slot(i).is_none()));
*forward_parked_best = ParkedCache::Stale;
*backward_parked_best = ParkedCache::Stale;
}
}
impl<S: MergeSource + crate::reseek::Reseekable, C: UserComparator + Clone>
crate::reseek::Reseekable for SeekingMerger<S, C>
{
fn reseek(&mut self, ctx: &crate::reseek::ReseekCtx) {
for source in &mut self.sources {
source.reseek(ctx);
}
if let Some(tree) = &mut self.forward_tree {
tree.clear();
}
if let Some(tree) = &mut self.backward_tree {
tree.clear();
}
self.forward_primed = false;
self.backward_primed = false;
self.forward_vacated.clear();
self.backward_vacated.clear();
self.forward_parked_best = ParkedCache::Stale;
self.backward_parked_best = ParkedCache::Stale;
self.pending_forward_error = None;
self.pending_backward_error = None;
}
}
impl<S: MergeSource, C: UserComparator + Clone> Iterator for SeekingMerger<S, C> {
type Item = IterItem;
fn next(&mut self) -> Option<Self::Item> {
if let Some(e) = self.pending_forward_error.take() {
return Some(Err(e));
}
if let Some(e) = self.pending_backward_error.take() {
return Some(Err(e));
}
if !self.forward_primed {
self.initialize_forward();
self.forward_primed = true;
if let Some(e) = self.pending_forward_error.take() {
return Some(Err(e));
}
}
let Self {
sources,
forward_tree,
backward_tree,
pending_forward_error,
comparator,
forward_vacated,
backward_vacated,
forward_parked_best,
backward_parked_best,
..
} = self;
let tree = forward_tree.as_mut()?;
*backward_parked_best = ParkedCache::Stale;
if let Some(slot) = pick_parked(
tree,
backward_tree.as_ref(),
forward_vacated,
forward_parked_best,
&build_min_cmp(comparator.clone()),
) && let Some(entry) = backward_tree.as_mut().and_then(|bt| bt.take_slot(slot))
{
backward_vacated.push(slot);
*forward_parked_best = ParkedCache::Stale;
return Some(Ok(entry.value));
}
let source = tree.winner_slot()?;
#[expect(
clippy::indexing_slicing,
reason = "source index < n_sources by construction"
)]
let next_pull = MergeSource::next(&mut sources[source]);
match next_pull {
Some(Ok(next_value)) => {
let old = tree.replace_min(MergerEntry {
source,
value: next_value,
});
Some(Ok(old.value))
}
Some(Err(e)) => {
let old = tree.pop_min()?;
pending_forward_error.get_or_insert(e);
forward_vacated.push(source);
*forward_parked_best = ParkedCache::Stale;
Some(Ok(old.value))
}
None => {
let old = tree.pop_min()?;
forward_vacated.push(source);
*forward_parked_best = ParkedCache::Stale;
Some(Ok(old.value))
}
}
}
}
impl<S: CoherentMergeSource, C: UserComparator + Clone> DoubleEndedIterator
for SeekingMerger<S, C>
{
fn next_back(&mut self) -> Option<Self::Item> {
if let Some(e) = self.pending_backward_error.take() {
return Some(Err(e));
}
if let Some(e) = self.pending_forward_error.take() {
return Some(Err(e));
}
if !self.backward_primed {
self.initialize_backward();
self.backward_primed = true;
if let Some(e) = self.pending_backward_error.take() {
return Some(Err(e));
}
}
let Self {
sources,
forward_tree,
backward_tree,
pending_backward_error,
comparator,
forward_vacated,
backward_vacated,
forward_parked_best,
backward_parked_best,
..
} = self;
let tree = backward_tree.as_mut()?;
*forward_parked_best = ParkedCache::Stale;
if let Some(slot) = pick_parked(
tree,
forward_tree.as_ref(),
backward_vacated,
backward_parked_best,
&build_max_cmp(comparator.clone()),
) && let Some(entry) = forward_tree.as_mut().and_then(|ft| ft.take_slot(slot))
{
forward_vacated.push(slot);
*backward_parked_best = ParkedCache::Stale;
return Some(Ok(entry.value));
}
let source = tree.winner_slot()?;
#[expect(
clippy::indexing_slicing,
reason = "source index < n_sources by construction"
)]
let next_pull = MergeSource::next_back(&mut sources[source]);
match next_pull {
Some(Ok(next_value)) => {
let old = tree.replace_min(MergerEntry {
source,
value: next_value,
});
Some(Ok(old.value))
}
Some(Err(e)) => {
let old = tree.pop_min()?;
pending_backward_error.get_or_insert(e);
backward_vacated.push(source);
*backward_parked_best = ParkedCache::Stale;
Some(Ok(old.value))
}
None => {
let old = tree.pop_min()?;
backward_vacated.push(source);
*backward_parked_best = ParkedCache::Stale;
Some(Ok(old.value))
}
}
}
}
#[cfg(test)]
#[expect(clippy::unwrap_used, reason = "test assertions")]
mod tests;