/*
 *
 * SPDX-FileCopyrightText: 2024 Michele Andreata
 * SPDX-FileCopyrightText: 2024 Sebastiano Vigna
 *
 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
 */

use super::SmallCounters;
use crate::prelude::*;
use crate::traits::{Backend, Word};
use crate::utils::SelectInWord;
use ambassador::Delegate;
use mem_dbg::{MemDbg, MemSize};
use num_primitive::PrimitiveInteger;

use crate::ambassador_impl_Index;
use crate::rank_sel::ambassador_impl_SmallCounters;
use crate::traits::ambassador_impl_Backend;
use crate::traits::bal_paren::ambassador_impl_BalParen;
use crate::traits::bit_vec_ops::ambassador_impl_BitLength;
use crate::traits::rank_sel::ambassador_impl_BitCount;
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_Select;
use crate::traits::rank_sel::ambassador_impl_SelectHinted;
use crate::traits::rank_sel::ambassador_impl_SelectUnchecked;
use crate::traits::rank_sel::ambassador_impl_SelectZeroHinted;
use std::ops::Deref;
use std::ops::Index;

// NOTE: to make parallel modifications with SelectSmall as easy as possible,
// "ones" are considered to be zeros in the following code.

/// A version of [`SelectSmall`] implementing [selection on zeros].
///
/// [selection on zeros]: crate::traits::SelectZero
///
/// # Examples
///
/// ```rust
/// # #[cfg(target_pointer_width = "64")]
/// # {
/// # use sux::{bit_vec, rank_small};
/// # use sux::rank_sel::{SelectSmall, SelectZeroSmall};
/// # use sux::traits::{Select, SelectZero};
/// let bits = bit_vec![0, 1, 0, 0, 1, 0, 1, 0];
/// let rank_small = rank_small![u64: 1; bits];
/// let sel = SelectZeroSmall::<1, 9, _>::new(rank_small);
///
/// assert_eq!(sel.select_zero(0), Some(0));
/// assert_eq!(sel.select_zero(1), Some(2));
/// assert_eq!(sel.select_zero(2), Some(3));
/// assert_eq!(sel.select_zero(3), Some(5));
/// assert_eq!(sel.select_zero(4), Some(7));
/// assert_eq!(sel.select_zero(5), None);
///
/// // One can also stack a SelectZeroSmall on top of a SelectSmall:
/// let rank_small = sel.into_inner();
/// let sel = SelectZeroSmall::<1, 9, _>::new(SelectSmall::<1, 9, _>::new(rank_small));
///
/// assert_eq!(sel.select_zero(0), Some(0));
/// assert_eq!(sel.select_zero(1), Some(2));
/// assert_eq!(sel.select_zero(2), Some(3));
/// assert_eq!(sel.select_zero(3), Some(5));
/// assert_eq!(sel.select_zero(4), Some(7));
/// assert_eq!(sel.select_zero(5), None);
///
/// assert_eq!(sel.select(0), Some(1));
/// assert_eq!(sel.select(1), Some(4));
/// assert_eq!(sel.select(2), Some(6));
/// assert_eq!(sel.select(3), None);
/// # }
/// ```
#[derive(Debug, Clone, MemSize, MemDbg, Delegate)]
#[cfg_attr(feature = "epserde", derive(epserde::Epserde))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[delegate(crate::traits::Backend, target = "small_counters")]
#[delegate(Index<usize>, target = "small_counters")]
#[delegate(crate::traits::rank_sel::BitCount, target = "small_counters")]
#[delegate(crate::traits::bit_vec_ops::BitLength, target = "small_counters")]
#[delegate(crate::traits::rank_sel::NumBits, target = "small_counters")]
#[delegate(crate::traits::rank_sel::Rank, target = "small_counters")]
#[delegate(crate::traits::rank_sel::RankHinted, target = "small_counters")]
#[delegate(crate::traits::rank_sel::RankUnchecked, target = "small_counters")]
#[delegate(crate::traits::rank_sel::RankZero, target = "small_counters")]
#[delegate(crate::traits::rank_sel::Select, target = "small_counters")]
#[delegate(crate::traits::rank_sel::SelectHinted, target = "small_counters")]
#[delegate(crate::traits::rank_sel::SelectUnchecked, target = "small_counters")]
#[delegate(crate::traits::rank_sel::SelectZeroHinted, target = "small_counters")]
#[delegate(crate::bal_paren::BalParen, target = "small_counters")]
#[delegate(crate::rank_sel::SmallCounters<NUM_U32S, COUNTER_WIDTH>, target = "small_counters")]
pub struct SelectZeroSmall<
    const NUM_U32S: usize,
    const COUNTER_WIDTH: usize,
    C,
    I = Box<[u32]>,
    O = Box<[usize]>,
> {
    small_counters: C,
    inventory: I,
    inventory_begin: O,
    log2_ones_per_inventory: usize,
}

impl<C: Backend + AsRef<[C::Word]>, const NUM_U32S: usize, const COUNTER_WIDTH: usize, I, O>
    AsRef<[C::Word]> for SelectZeroSmall<NUM_U32S, COUNTER_WIDTH, C, I, O>
{
    #[inline(always)]
    fn as_ref(&self) -> &[C::Word] {
        self.small_counters.as_ref()
    }
}

impl<const NUM_U32S: usize, const COUNTER_WIDTH: usize, C, I, O> Deref
    for SelectZeroSmall<NUM_U32S, COUNTER_WIDTH, C, I, O>
{
    type Target = C;

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        &self.small_counters
    }
}

impl<const NUM_U32S: usize, const COUNTER_WIDTH: usize, C, I, O>
    SelectZeroSmall<NUM_U32S, COUNTER_WIDTH, C, I, O>
{
    // On 64-bit, superblocks contain 2^32 bits.
    // On 32-bit, a single superblock covers the entire address space.
    #[cfg(target_pointer_width = "64")]
    const SUPERBLOCK_BIT_SIZE: usize = 1 << 32;
    #[cfg(not(target_pointer_width = "64"))]
    const SUPERBLOCK_BIT_SIZE: usize = usize::MAX;
    const BLOCK_BIT_SIZE: usize = 1 << COUNTER_WIDTH;
    /// Bits in each inventory `u32` used for the block index within the superblock.
    /// The remaining high `COUNTER_WIDTH` bits store the midpoint-block delta.
    const BLOCK_IDX_BITS: usize = 32 - COUNTER_WIDTH;
    const BLOCK_IDX_MASK: usize = (1 << Self::BLOCK_IDX_BITS) - 1;

    /// Returns the underlying rank-small structure, consuming this structure.
    pub fn into_inner(self) -> C {
        self.small_counters
    }
}

impl<const NUM_U32S: usize, const COUNTER_WIDTH: usize, C: BitLength, I, O>
    SelectZeroSmall<NUM_U32S, COUNTER_WIDTH, C, I, O>
{
    /// Returns the number of bits in the underlying bit vector.
    ///
    /// This method is equivalent to [`BitLength::len`], but it is provided to
    /// reduce ambiguity in method resolution.
    #[inline(always)]
    pub fn len(&self) -> usize {
        self.small_counters.len()
    }
}

macro_rules! impl_select_zero_small {
    ($NUM_U32S: literal; $COUNTER_WIDTH: literal) => {
        impl<
            C: SmallCounters<$NUM_U32S, $COUNTER_WIDTH>
                + Backend<Word: Word + SelectInWord>
                + AsRef<[C::Word]>
                + BitLength
                + NumBits
                + SelectZeroHinted,
        > SelectZeroSmall<$NUM_U32S, $COUNTER_WIDTH, C>
        {
            /// Creates a new selection structure with four [`RankSmall`]
            /// blocks per inventory on average.
            #[must_use]
            pub fn new(small_counters: C) -> Self {
                Self::with_inv(small_counters, 4)
            }

            /// Creates a new selection structure with a given number of
            /// [`RankSmall`] blocks per inventory on average.
            #[must_use]
            pub fn with_inv(small_counters: C, blocks_per_inv: usize) -> Self {
                let num_bits = small_counters.len();
                let num_ones = small_counters.len() - small_counters.num_ones();

                let target_inventory_span = blocks_per_inv * Self::BLOCK_BIT_SIZE;
                let log2_ones_per_inventory = (num_ones as u64 * target_inventory_span as u64)
                    .div_ceil(num_bits.max(1) as u64)
                    .max(1)
                    .ilog2() as usize;

                Self::_new(small_counters, num_ones, log2_ones_per_inventory)
            }

            fn _new(small_counters: C, num_ones: usize, log2_ones_per_inventory: usize) -> Self {
                let ones_per_inventory = 1 << log2_ones_per_inventory;
                let half_ones = ones_per_inventory >> 1;
                let step = half_ones.max(1);

                let inventory_size = num_ones.div_ceil(ones_per_inventory);
                let mut inventory = Vec::<u32>::with_capacity(inventory_size + 1);

                // The inventory_begin vector stores the index of the first element of
                // each superblock in the inventory.
                let mut inventory_begin =
                    Vec::<usize>::with_capacity(small_counters.upper_counts().len() + 1);

                let mut past_ones: usize = 0;
                let mut next_quantum: usize = 0;

                let bits_per_word = C::Word::BITS as usize;
                for superblock in small_counters
                    .as_ref()
                    .chunks(Self::SUPERBLOCK_BIT_SIZE / bits_per_word)
                {
                    let mut first = true;
                    for (i, word) in superblock.iter().copied().map(|b| !b).enumerate() {
                        let ones_in_word = (word.count_ones() as usize).min(num_ones - past_ones);

                        while past_ones + ones_in_word > next_quantum {
                            let in_word_index = word.select_in_word(next_quantum - past_ones);
                            let in_superblock_index = i * bits_per_word + in_word_index;
                            let block_in_superblock = in_superblock_index / Self::BLOCK_BIT_SIZE;

                            if half_ones == 0 || next_quantum & half_ones == 0 {
                                if first {
                                    inventory_begin.push(inventory.len() as usize);
                                    first = false;
                                }
                                inventory.push(block_in_superblock as u32);
                            } else {
                                let last = inventory.last_mut().unwrap();
                                let delta =
                                    block_in_superblock - (*last as usize & Self::BLOCK_IDX_MASK);
                                debug_assert!(
                                    delta < Self::BLOCK_BIT_SIZE,
                                    "midpoint delta {delta} overflows COUNTER_WIDTH bits"
                                );
                                *last |= (delta as u32) << Self::BLOCK_IDX_BITS;
                            }
                            next_quantum += step;
                        }

                        past_ones += ones_in_word;
                    }
                }
                assert_eq!(num_ones, past_ones);

                if inventory.is_empty() {
                    inventory.push(0);
                    inventory_begin.push(0 as usize);
                } else {
                    inventory_begin.push(small_counters.as_ref().len() as usize);
                }

                let inventory = inventory.into_boxed_slice();
                let inventory_begin = inventory_begin.into_boxed_slice();

                Self {
                    small_counters,
                    inventory,
                    inventory_begin,
                    log2_ones_per_inventory,
                }
            }
        }

        impl<
            C: SmallCounters<$NUM_U32S, $COUNTER_WIDTH>
                + Backend<Word: Word + SelectInWord>
                + AsRef<[C::Word]>
                + BitLength
                + NumBits
                + SelectZeroHinted,
        > SelectZeroUnchecked for SelectZeroSmall<$NUM_U32S, $COUNTER_WIDTH, C>
        {
            /// # Safety
            ///
            /// `rank` must be between zero (included) and the number of zeros in the
            /// underlying bit vector (excluded).
            unsafe fn select_zero_unchecked(&self, rank: usize) -> usize {
                let upper_counts = self.small_counters.upper_counts();
                let counts = self.small_counters.counts();

                let upper_block_idx = upper_counts.linear_partition_point(|i, &x| {
                    i * Self::SUPERBLOCK_BIT_SIZE - x as usize <= rank
                }) - 1;
                let upper_rank_ones =
                    *unsafe { upper_counts.get_unchecked(upper_block_idx) } as usize;

                let inventory = self.inventory.as_ref();
                let inventory_begin = self.inventory_begin.as_ref();

                // we find the two inventory entries containing the rank and from that
                // the blocks containing the rank.
                // if the span of the two entries is not contained in the same upper block
                // we clip the span to the upper block boundaries since we know that
                // the rank is in the upper block with index upper_block_idx.

                let inv_idx = rank >> self.log2_ones_per_inventory;
                let inv_upper_block_idx =
                    inventory_begin.linear_partition_point(|_, &x| x as usize <= inv_idx) - 1;
                let half_ones = 1usize << self.log2_ones_per_inventory >> 1;
                // Branchless: compute a mask that is all-ones in the second half
                // and all-zeros in the first half, avoiding a ~50/50 branch that
                // would cause systematic mispredictions.
                let second_half_mask = ((rank & half_ones != 0) as usize).wrapping_neg();
                let mut block_idx = if inv_upper_block_idx == upper_block_idx {
                    let inv_entry = *unsafe { inventory.get_unchecked(inv_idx) } as usize;
                    let base_block = inv_entry & Self::BLOCK_IDX_MASK;
                    let mid_delta = inv_entry >> Self::BLOCK_IDX_BITS;
                    base_block + (mid_delta & second_half_mask)
                } else {
                    0
                } + upper_block_idx
                    * (Self::SUPERBLOCK_BIT_SIZE / Self::BLOCK_BIT_SIZE);

                // Prefetch all subblocks of the approximate target block now,
                // so the bit-vector DRAM fetch can proceed in parallel with
                // the upcoming counts DRAM fetch (same idea as rank_unchecked).
                {
                    const NUM_SUBBLOCKS: usize = match $NUM_U32S {
                        1 => 4,
                        _ => 8,
                    };
                    let bits_ref: &[C::Word] = self.as_ref();
                    let words_per_subblock =
                        Self::BLOCK_BIT_SIZE / (NUM_SUBBLOCKS * C::Word::BITS as usize);
                    let base_word = block_idx * (Self::BLOCK_BIT_SIZE / C::Word::BITS as usize);
                    for k in 0..NUM_SUBBLOCKS {
                        crate::utils::prefetch_index(bits_ref, base_word + k * words_per_subblock);
                    }
                }

                let last_block_idx;
                if inv_idx + 1 < inventory.len() {
                    // linear_partition_point returns the first index where
                    // the predicate is false; subtracting 1 gives the last
                    // superblock whose inventory start is <= inv_idx + 1,
                    // i.e., the superblock containing inventory[inv_idx + 1].
                    let next_inv_upper_block_idx = inventory_begin
                        .linear_partition_point(|_, &x| x as usize <= inv_idx + 1)
                        - 1;
                    last_block_idx = if next_inv_upper_block_idx == upper_block_idx {
                        let next_base_block = (*unsafe { inventory.get_unchecked(inv_idx + 1) }
                            as usize)
                            & Self::BLOCK_IDX_MASK;
                        next_base_block
                            + 1
                            + upper_block_idx * (Self::SUPERBLOCK_BIT_SIZE / Self::BLOCK_BIT_SIZE)
                    } else {
                        (upper_block_idx + 1) * (Self::SUPERBLOCK_BIT_SIZE / Self::BLOCK_BIT_SIZE)
                    };
                } else {
                    // Since we use 32-bit inventory entries, we cannot add
                    // a sentinel with value equal to the number of bits
                    // (which may exceed 2^32). Thus, we handle the last
                    // inventory entry as a special case.
                    last_block_idx = self.len().div_ceil(Self::BLOCK_BIT_SIZE);
                }

                debug_assert!(block_idx < counts.len());

                debug_assert!(
                    block_idx <= last_block_idx,
                    "{}, {}",
                    block_idx,
                    last_block_idx
                );

                debug_assert!(block_idx < last_block_idx);

                block_idx += counts[block_idx..last_block_idx].linear_partition_point(|i, x| {
                    ((block_idx + i) * Self::BLOCK_BIT_SIZE)
                        - (upper_rank_ones + x.absolute as usize)
                        <= rank
                }) - 1;

                let block_count = unsafe { counts.get_unchecked(block_idx) };
                let hint_pos = block_idx * Self::BLOCK_BIT_SIZE;
                let hint_rank = hint_pos - (upper_rank_ones + block_count.absolute as usize);

                unsafe { self.complete_select(block_count, hint_pos, rank, hint_rank) }
            }
        }

        impl<
            C: SmallCounters<$NUM_U32S, $COUNTER_WIDTH>
                + Backend<Word: Word + SelectInWord>
                + AsRef<[C::Word]>
                + BitLength
                + NumBits
                + SelectZeroHinted,
        > SelectZero for SelectZeroSmall<$NUM_U32S, $COUNTER_WIDTH, C>
        {
        }
    };
}

impl<
    C: SmallCounters<2, 9>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits,
> SelectZeroSmall<2, 9, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<2, 9>,
        mut hint_pos: usize,
        rank: usize,
        hint_rank: usize,
    ) -> usize {
        const ONES_STEP_9: u64 = (1_u64 << 0)
            | (1_u64 << 9)
            | (1_u64 << 18)
            | (1_u64 << 27)
            | (1_u64 << 36)
            | (1_u64 << 45)
            | (1_u64 << 54);
        const MSBS_STEP_9: u64 = 0x100_u64 * ONES_STEP_9;
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 9) / 8;
        const POS_STEP_9: u64 = (SUBBLOCK_BIT_SIZE << (6 * 9))
            | ((2 * SUBBLOCK_BIT_SIZE) << (5 * 9))
            | ((3 * SUBBLOCK_BIT_SIZE) << (4 * 9))
            | ((4 * SUBBLOCK_BIT_SIZE) << (3 * 9))
            | ((5 * SUBBLOCK_BIT_SIZE) << (2 * 9))
            | ((6 * SUBBLOCK_BIT_SIZE) << 9)
            | (7 * SUBBLOCK_BIT_SIZE);

        macro_rules! ULEQ_STEP_9 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_9) - (($x) & !MSBS_STEP_9)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_9)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_9 = rank_in_block as u64 * ONES_STEP_9;
        let relative = POS_STEP_9 - block_count.all_rel();
        let offset_in_block = ULEQ_STEP_9!(relative, rank_in_block_step_9).count_ones() as usize;

        let rank_in_word = rank_in_block
            - (offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block));
        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);

        hint_pos
            + (!*unsafe {
                self.small_counters
                    .as_ref()
                    .get_unchecked(hint_pos / C::Word::BITS as usize)
            })
            .select_in_word(rank_in_word)
    }
}

impl<
    C: SmallCounters<1, 9>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits
        + SelectZeroHinted,
> SelectZeroSmall<1, 9, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<1, 9>,
        mut hint_pos: usize,
        rank: usize,
        mut hint_rank: usize,
    ) -> usize {
        const ONES_STEP_9: u64 = (1_u64 << 0) | (1_u64 << 9) | (1_u64 << 18);
        const MSBS_STEP_9: u64 = 0x100_u64 * ONES_STEP_9;
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 9) / 4;
        // We cannot put this const together with the rest because we need
        // to use it for defining POS_STEP_9.
        const POS_STEP_9: u64 =
            (SUBBLOCK_BIT_SIZE << 18) | ((2 * SUBBLOCK_BIT_SIZE) << 9) | (3 * SUBBLOCK_BIT_SIZE);

        macro_rules! ULEQ_STEP_9 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_9) - (($x) & !MSBS_STEP_9)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_9)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_9 = rank_in_block as u64 * ONES_STEP_9;
        let relative = POS_STEP_9 - block_count.all_rel();

        let offset_in_block = ULEQ_STEP_9!(relative, rank_in_block_step_9).count_ones() as usize;

        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);
        hint_rank +=
            offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block);

        let subblock_words = (SUBBLOCK_BIT_SIZE as usize) / C::Word::BITS as usize;
        unsafe {
            match subblock_words {
                1 => self.select_zero_hinted::<1>(rank, hint_pos, hint_rank),
                2 => self.select_zero_hinted::<2>(rank, hint_pos, hint_rank),
                4 => self.select_zero_hinted::<4>(rank, hint_pos, hint_rank),
                8 => self.select_zero_hinted::<8>(rank, hint_pos, hint_rank),
                16 => self.select_zero_hinted::<16>(rank, hint_pos, hint_rank),
                32 => self.select_zero_hinted::<32>(rank, hint_pos, hint_rank),
                _ => self.select_zero_hinted::<{ usize::MAX }>(rank, hint_pos, hint_rank),
            }
        }
    }
}

impl<
    C: SmallCounters<1, 10>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits
        + SelectZeroHinted,
> SelectZeroSmall<1, 10, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<1, 10>,
        mut hint_pos: usize,
        rank: usize,
        mut hint_rank: usize,
    ) -> usize {
        const ONES_STEP_10: u64 = (1_u64 << 0) | (1_u64 << 10) | (1_u64 << 20);
        const MSBS_STEP_10: u64 = 0x200_u64 * ONES_STEP_10;
        // We cannot put this const together with the rest because we need
        // to use it for defining POS_STEP_10.
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 10) / 4;
        const POS_STEP_10: u64 =
            (SUBBLOCK_BIT_SIZE << 20) | ((2 * SUBBLOCK_BIT_SIZE) << 10) | (3 * SUBBLOCK_BIT_SIZE);

        macro_rules! ULEQ_STEP_10 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_10) - (($x) & !MSBS_STEP_10)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_10)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_10 = rank_in_block as u64 * ONES_STEP_10;
        let relative = POS_STEP_10 - block_count.all_rel();

        let offset_in_block = ULEQ_STEP_10!(relative, rank_in_block_step_10).count_ones() as usize;

        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);
        hint_rank +=
            offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block);

        let subblock_words = (SUBBLOCK_BIT_SIZE as usize) / C::Word::BITS as usize;
        unsafe {
            match subblock_words {
                1 => self.select_zero_hinted::<1>(rank, hint_pos, hint_rank),
                2 => self.select_zero_hinted::<2>(rank, hint_pos, hint_rank),
                4 => self.select_zero_hinted::<4>(rank, hint_pos, hint_rank),
                8 => self.select_zero_hinted::<8>(rank, hint_pos, hint_rank),
                16 => self.select_zero_hinted::<16>(rank, hint_pos, hint_rank),
                32 => self.select_zero_hinted::<32>(rank, hint_pos, hint_rank),
                _ => self.select_zero_hinted::<{ usize::MAX }>(rank, hint_pos, hint_rank),
            }
        }
    }
}

impl<
    C: SmallCounters<1, 11>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits
        + SelectZeroHinted,
> SelectZeroSmall<1, 11, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<1, 11>,
        mut hint_pos: usize,
        rank: usize,
        mut hint_rank: usize,
    ) -> usize {
        const ONES_STEP_11: u64 = (1_u64 << 0) | (1_u64 << 11) | (1_u64 << 22);
        const MSBS_STEP_11: u64 = 0x400_u64 * ONES_STEP_11;
        // We cannot put this const together with the rest because we need
        // to use it for defining POS_STEP_11.
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 11) / 4;
        const POS_STEP_11: u64 =
            (SUBBLOCK_BIT_SIZE << 22) | ((2 * SUBBLOCK_BIT_SIZE) << 11) | (3 * SUBBLOCK_BIT_SIZE);

        macro_rules! ULEQ_STEP_11 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_11) - (($x) & !MSBS_STEP_11)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_11)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_11 = rank_in_block as u64 * ONES_STEP_11;
        let relative = POS_STEP_11 - block_count.all_rel();

        let offset_in_block = ULEQ_STEP_11!(relative, rank_in_block_step_11).count_ones() as usize;

        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);
        hint_rank +=
            offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block);

        let subblock_words = (SUBBLOCK_BIT_SIZE as usize) / C::Word::BITS as usize;
        unsafe {
            match subblock_words {
                1 => self.select_zero_hinted::<1>(rank, hint_pos, hint_rank),
                2 => self.select_zero_hinted::<2>(rank, hint_pos, hint_rank),
                4 => self.select_zero_hinted::<4>(rank, hint_pos, hint_rank),
                8 => self.select_zero_hinted::<8>(rank, hint_pos, hint_rank),
                16 => self.select_zero_hinted::<16>(rank, hint_pos, hint_rank),
                32 => self.select_zero_hinted::<32>(rank, hint_pos, hint_rank),
                _ => self.select_zero_hinted::<{ usize::MAX }>(rank, hint_pos, hint_rank),
            }
        }
    }
}

impl<
    C: SmallCounters<3, 13>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits
        + SelectZeroHinted,
> SelectZeroSmall<3, 13, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<3, 13>,
        mut hint_pos: usize,
        rank: usize,
        mut hint_rank: usize,
    ) -> usize {
        const ONES_STEP_13: u128 = (1_u128 << 0)
            | (1_u128 << 13)
            | (1_u128 << 26)
            | (1_u128 << 39)
            | (1_u128 << 52)
            | (1_u128 << 65)
            | (1_u128 << 78);
        const MSBS_STEP_13: u128 = 0x1000_u128 * ONES_STEP_13;
        // We cannot put this const together with the rest because we need
        // to use it for defining POS_STEP_13.
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 13) / 8;
        const POS_STEP_13: u128 = ((SUBBLOCK_BIT_SIZE as u128) << 78)
            | ((2 * (SUBBLOCK_BIT_SIZE as u128)) << 65)
            | ((3 * (SUBBLOCK_BIT_SIZE as u128)) << 52)
            | ((4 * (SUBBLOCK_BIT_SIZE as u128)) << 39)
            | ((5 * (SUBBLOCK_BIT_SIZE as u128)) << 26)
            | ((6 * (SUBBLOCK_BIT_SIZE as u128)) << 13)
            | (7 * (SUBBLOCK_BIT_SIZE as u128));

        macro_rules! ULEQ_STEP_13 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_13) - (($x) & !MSBS_STEP_13)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_13)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_13 = rank_in_block as u128 * ONES_STEP_13;
        let relative = POS_STEP_13 - block_count.all_rel();

        let offset_in_block = ULEQ_STEP_13!(relative, rank_in_block_step_13).count_ones() as usize;

        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);
        hint_rank +=
            offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block);

        let subblock_words = (SUBBLOCK_BIT_SIZE as usize) / C::Word::BITS as usize;
        unsafe {
            match subblock_words {
                1 => self.select_zero_hinted::<1>(rank, hint_pos, hint_rank),
                2 => self.select_zero_hinted::<2>(rank, hint_pos, hint_rank),
                4 => self.select_zero_hinted::<4>(rank, hint_pos, hint_rank),
                8 => self.select_zero_hinted::<8>(rank, hint_pos, hint_rank),
                16 => self.select_zero_hinted::<16>(rank, hint_pos, hint_rank),
                32 => self.select_zero_hinted::<32>(rank, hint_pos, hint_rank),
                _ => self.select_zero_hinted::<{ usize::MAX }>(rank, hint_pos, hint_rank),
            }
        }
    }
}

// 64-bit word variants
// One invocation per unique (NUM_U32S, COUNTER_WIDTH) pair.
// The impl is generic over C, so it covers both u32 and u64 backends.
impl_select_zero_small!(2; 9);
impl_select_zero_small!(1; 9);
impl_select_zero_small!(1; 10);
impl_select_zero_small!(1; 11);
impl_select_zero_small!(3; 13);

impl<
    C: SmallCounters<2, 8>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits,
> SelectZeroSmall<2, 8, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<2, 8>,
        mut hint_pos: usize,
        rank: usize,
        hint_rank: usize,
    ) -> usize {
        const ONES_STEP_8: u64 = (1_u64 << 0)
            | (1_u64 << 8)
            | (1_u64 << 16)
            | (1_u64 << 24)
            | (1_u64 << 32)
            | (1_u64 << 40)
            | (1_u64 << 48);
        const MSBS_STEP_8: u64 = 0x80_u64 * ONES_STEP_8;
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 8) / 8;
        const POS_STEP_8: u64 = (SUBBLOCK_BIT_SIZE << (6 * 8))
            | ((2 * SUBBLOCK_BIT_SIZE) << (5 * 8))
            | ((3 * SUBBLOCK_BIT_SIZE) << (4 * 8))
            | ((4 * SUBBLOCK_BIT_SIZE) << (3 * 8))
            | ((5 * SUBBLOCK_BIT_SIZE) << (2 * 8))
            | ((6 * SUBBLOCK_BIT_SIZE) << 8)
            | (7 * SUBBLOCK_BIT_SIZE);

        macro_rules! ULEQ_STEP_8 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_8) - (($x) & !MSBS_STEP_8)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_8)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_8 = rank_in_block as u64 * ONES_STEP_8;
        let relative = POS_STEP_8 - block_count.all_rel();
        let offset_in_block = ULEQ_STEP_8!(relative, rank_in_block_step_8).count_ones() as usize;

        let rank_in_word = rank_in_block
            - (offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block));
        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);

        hint_pos
            + (!*unsafe {
                self.small_counters
                    .as_ref()
                    .get_unchecked(hint_pos / C::Word::BITS as usize)
            })
            .select_in_word(rank_in_word)
    }
}

impl<
    C: SmallCounters<1, 8>
        + Backend<Word: Word + SelectInWord>
        + AsRef<[C::Word]>
        + BitLength
        + NumBits
        + SelectZeroHinted,
> SelectZeroSmall<1, 8, C>
{
    #[inline(always)]
    unsafe fn complete_select(
        &self,
        block_count: &Block32Counters<1, 8>,
        mut hint_pos: usize,
        rank: usize,
        mut hint_rank: usize,
    ) -> usize {
        const ONES_STEP_8: u64 = (1_u64 << 0) | (1_u64 << 8) | (1_u64 << 16);
        const MSBS_STEP_8: u64 = 0x80_u64 * ONES_STEP_8;
        const SUBBLOCK_BIT_SIZE: u64 = (1u64 << 8) / 4;
        const POS_STEP_8: u64 =
            (SUBBLOCK_BIT_SIZE << 16) | ((2 * SUBBLOCK_BIT_SIZE) << 8) | (3 * SUBBLOCK_BIT_SIZE);

        macro_rules! ULEQ_STEP_8 {
            ($x:ident, $y:ident) => {
                (((((($y) | MSBS_STEP_8) - (($x) & !MSBS_STEP_8)) | ($x ^ $y)) ^ ($x & !$y))
                    & MSBS_STEP_8)
            };
        }

        let rank_in_block = rank - hint_rank;
        let rank_in_block_step_8 = rank_in_block as u64 * ONES_STEP_8;
        let relative = POS_STEP_8 - block_count.all_rel();

        let offset_in_block = ULEQ_STEP_8!(relative, rank_in_block_step_8).count_ones() as usize;

        hint_pos += offset_in_block * (SUBBLOCK_BIT_SIZE as usize);
        hint_rank +=
            offset_in_block * (SUBBLOCK_BIT_SIZE as usize) - block_count.rel(offset_in_block);

        let subblock_words = (SUBBLOCK_BIT_SIZE as usize) / C::Word::BITS as usize;
        unsafe {
            match subblock_words {
                1 => self.select_zero_hinted::<1>(rank, hint_pos, hint_rank),
                2 => self.select_zero_hinted::<2>(rank, hint_pos, hint_rank),
                4 => self.select_zero_hinted::<4>(rank, hint_pos, hint_rank),
                8 => self.select_zero_hinted::<8>(rank, hint_pos, hint_rank),
                16 => self.select_zero_hinted::<16>(rank, hint_pos, hint_rank),
                32 => self.select_zero_hinted::<32>(rank, hint_pos, hint_rank),
                _ => self.select_zero_hinted::<{ usize::MAX }>(rank, hint_pos, hint_rank),
            }
        }
    }
}

// 32-bit word variants
impl_select_zero_small!(2; 8);
impl_select_zero_small!(1; 8);

/// A trait providing the semantics of [`partition_point`], but using a
/// linear search.
///
/// [`partition_point`]: slice::partition_point
trait LinearPartitionPointWithIndexExt<T>: AsRef<[T]> {
    fn linear_partition_point<P>(&self, mut pred: P) -> usize
    where
        P: FnMut(usize, &T) -> bool,
    {
        let as_ref = self.as_ref();
        as_ref
            .iter()
            .enumerate()
            .position(|(i, x)| !pred(i, x))
            .unwrap_or(as_ref.len())
    }
}

impl<T> LinearPartitionPointWithIndexExt<T> for [T] {}
