use alloc::vec::Vec;
#[derive(Debug)]
pub struct PackedArray<const BITS: u32> {
words: Vec<usize>,
len: usize,
}
impl<const BITS: u32> PackedArray<BITS> {
#[inline(always)]
fn items_per_word() -> usize {
(usize::BITS / BITS) as usize
}
#[inline(always)]
fn mask() -> usize {
(1usize << BITS) - 1
}
#[inline(always)]
fn word_of(index: usize) -> usize {
index / Self::items_per_word()
}
#[inline(always)]
fn offset_of(index: usize) -> usize {
(index % Self::items_per_word()) * BITS as usize
}
#[inline(always)]
fn words_for(len: usize) -> usize {
if len == 0 {
0
} else {
Self::word_of(len - 1) + 1
}
}
const WIDTH_OK: () = assert!(
BITS == 1 || BITS == 2 || BITS == 4,
"PackedArray BITS must be 1, 2 or 4"
);
#[inline]
pub fn new() -> Self {
let () = Self::WIDTH_OK;
Self {
words: Vec::new(),
len: 0,
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
let () = Self::WIDTH_OK;
Self {
words: Vec::with_capacity(Self::words_for(capacity)),
len: 0,
}
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn capacity(&self) -> usize {
self.words.capacity().saturating_mul(Self::items_per_word())
}
pub fn reserve(&mut self, additional: usize) {
let target_words = Self::words_for(self.len.saturating_add(additional));
if target_words > self.words.capacity() {
self.words.reserve(target_words - self.words.len());
}
}
pub fn reserve_exact(&mut self, additional: usize) {
let target_words = Self::words_for(self.len.saturating_add(additional));
if target_words > self.words.capacity() {
self.words.reserve_exact(target_words - self.words.len());
}
}
pub fn shrink_to_fit(&mut self) {
self.words.truncate(Self::words_for(self.len));
self.words.shrink_to_fit();
}
#[inline]
pub unsafe fn set_len(&mut self, new_len: usize) {
debug_assert!(Self::words_for(new_len) <= self.words.len());
self.len = new_len;
}
#[inline]
pub fn get(&self, index: usize) -> usize {
assert!(
index < self.len,
"index out of bounds: the len is {} but the index is {}",
self.len,
index
);
unsafe { self.get_unchecked(index) }
}
#[inline]
pub fn set(&mut self, index: usize, value: usize) {
assert!(
index < self.len,
"index out of bounds: the len is {} but the index is {}",
self.len,
index
);
unsafe { self.set_unchecked(index, value) }
}
#[inline(always)]
pub unsafe fn get_unchecked(&self, index: usize) -> usize {
debug_assert!(
Self::word_of(index) < self.words.len(),
"lane {index} has no backing word"
);
let word = unsafe { *self.words.get_unchecked(Self::word_of(index)) };
(word >> Self::offset_of(index)) & Self::mask()
}
#[inline(always)]
pub unsafe fn set_unchecked(&mut self, index: usize, value: usize) {
debug_assert!(
Self::word_of(index) < self.words.len(),
"lane {index} has no backing word"
);
let off = Self::offset_of(index);
let slot =
unsafe { self.words.get_unchecked_mut(Self::word_of(index)) };
*slot &= !(Self::mask() << off);
*slot |= (value & Self::mask()) << off;
}
#[inline]
pub fn push(&mut self, value: usize) {
let off = Self::offset_of(self.len);
let bits = (value & Self::mask()) << off;
if let Some(w) = self.words.get_mut(Self::word_of(self.len)) {
*w = (*w & !(Self::mask() << off)) | bits;
} else {
debug_assert!(off == 0);
self.words.push(bits);
}
self.len += 1;
}
#[inline]
pub fn pop(&mut self) -> Option<usize> {
if self.len == 0 {
return None;
}
let value = unsafe { self.get_unchecked(self.len - 1) };
self.len -= 1;
self.words.truncate(Self::words_for(self.len));
Some(value)
}
pub fn truncate(&mut self, new_len: usize) {
if new_len >= self.len {
return;
}
self.len = new_len;
self.words.truncate(Self::words_for(new_len));
}
pub fn clear(&mut self) {
self.words.clear();
self.len = 0;
}
pub fn append(&mut self, other: &mut Self) {
let other_len = other.len;
if other_len == 0 {
return;
}
let per = Self::items_per_word();
self.reserve(other_len);
if self.len % per == 0 && self.words.len() == Self::words_for(self.len)
{
self.words
.extend_from_slice(&other.words[..Self::words_for(other_len)]);
self.len += other_len;
} else {
for i in 0..other_len {
self.push(unsafe { other.get_unchecked(i) });
}
}
other.clear();
}
pub fn extend_from_packed(
&mut self,
other: &Self,
start: usize,
len: usize,
) {
assert!(
start <= other.len && len <= other.len - start,
"source range out of bounds: the len is {} but the range is \
{start}..{start}+{len}",
other.len
);
if len == 0 {
return;
}
let per = Self::items_per_word();
self.reserve(len);
if self.len % per == 0
&& start % per == 0
&& self.words.len() == Self::words_for(self.len)
{
let first = start / per;
let nwords = Self::words_for(len);
self.words
.extend_from_slice(&other.words[first..first + nwords]);
self.len += len;
} else {
for i in 0..len {
self.push(unsafe { other.get_unchecked(start + i) });
}
}
}
pub fn extend_fill(&mut self, value: usize, count: usize) {
if count == 0 {
return;
}
let per = Self::items_per_word();
let mask = Self::mask();
let v = value & mask;
self.reserve(count);
let mut filled = 0;
while filled < count && self.len % per != 0 {
self.push(v);
filled += 1;
}
let full_words = (count - filled) / per;
if full_words > 0 && self.words.len() == Self::words_for(self.len) {
let rep = v.wrapping_mul(usize::MAX / mask);
self.words.resize(self.words.len() + full_words, rep);
self.len += full_words * per;
filled += full_words * per;
}
while filled < count {
self.push(v);
filled += 1;
}
}
#[inline(always)]
fn bit_span_mask(lo: usize, hi: usize) -> usize {
let width = hi - lo;
if width == usize::BITS as usize {
usize::MAX
} else {
((1usize << width) - 1) << lo
}
}
pub fn fill_range(&mut self, start: usize, len: usize, value: usize) {
assert!(
start <= self.len && len <= self.len - start,
"fill range out of bounds: the len is {} but the range is \
{start}..{start}+{len}",
self.len
);
if len == 0 {
return;
}
let mask = Self::mask();
let rep = (value & mask).wrapping_mul(usize::MAX / mask);
let end = start + len;
let mut wi = Self::word_of(start);
let last = Self::word_of(end - 1);
let head_lo = Self::offset_of(start);
let tail_hi = Self::offset_of(end - 1) + BITS as usize;
let word_bits = usize::BITS as usize;
if wi == last {
let m = Self::bit_span_mask(head_lo, tail_hi);
let w = &mut self.words[wi];
*w = (*w & !m) | (rep & m);
return;
}
if head_lo != 0 {
let m = Self::bit_span_mask(head_lo, word_bits);
let w = &mut self.words[wi];
*w = (*w & !m) | (rep & m);
wi += 1;
}
let interior_end = if tail_hi == word_bits { last + 1 } else { last };
for w in &mut self.words[wi..interior_end] {
*w = rep;
}
if tail_hi != word_bits {
let m = Self::bit_span_mask(0, tail_hi);
let w = &mut self.words[last];
*w = (*w & !m) | (rep & m);
}
}
#[inline(always)]
fn read_bits(words: &[usize], bit: usize, n: usize) -> usize {
let word_bits = usize::BITS as usize;
let w = bit / word_bits;
let r = bit % word_bits;
let val = if r == 0 {
words[w]
} else if r + n <= word_bits {
words[w] >> r
} else {
(words[w] >> r) | (words[w + 1] << (word_bits - r))
};
if n == word_bits {
val
} else {
val & ((1usize << n) - 1)
}
}
#[inline(always)]
fn write_bits(words: &mut [usize], bit: usize, n: usize, val: usize) {
let word_bits = usize::BITS as usize;
let w = bit / word_bits;
let r = bit % word_bits;
let m = if n == word_bits {
usize::MAX
} else {
(1usize << n) - 1
};
let v = val & m;
words[w] = (words[w] & !(m << r)) | (v << r);
if r + n > word_bits {
let spill_mask = (1usize << (r + n - word_bits)) - 1;
words[w + 1] =
(words[w + 1] & !spill_mask) | (v >> (word_bits - r));
}
}
pub fn copy_lanes(&mut self, src: usize, dst: usize, count: usize) {
assert!(
src <= self.len
&& count <= self.len - src
&& dst <= self.len
&& count <= self.len - dst,
"copy range out of bounds: the len is {} but the ranges are \
{src}..{src}+{count} and {dst}..{dst}+{count}",
self.len
);
if count == 0 || src == dst {
return;
}
let word_bits = usize::BITS as usize;
let b = BITS as usize;
let total = count * b;
let sbit = src * b;
let dbit = dst * b;
if dbit < sbit {
let mut done = 0;
while done < total {
let n = word_bits.min(total - done);
let v = Self::read_bits(&self.words, sbit + done, n);
Self::write_bits(&mut self.words, dbit + done, n, v);
done += n;
}
} else {
let mut left = total;
while left > 0 {
let n = word_bits.min(left);
left -= n;
let v = Self::read_bits(&self.words, sbit + left, n);
Self::write_bits(&mut self.words, dbit + left, n, v);
}
}
}
pub fn range_eq(
&self,
start: usize,
other: &Self,
other_start: usize,
len: usize,
) -> bool {
assert!(
start <= self.len
&& len <= self.len - start
&& other_start <= other.len
&& len <= other.len - other_start,
"compare range out of bounds"
);
if len == 0 {
return true;
}
let per = Self::items_per_word();
if start % per != 0 || other_start % per != 0 {
return (0..len)
.all(|i| self.get(start + i) == other.get(other_start + i));
}
let (sw, ow) = (start / per, other_start / per);
let full = len / per;
if self.words[sw..sw + full] != other.words[ow..ow + full] {
return false;
}
let tail = len % per;
if tail == 0 {
return true;
}
let mask = (1usize << (tail * BITS as usize)) - 1;
(self.words[sw + full] & mask) == (other.words[ow + full] & mask)
}
pub fn count_in(&self, start: usize, len: usize, value: usize) -> usize {
assert!(
start <= self.len && len <= self.len - start,
"count range out of bounds: the len is {} but the range is \
{start}..{start}+{len}",
self.len
);
let per = Self::items_per_word();
let mut total = 0usize;
let mut idx = start;
let end = start + len;
while idx < end {
let wstart = (idx / per) * per;
let wend = wstart + per;
if idx == wstart && wend <= end {
total += count_word_in::<BITS>(self.words[wstart / per], value);
idx = wend;
} else {
let stop = wend.min(end);
for i in idx..stop {
if unsafe { self.get_unchecked(i) } == value {
total += 1;
}
}
idx = stop;
}
}
total
}
}
#[doc(hidden)]
#[inline]
pub fn count_word_in<const BITS: u32>(word: usize, value: usize) -> usize {
if BITS == 1 {
if value != 0 {
word.count_ones() as usize
} else {
word.count_zeros() as usize
}
} else {
let mask = (1usize << BITS) - 1;
let rep1 = usize::MAX / mask;
let highrep = rep1 << (BITS - 1);
let target = value & mask;
let rep_target = target.wrapping_mul(rep1);
let u = word ^ rep_target;
let mut collapsed = u & rep1;
let mut j: u32 = 1;
while j < BITS {
collapsed |= (u >> j) & rep1;
j += 1;
}
let nz_high = collapsed << (BITS - 1);
let matches = !nz_high & highrep;
matches.count_ones() as usize
}
}
impl<const BITS: u32> Default for PackedArray<BITS> {
#[inline]
fn default() -> Self {
Self::new()
}
}
impl<const BITS: u32> Clone for PackedArray<BITS> {
fn clone(&self) -> Self {
Self {
words: self.words.clone(),
len: self.len,
}
}
}
pub trait BitPack: Clone + Default + core::fmt::Debug + Sized {
fn new() -> Self;
fn with_capacity(capacity: usize) -> Self;
fn len(&self) -> usize;
fn is_empty(&self) -> bool;
fn capacity(&self) -> usize;
fn reserve(&mut self, additional: usize);
fn reserve_exact(&mut self, additional: usize);
fn shrink_to_fit(&mut self);
fn truncate(&mut self, new_len: usize);
unsafe fn set_len(&mut self, new_len: usize);
fn get(&self, index: usize) -> usize;
unsafe fn get_unchecked(&self, index: usize) -> usize;
fn word(&self, index: usize) -> usize;
fn set(&mut self, index: usize, value: usize);
unsafe fn set_unchecked(&mut self, index: usize, value: usize);
fn push(&mut self, value: usize);
fn pop(&mut self) -> Option<usize>;
fn clear(&mut self);
fn append(&mut self, other: &mut Self);
fn extend_from_packed(&mut self, other: &Self, start: usize, len: usize);
fn extend_fill(&mut self, value: usize, count: usize);
fn fill_range(&mut self, start: usize, len: usize, value: usize);
fn copy_lanes(&mut self, src: usize, dst: usize, count: usize);
fn range_eq(
&self,
start: usize,
other: &Self,
other_start: usize,
len: usize,
) -> bool;
fn count_in(&self, start: usize, len: usize, value: usize) -> usize;
}
impl<const BITS: u32> BitPack for PackedArray<BITS> {
#[inline]
fn new() -> Self {
PackedArray::new()
}
#[inline]
fn with_capacity(capacity: usize) -> Self {
PackedArray::with_capacity(capacity)
}
#[inline]
fn len(&self) -> usize {
PackedArray::len(self)
}
#[inline]
fn is_empty(&self) -> bool {
PackedArray::is_empty(self)
}
#[inline]
fn capacity(&self) -> usize {
PackedArray::capacity(self)
}
#[inline]
fn reserve(&mut self, additional: usize) {
PackedArray::reserve(self, additional);
}
#[inline]
fn reserve_exact(&mut self, additional: usize) {
PackedArray::reserve_exact(self, additional);
}
#[inline]
fn shrink_to_fit(&mut self) {
PackedArray::shrink_to_fit(self);
}
#[inline]
fn truncate(&mut self, new_len: usize) {
PackedArray::truncate(self, new_len);
}
#[inline]
unsafe fn set_len(&mut self, new_len: usize) {
unsafe { PackedArray::set_len(self, new_len) };
}
#[inline]
fn get(&self, index: usize) -> usize {
PackedArray::get(self, index)
}
#[inline(always)]
unsafe fn get_unchecked(&self, index: usize) -> usize {
unsafe { PackedArray::get_unchecked(self, index) }
}
#[inline]
fn word(&self, index: usize) -> usize {
debug_assert!(index < self.words.len());
self.words[index]
}
#[inline]
fn set(&mut self, index: usize, value: usize) {
PackedArray::set(self, index, value);
}
#[inline(always)]
unsafe fn set_unchecked(&mut self, index: usize, value: usize) {
unsafe { PackedArray::set_unchecked(self, index, value) };
}
#[inline]
fn push(&mut self, value: usize) {
PackedArray::push(self, value);
}
#[inline]
fn pop(&mut self) -> Option<usize> {
PackedArray::pop(self)
}
#[inline]
fn clear(&mut self) {
PackedArray::clear(self);
}
#[inline]
fn append(&mut self, other: &mut Self) {
PackedArray::append(self, other);
}
#[inline]
fn extend_from_packed(&mut self, other: &Self, start: usize, len: usize) {
PackedArray::extend_from_packed(self, other, start, len);
}
#[inline]
fn extend_fill(&mut self, value: usize, count: usize) {
PackedArray::extend_fill(self, value, count);
}
#[inline]
fn fill_range(&mut self, start: usize, len: usize, value: usize) {
PackedArray::fill_range(self, start, len, value);
}
#[inline]
fn copy_lanes(&mut self, src: usize, dst: usize, count: usize) {
PackedArray::copy_lanes(self, src, dst, count);
}
#[inline]
fn range_eq(
&self,
start: usize,
other: &Self,
other_start: usize,
len: usize,
) -> bool {
PackedArray::range_eq(self, start, other, other_start, len)
}
#[inline]
fn count_in(&self, start: usize, len: usize, value: usize) -> usize {
PackedArray::count_in(self, start, len, value)
}
}