use crate::utils::SelectInWord;
use ambassador::Delegate;
use mem_dbg::{MemDbg, MemSize};
use std::{
cmp::{max, min},
ops::Deref,
};
use crate::{
prelude::{BitCount, BitLength, Select, SelectHinted},
traits::{
NumBits, Rank, RankHinted, RankUnchecked, RankZero, SelectUnchecked, SelectZero,
SelectZeroHinted, SelectZeroUnchecked,
},
};
use crate::ambassador_impl_AsRef;
use crate::ambassador_impl_Index;
use crate::traits::rank_sel::ambassador_impl_BitCount;
use crate::traits::rank_sel::ambassador_impl_BitLength;
use crate::traits::rank_sel::ambassador_impl_NumBits;
use crate::traits::rank_sel::ambassador_impl_Rank;
use crate::traits::rank_sel::ambassador_impl_RankHinted;
use crate::traits::rank_sel::ambassador_impl_RankUnchecked;
use crate::traits::rank_sel::ambassador_impl_RankZero;
use crate::traits::rank_sel::ambassador_impl_SelectHinted;
use crate::traits::rank_sel::ambassador_impl_SelectZero;
use crate::traits::rank_sel::ambassador_impl_SelectZeroHinted;
use crate::traits::rank_sel::ambassador_impl_SelectZeroUnchecked;
use std::ops::Index;
#[derive(Debug, Clone, Copy, MemDbg, MemSize, Delegate)]
#[cfg_attr(feature = "epserde", derive(epserde::Epserde))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[delegate(AsRef<[usize]>, target = "bits")]
#[delegate(Index<usize>, target = "bits")]
#[delegate(crate::traits::rank_sel::BitCount, target = "bits")]
#[delegate(crate::traits::rank_sel::BitLength, target = "bits")]
#[delegate(crate::traits::rank_sel::NumBits, target = "bits")]
#[delegate(crate::traits::rank_sel::Rank, target = "bits")]
#[delegate(crate::traits::rank_sel::RankHinted<64>, target = "bits")]
#[delegate(crate::traits::rank_sel::RankUnchecked, target = "bits")]
#[delegate(crate::traits::rank_sel::RankZero, target = "bits")]
#[delegate(crate::traits::rank_sel::SelectHinted, target = "bits")]
#[delegate(crate::traits::rank_sel::SelectZero, target = "bits")]
#[delegate(crate::traits::rank_sel::SelectZeroHinted, target = "bits")]
#[delegate(crate::traits::rank_sel::SelectZeroUnchecked, target = "bits")]
pub struct SelectAdapt<B, I = Box<[usize]>> {
bits: B,
inventory: I,
spill: I,
log2_ones_per_inventory: usize,
log2_ones_per_sub16: usize,
log2_u64_per_subinventory: usize,
ones_per_inventory_mask: usize,
ones_per_sub16_mask: usize,
}
impl<B, I> Deref for SelectAdapt<B, I> {
type Target = B;
fn deref(&self) -> &Self::Target {
&self.bits
}
}
pub(super) trait Inventory {
fn is_u16_span(&self) -> bool;
fn is_u32_span(&self) -> bool;
fn is_u64_span(&self) -> bool;
fn set_u16_span(&mut self);
fn set_u32_span(&mut self);
fn set_u64_span(&mut self);
fn get(&self) -> usize;
}
impl Inventory for usize {
#[inline(always)]
fn is_u16_span(&self) -> bool {
*self as isize >= 0
}
#[inline(always)]
fn is_u32_span(&self) -> bool {
*self >> 62 == 2
}
#[inline(always)]
fn is_u64_span(&self) -> bool {
*self >> 62 == 3
}
#[inline(always)]
fn set_u16_span(&mut self) {}
#[inline(always)]
fn set_u32_span(&mut self) {
*self |= 1 << 63;
}
#[inline(always)]
fn set_u64_span(&mut self) {
*self |= 3 << 62;
}
#[inline(always)]
fn get(&self) -> usize {
*self & 0x3FFF_FFFF_FFFF_FFFF
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum SpanType {
U16,
U32,
U64,
}
impl SpanType {
pub fn from_span(x: usize) -> SpanType {
match x {
0..=0x10000 => SpanType::U16,
0x10001..=0x100000000 => SpanType::U32,
_ => SpanType::U64,
}
}
}
impl<B, I> SelectAdapt<B, I> {
pub fn into_inner(self) -> B {
self.bits
}
#[inline(always)]
const fn log2_ones_per_sub32(span: usize, log2_ones_per_sub16: usize) -> usize {
debug_assert!(span >= 1 << 16);
log2_ones_per_sub16.saturating_sub((span >> 15).ilog2() as usize + 1)
}
pub unsafe fn map<C>(self, f: impl FnOnce(B) -> C) -> SelectAdapt<C, I>
where
C: SelectHinted,
{
SelectAdapt {
bits: f(self.bits),
inventory: self.inventory,
spill: self.spill,
log2_ones_per_inventory: self.log2_ones_per_inventory,
log2_ones_per_sub16: self.log2_ones_per_sub16,
log2_u64_per_subinventory: self.log2_u64_per_subinventory,
ones_per_inventory_mask: self.ones_per_inventory_mask,
ones_per_sub16_mask: self.ones_per_sub16_mask,
}
}
pub const DEFAULT_TARGET_INVENTORY_SPAN: usize = 8192;
}
impl<B: BitLength, C> SelectAdapt<B, C> {
#[inline(always)]
pub fn len(&self) -> usize {
BitLength::len(self)
}
}
impl<B: AsRef<[usize]> + BitCount> SelectAdapt<B, Box<[usize]>> {
pub fn new(bits: B, max_log2_u64_per_subinv: usize) -> Self {
Self::with_span(
bits,
Self::DEFAULT_TARGET_INVENTORY_SPAN,
max_log2_u64_per_subinv,
)
}
pub fn with_span(
bits: B,
target_inventory_span: usize,
max_log2_u64_per_subinventory: usize,
) -> Self {
let num_bits = max(1usize, bits.len());
let num_ones = bits.count_ones();
let log2_ones_per_inventory = (num_ones * target_inventory_span)
.div_ceil(num_bits)
.max(1)
.ilog2() as usize;
Self::_new(
bits,
num_ones,
log2_ones_per_inventory,
max_log2_u64_per_subinventory,
)
}
pub fn with_inv(
bits: B,
log2_ones_per_inventory: usize,
max_log2_u64_per_subinventory: usize,
) -> Self {
let num_ones = bits.count_ones();
Self::_new(
bits,
num_ones,
log2_ones_per_inventory,
max_log2_u64_per_subinventory,
)
}
fn _new(
bits: B,
num_ones: usize,
log2_ones_per_inventory: usize,
max_log2_u64_per_subinventory: usize,
) -> Self {
let num_bits = max(1, bits.len());
let ones_per_inventory = 1 << log2_ones_per_inventory;
let ones_per_inventory_mask = ones_per_inventory - 1;
let inventory_size = num_ones.div_ceil(ones_per_inventory);
let log2_u64_per_subinventory =
max_log2_u64_per_subinventory.min(log2_ones_per_inventory.saturating_sub(2));
let u64_per_subinventory = 1 << log2_u64_per_subinventory;
let u64_per_inventory = u64_per_subinventory + 1;
let log2_ones_per_sub16 =
log2_ones_per_inventory.saturating_sub(log2_u64_per_subinventory + 2);
let ones_per_sub16 = 1 << log2_ones_per_sub16;
let ones_per_sub16_mask = ones_per_sub16 - 1;
let inventory_words = inventory_size * u64_per_inventory + 1;
let mut inventory = Vec::with_capacity(inventory_words);
let mut past_ones = 0;
let mut next_quantum = 0;
let mut spilled = 0;
for (i, word) in bits.as_ref().iter().copied().enumerate() {
let ones_in_word = word.count_ones() as usize;
while past_ones + ones_in_word > next_quantum {
let in_word_index = word.select_in_word(next_quantum - past_ones);
let index = (i * usize::BITS as usize) + in_word_index;
inventory.push(index);
inventory.resize(inventory.len() + u64_per_subinventory, 0);
next_quantum += ones_per_inventory;
}
past_ones += ones_in_word;
}
assert_eq!(past_ones, num_ones);
inventory.push(num_bits);
assert_eq!(inventory.len(), inventory_words);
for (i, inv) in inventory[..inventory_size * u64_per_inventory]
.iter()
.copied()
.step_by(u64_per_inventory)
.enumerate()
{
let start = inv;
let span = inventory[i * u64_per_inventory + u64_per_inventory] - start;
past_ones = i * ones_per_inventory;
let ones = min(num_ones - past_ones, ones_per_inventory);
debug_assert!(start + span == num_bits || ones == ones_per_inventory);
match SpanType::from_span(span) {
SpanType::U32 => {
let log2_ones_per_sub32 = Self::log2_ones_per_sub32(span, log2_ones_per_sub16);
let num_u32s = ones.div_ceil(1 << log2_ones_per_sub32);
let num_u64s = num_u32s.div_ceil(2);
let spilled_u64s = num_u64s.saturating_sub(u64_per_subinventory - 1);
spilled += spilled_u64s;
}
SpanType::U64 => {
spilled += (ones - 1).saturating_sub(u64_per_subinventory - 1);
}
_ => {}
}
}
let spill_size = spilled;
let mut inventory: Box<[usize]> = inventory.into();
let mut spill: Box<[usize]> = vec![0; spill_size].into();
spilled = 0;
let locally_stored_u32s = 2 * (u64_per_subinventory - 1);
for inventory_idx in 0..inventory_size {
let start_inv_idx = inventory_idx * u64_per_inventory;
let end_inv_idx = start_inv_idx + u64_per_inventory;
let start_bit_idx = inventory[start_inv_idx];
let end_bit_idx = inventory[end_inv_idx];
let span = end_bit_idx - start_bit_idx;
let span_type = SpanType::from_span(span);
let mut past_ones = inventory_idx * ones_per_inventory;
let mut next_quantum = past_ones;
let log2_quantum;
match span_type {
SpanType::U16 => {
log2_quantum = log2_ones_per_sub16;
inventory[start_inv_idx].set_u16_span();
}
SpanType::U32 => {
log2_quantum = Self::log2_ones_per_sub32(span, log2_ones_per_sub16);
inventory[start_inv_idx].set_u32_span();
inventory[start_inv_idx + 1] = spilled;
}
SpanType::U64 => {
log2_quantum = 0;
inventory[start_inv_idx].set_u64_span();
inventory[start_inv_idx + 1] = spilled;
}
}
let quantum = 1 << log2_quantum;
let mut subinventory_idx = 1;
next_quantum += quantum;
let mut word_idx = start_bit_idx / usize::BITS as usize;
let end_word_idx = end_bit_idx.div_ceil(usize::BITS as usize);
let bit_idx = start_bit_idx % usize::BITS as usize;
let mut word = (bits.as_ref()[word_idx] >> bit_idx) << bit_idx;
'outer: loop {
let ones_in_word = word.count_ones() as usize;
while past_ones + ones_in_word > next_quantum {
debug_assert!(next_quantum <= end_bit_idx);
let in_word_index = word.select_in_word(next_quantum - past_ones);
let bit_index = (word_idx * usize::BITS as usize) + in_word_index;
if bit_index >= end_bit_idx {
break 'outer;
}
let sub_offset = bit_index - start_bit_idx;
match span_type {
SpanType::U16 => {
let subinventory: &mut [u16] = unsafe {
inventory[start_inv_idx + 1..end_inv_idx].align_to_mut().1
};
subinventory[subinventory_idx] = sub_offset as u16;
subinventory_idx += 1;
if subinventory_idx << log2_quantum == ones_per_inventory {
break 'outer;
}
}
SpanType::U32 => {
if subinventory_idx < locally_stored_u32s {
let subinventory: &mut [u32] = unsafe {
inventory[start_inv_idx + 2..end_inv_idx].align_to_mut().1
};
debug_assert_eq!(subinventory[subinventory_idx], 0);
subinventory[subinventory_idx] = sub_offset as u32;
} else {
let u32_spill: &mut [u32] =
unsafe { spill[spilled..].align_to_mut().1 };
debug_assert_eq!(
u32_spill[subinventory_idx - locally_stored_u32s],
0
);
u32_spill[subinventory_idx - locally_stored_u32s] =
sub_offset as u32;
}
subinventory_idx += 1;
if subinventory_idx << log2_quantum == ones_per_inventory {
break 'outer;
}
}
SpanType::U64 => {
if subinventory_idx < u64_per_subinventory {
inventory[start_inv_idx + 1 + subinventory_idx] = bit_index;
subinventory_idx += 1;
} else {
assert!(spilled < spill_size);
spill[spilled] = bit_index;
spilled += 1;
}
if subinventory_idx == ones_per_inventory {
break 'outer;
}
}
}
next_quantum += quantum;
}
past_ones += ones_in_word;
word_idx += 1;
if word_idx == end_word_idx {
break;
}
word = bits.as_ref()[word_idx];
}
if span_type == SpanType::U32 {
spilled += subinventory_idx
.saturating_sub(locally_stored_u32s)
.div_ceil(2);
}
}
assert_eq!(spilled, spill_size);
Self {
bits,
inventory,
spill,
log2_ones_per_inventory,
log2_ones_per_sub16,
log2_u64_per_subinventory,
ones_per_inventory_mask,
ones_per_sub16_mask,
}
}
}
impl<B: AsRef<[usize]> + BitLength + SelectHinted, I: AsRef<[usize]>> SelectUnchecked
for SelectAdapt<B, I>
{
unsafe fn select_unchecked(&self, rank: usize) -> usize {
unsafe {
let inventory = self.inventory.as_ref();
let inventory_index = rank >> self.log2_ones_per_inventory;
let inventory_start_pos =
(inventory_index << self.log2_u64_per_subinventory) + inventory_index;
let inventory_rank = { *inventory.get_unchecked(inventory_start_pos) };
let subrank = rank & self.ones_per_inventory_mask;
if inventory_rank.is_u16_span() {
let subinventory = inventory
.get_unchecked(inventory_start_pos + 1..)
.align_to::<u16>()
.1;
debug_assert!(subrank >> self.log2_ones_per_sub16 < subinventory.len());
let hint_pos = inventory_rank
+ *subinventory.get_unchecked(subrank >> self.log2_ones_per_sub16) as usize;
let residual = subrank & self.ones_per_sub16_mask;
return self.bits.select_hinted(rank, hint_pos, rank - residual);
}
let u64_per_subinventory = 1 << self.log2_u64_per_subinventory;
if inventory_rank.is_u32_span() {
let inventory_rank = inventory_rank.get();
let span = (*inventory
.get_unchecked(inventory_start_pos + u64_per_subinventory + 1))
.get()
- inventory_rank;
let log2_ones_per_sub32 = Self::log2_ones_per_sub32(span, self.log2_ones_per_sub16);
let hint_pos = if subrank >> log2_ones_per_sub32 < (u64_per_subinventory - 1) * 2 {
let u32s = inventory
.get_unchecked(inventory_start_pos + 2..)
.align_to::<u32>()
.1;
inventory_rank + *u32s.get_unchecked(subrank >> log2_ones_per_sub32) as usize
} else {
let start_spill_idx = *inventory.get_unchecked(inventory_start_pos + 1);
let spilled_u32s = self
.spill
.as_ref()
.get_unchecked(start_spill_idx..)
.align_to::<u32>()
.1;
inventory_rank
+ *spilled_u32s.get_unchecked(
(subrank >> log2_ones_per_sub32) - (u64_per_subinventory - 1) * 2,
) as usize
};
let residual = subrank & ((1 << log2_ones_per_sub32) - 1);
return self.bits.select_hinted(rank, hint_pos, rank - residual);
}
debug_assert!(inventory_rank.is_u64_span());
let inventory_rank = inventory_rank.get();
if subrank < u64_per_subinventory {
if subrank == 0 {
return inventory_rank;
}
return *inventory.get_unchecked(inventory_start_pos + 1 + subrank);
}
let spill_idx = { *inventory.get_unchecked(inventory_start_pos + 1) } + subrank
- u64_per_subinventory;
debug_assert!(spill_idx < self.spill.as_ref().len());
*self.spill.as_ref().get_unchecked(spill_idx)
}
}
}
impl<B: SelectHinted + AsRef<[usize]> + NumBits, I: AsRef<[usize]>> Select for SelectAdapt<B, I> {}
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use super::*;
use crate::bits::BitVec;
use crate::traits::AddNumBits;
use crate::traits::BitVecOpsMut;
use rand::rngs::SmallRng;
use rand::{RngExt, SeedableRng};
#[test]
fn test_sub64s() {
let len = 5_000_000_000;
let mut rng = SmallRng::seed_from_u64(0);
let mut bits = BitVec::new(len);
let mut pos = BTreeSet::new();
for _ in 0..(1 << 13) / 4 * 3 {
let p = rng.random_range(0..len);
if pos.insert(p) {
bits.set(p, true);
}
}
let bits: AddNumBits<BitVec> = bits.into();
for m in [0, 3, 16] {
let simple = SelectAdapt::with_inv(&bits, 13, m);
assert!(simple.inventory[0].is_u64_span());
for (i, &p) in pos.iter().enumerate() {
assert_eq!(simple.select(i), Some(p));
}
assert_eq!(simple.select(pos.len()), None);
}
}
}