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,
lexicographic: bool,
}
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 {
let keys = if self.lexicographic {
a.value.key.cmp(&b.value.key)
} else {
a.value.key.compare_with(&b.value.key, &self.comparator)
};
keys.then_with(|| a.source.cmp(&b.source))
}
}
#[derive(Clone)]
struct MaxCmp<C: UserComparator + Clone> {
comparator: C,
lexicographic: bool,
}
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 {
let keys = if self.lexicographic {
b.value.key.cmp(&a.value.key)
} else {
b.value.key.compare_with(&a.value.key, &self.comparator)
};
keys.then_with(|| b.source.cmp(&a.source))
}
}
fn build_min_cmp<C: UserComparator + Clone>(comparator: C) -> MinCmp<C> {
let lexicographic = comparator.is_lexicographic();
MinCmp {
comparator,
lexicographic,
}
}
fn build_max_cmp<C: UserComparator + Clone>(comparator: C) -> MaxCmp<C> {
let lexicographic = comparator.is_lexicographic();
MaxCmp {
comparator,
lexicographic,
}
}
#[derive(Clone, Copy)]
enum ParkedCache {
Stale,
Empty,
Best(usize),
}
fn pick_parked<A, O>(
active: &LoserTree<MergerEntry, A>,
other: Option<&LoserTree<MergerEntry, O>>,
vacated: &[usize],
cache: &mut ParkedCache,
) -> Option<usize>
where
A: crate::loser_tree::EntryComparator<MergerEntry>,
O: crate::loser_tree::EntryComparator<MergerEntry>,
{
let other = other?;
if vacated.is_empty() || other.is_empty() {
return None;
}
let cmp = active.comparator();
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,
backward_jumped: Vec<usize>,
backward_set_aside: Vec<MergerEntry>,
back_run_slot: usize,
back_run: u32,
}
const BACK_RUN_BEFORE_JUMP: u32 = 8;
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,
backward_jumped: Vec::new(),
backward_set_aside: Vec::new(),
back_run_slot: usize::MAX,
back_run: 0,
}
}
#[inline]
fn note_back_item(back_run_slot: &mut usize, back_run: &mut u32, slot: usize) {
if *back_run_slot == slot {
*back_run += 1;
} else {
*back_run_slot = slot;
*back_run = 1;
}
}
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;
self.back_run = 0;
}
}
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,
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,
) && 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: MergeSource + crate::skip_key::SkipKey, C: UserComparator + Clone> crate::skip_key::SkipKey
for SeekingMerger<S, C>
{
fn skip_front_past(&mut self, key: &[u8]) -> crate::Result<bool> {
if self.backward_primed
|| !self.forward_primed
|| self.pending_forward_error.is_some()
|| self.pending_backward_error.is_some()
{
return Ok(false);
}
let Self {
sources,
forward_tree,
pending_forward_error,
forward_vacated,
forward_parked_best,
..
} = self;
let Some(tree) = forward_tree.as_mut() else {
return Ok(false);
};
while let Some(slot) = tree.winner_slot() {
let at_key = tree.peek_slot(slot).is_some_and(|head| {
crate::comparator::same_user_key(&head.value.key.user_key, key)
});
if !at_key {
break;
}
let Some(source) = sources.get_mut(slot) else {
break;
};
match source.skip_front_past(key) {
Ok(true) => {}
Ok(false) => return Ok(false),
Err(e) => {
tree.pop_min();
pending_forward_error.get_or_insert(e);
forward_vacated.push(slot);
*forward_parked_best = ParkedCache::Stale;
continue;
}
}
match MergeSource::next(source) {
Some(Ok(value)) => {
tree.replace_min(MergerEntry {
source: slot,
value,
});
}
Some(Err(e)) => {
tree.pop_min();
pending_forward_error.get_or_insert(e);
forward_vacated.push(slot);
*forward_parked_best = ParkedCache::Stale;
}
None => {
tree.pop_min();
forward_vacated.push(slot);
*forward_parked_best = ParkedCache::Stale;
}
}
}
Ok(true)
}
fn skip_back_to_newest(&mut self, key: &[u8], seqno: crate::SeqNo) -> crate::Result<bool> {
if self.back_run < BACK_RUN_BEFORE_JUMP
|| self.forward_primed
|| !self.backward_primed
|| self.pending_forward_error.is_some()
|| self.pending_backward_error.is_some()
{
return Ok(false);
}
let Self {
sources,
backward_tree,
pending_backward_error,
backward_vacated,
backward_parked_best,
backward_jumped,
backward_set_aside,
..
} = self;
let Some(tree) = backward_tree.as_mut() else {
return Ok(false);
};
let mut moved = false;
while let Some(slot) = tree.winner_slot() {
let at_key = tree.peek_slot(slot).is_some_and(|head| {
crate::comparator::same_user_key(&head.value.key.user_key, key)
});
if !at_key {
break;
}
if backward_jumped.contains(&slot) {
if let Some(head) = tree.pop_min() {
backward_set_aside.push(head);
}
continue;
}
backward_jumped.push(slot);
let Some(source) = sources.get_mut(slot) else {
break;
};
match source.skip_back_to_newest(key, seqno) {
Ok(true) => {}
Ok(false) => continue,
Err(e) => {
tree.pop_min();
pending_backward_error.get_or_insert(e);
backward_vacated.push(slot);
*backward_parked_best = ParkedCache::Stale;
moved = true;
continue;
}
}
moved = true;
match MergeSource::next_back(source) {
Some(Ok(value)) => {
tree.replace_min(MergerEntry {
source: slot,
value,
});
}
Some(Err(e)) => {
tree.pop_min();
pending_backward_error.get_or_insert(e);
backward_vacated.push(slot);
*backward_parked_best = ParkedCache::Stale;
}
None => {
tree.pop_min();
backward_vacated.push(slot);
*backward_parked_best = ParkedCache::Stale;
}
}
}
backward_jumped.clear();
if !backward_set_aside.is_empty() {
tree.restore(backward_set_aside.drain(..).map(|head| (head.source, head)));
}
Ok(moved)
}
}
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,
forward_vacated,
backward_vacated,
forward_parked_best,
backward_parked_best,
back_run_slot,
back_run,
..
} = 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,
) && let Some(entry) = forward_tree.as_mut().and_then(|ft| ft.take_slot(slot))
{
forward_vacated.push(slot);
*backward_parked_best = ParkedCache::Stale;
Self::note_back_item(back_run_slot, back_run, slot);
return Some(Ok(entry.value));
}
let source = tree.winner_slot()?;
Self::note_back_item(back_run_slot, back_run, source);
#[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;