binius_utils/
bitwise.rs

1// Copyright 2025 Irreducible Inc.
2
3use std::{
4	marker::PhantomData,
5	ops::{BitAnd, Shr},
6};
7
8use trait_set::trait_set;
9
10use super::random_access_sequence::RandomAccessSequence;
11
12// A trait alias for a type that can act as a bitmask.
13//
14// It is intended to be a drop in constraint for a primitive integer type.
15// Take note that `Shr` implementation sign extends signed types.
16trait_set! {
17	pub trait Bitwise =
18		BitAnd<Output=Self> +
19		Shr<usize, Output=Self> +
20		From<u8> +
21		PartialEq<Self> +
22		Sized +
23		Sync +
24		Copy;
25}
26
27/// An adaptor structure that wraps a slice of integers and presents a
28/// random access sequence of bits at a given offset.
29pub struct BitSelector<B: Bitwise, S: AsRef<[B]>> {
30	bit_offset: usize,
31	slice: S,
32	_b_marker: PhantomData<B>,
33}
34
35impl<B: Bitwise, S: AsRef<[B]>> BitSelector<B, S> {
36	pub fn new(bit_offset: usize, slice: S) -> Self {
37		Self {
38			bit_offset,
39			slice,
40			_b_marker: PhantomData,
41		}
42	}
43}
44
45impl<B: Bitwise, S: AsRef<[B]>> RandomAccessSequence<bool> for BitSelector<B, S> {
46	#[inline(always)]
47	fn len(&self) -> usize {
48		self.slice.as_ref().len()
49	}
50
51	#[inline(always)]
52	fn get(&self, index: usize) -> bool {
53		(self.slice.as_ref()[index] >> self.bit_offset) & B::from(1u8) != B::from(0u8)
54	}
55
56	#[inline(always)]
57	unsafe fn get_unchecked(&self, index: usize) -> bool {
58		unsafe {
59			(*self.slice.as_ref().get_unchecked(index) >> self.bit_offset) & B::from(1u8)
60				!= B::from(0u8)
61		}
62	}
63}
64
65#[cfg(test)]
66mod tests {
67	use super::*;
68
69	#[test]
70	fn test_bit_selector_on_sequential_integer_range() {
71		let log_n = 10;
72
73		let integers = (0u16..1 << log_n).collect::<Vec<_>>();
74
75		for bit_offset in 0..log_n {
76			let selector = BitSelector::new(bit_offset, &integers);
77
78			for i in 0..1 << log_n {
79				assert_eq!(selector.get(i), (i >> bit_offset) & 1 != 0);
80			}
81		}
82	}
83}