extern crate alloc;
use alloc::vec::Vec;
use super::{Dot, NIL, NodeKind, OrderThread};
pub(super) fn nth_set_bit(mask: u64, k: u32) -> u32 {
debug_assert!(k < mask.count_ones());
let mut low = 0u32;
let mut remaining = k;
let mut mask = mask;
let mut width = 32u32;
while width > 0 {
let low_half = mask & ((1u64 << width) - 1);
let ones = low_half.count_ones();
if remaining < ones {
mask = low_half;
} else {
remaining -= ones;
low += width;
mask >>= width;
}
width /= 2;
}
low
}
impl OrderThread {
pub(in crate::metis::rhapsody) fn slot_len(&self) -> usize {
if self.root == NIL {
0
} else {
self.nodes[self.root as usize].slots
}
}
pub(in crate::metis::rhapsody) fn visible_len(&self) -> usize {
if self.root == NIL {
0
} else {
self.nodes[self.root as usize].visible
}
}
pub(super) fn locate(&self, dot: Dot) -> Option<(u32, usize, u32)> {
let (&base_dot, &leaf) = self
.index
.range(..=dot)
.next_back()
.filter(|entry| entry.0.station() == dot.station())?;
let offset = dot.counter().checked_sub(base_dot.counter())?;
let NodeKind::Leaf(frags) = &self.nodes[leaf as usize].kind else {
unreachable!("the fragment index points at leaves only");
};
let position = frags.iter().position(|frag| frag.base_dot() == base_dot)?;
(offset < u64::from(frags[position].len)).then(|| {
(
leaf,
position,
u32::try_from(offset).expect("fragment offsets fit FRAG_CAP"),
)
})
}
pub(in crate::metis::rhapsody) fn position_of(&self, dot: Dot) -> Option<(usize, usize)> {
let (leaf, frag_position, offset) = self.locate(dot)?;
let NodeKind::Leaf(frags) = &self.nodes[leaf as usize].kind else {
unreachable!("locate returns leaves only");
};
let frag = &frags[frag_position];
let below = (1u64 << offset) - 1;
let mut slot_position = offset as usize;
let mut visible_before = (frag.visible & below).count_ones() as usize;
for earlier in &frags[..frag_position] {
slot_position += earlier.slots();
visible_before += earlier.visible_count();
}
let mut child = leaf;
let mut parent = self.nodes[leaf as usize].parent;
while parent != NIL {
let NodeKind::Internal(children) = &self.nodes[parent as usize].kind else {
unreachable!("a parent is internal");
};
for &sibling in children.iter().take_while(|&&candidate| candidate != child) {
let node = &self.nodes[sibling as usize];
slot_position += node.slots;
visible_before += node.visible;
}
child = parent;
parent = self.nodes[parent as usize].parent;
}
Some((slot_position, visible_before))
}
pub(in crate::metis::rhapsody) fn order_at(&self, k: usize) -> Option<Dot> {
if self.root == NIL || k >= self.nodes[self.root as usize].visible {
return None;
}
let mut id = self.root;
let mut remaining = k;
loop {
match &self.nodes[id as usize].kind {
NodeKind::Internal(children) => {
let mut next = None;
for &child in children {
let visible = self.nodes[child as usize].visible;
if remaining < visible {
next = Some(child);
break;
}
remaining -= visible;
}
id = next.expect("the root aggregate bounds the descent");
}
NodeKind::Leaf(frags) => {
for frag in frags {
let visible = frag.visible_count();
if remaining < visible {
let bit = nth_set_bit(
frag.visible,
u32::try_from(remaining).expect("fragment counts fit u32"),
);
return Some(Dot::new(
frag.station,
frag.base.saturating_add(u64::from(bit)),
));
}
remaining -= visible;
}
unreachable!("leaf aggregates bound the scan");
}
}
}
}
pub(in crate::metis::rhapsody) fn slots(&self) -> impl Iterator<Item = (Dot, bool)> + '_ {
let mut stack: Vec<u32> = Vec::new();
if self.root != NIL {
stack.push(self.root);
}
let mut pending: Vec<(Dot, bool)> = Vec::new();
core::iter::from_fn(move || {
loop {
if let Some(slot) = pending.pop() {
return Some(slot);
}
let id = stack.pop()?;
match &self.nodes[id as usize].kind {
NodeKind::Internal(children) => stack.extend(children.iter().rev().copied()),
NodeKind::Leaf(frags) => {
for frag in frags.iter().rev() {
for offset in (0..u64::from(frag.len)).rev() {
let dot = Dot::new(frag.station, frag.base.saturating_add(offset));
pending.push((dot, frag.visible & (1u64 << offset) != 0));
}
}
}
}
}
})
}
#[cfg(any(test, feature = "instrumentation"))]
pub(in crate::metis::rhapsody) fn fragments(&self) -> usize {
self.index.len()
}
}