azalea_core/
bitset.rs

1use std::{
2    io::{self, Cursor, Write},
3    ops::Range,
4};
5
6use azalea_buf::{AzBuf, BufReadError};
7
8/// Represents Java's BitSet, a list of bits.
9#[derive(AzBuf, Clone, Debug, Default, Eq, Hash, PartialEq)]
10pub struct BitSet {
11    data: Box<[u64]>,
12}
13
14/// `log2(64)`.
15const LOG2_BITS_PER_WORD: usize = 6;
16
17// the Index trait requires us to return a reference, but we can't do that
18impl BitSet {
19    #[inline]
20    pub fn new(num_bits: usize) -> Self {
21        BitSet {
22            data: vec![0; num_bits.div_ceil(64)].into(),
23        }
24    }
25
26    /// Returns the bit at the given index.
27    ///
28    /// # Panics
29    ///
30    /// Panics if the index is out of bounds. Use [`Self::get`] for a
31    /// non-panicking version.
32    #[inline]
33    pub fn index(&self, index: usize) -> bool {
34        self.get(index).unwrap_or_else(|| {
35            let len = self.len();
36            panic!("index out of bounds: the len is {len} but the index is {index}")
37        })
38    }
39
40    #[inline]
41    pub fn get(&self, index: usize) -> Option<bool> {
42        self.data
43            .get(index / 64)
44            .map(|word| (word & (1u64 << (index % 64))) != 0)
45    }
46
47    /// Zeros the bits in the given range.
48    ///
49    /// # Panics
50    ///
51    /// Panics if the range is invalid (start > end).
52    pub fn clear(&mut self, range: Range<usize>) {
53        assert!(
54            range.start <= range.end,
55            "Range ends before it starts; {} must be less than or equal to {}",
56            range.start,
57            range.end
58        );
59
60        let from_idx = range.start;
61        let mut to_idx = range.end;
62
63        if from_idx == to_idx {
64            return;
65        }
66
67        let start_word_idx = self.word_index(from_idx);
68        if start_word_idx >= self.data.len() {
69            return;
70        }
71
72        let mut end_word_idx = self.word_index(to_idx - 1);
73        if end_word_idx >= self.data.len() {
74            to_idx = self.len();
75            end_word_idx = self.data.len() - 1;
76        }
77
78        let first_word_mask = u64::MAX.wrapping_shl(
79            from_idx
80                .try_into()
81                .expect("from_index shouldn't be larger than u32"),
82        );
83        let last_word_mask = u64::MAX.wrapping_shr((64 - (to_idx % 64)) as u32);
84        if start_word_idx == end_word_idx {
85            // one word
86            self.data[start_word_idx] &= !(first_word_mask & last_word_mask);
87        } else {
88            // multiple words
89            self.data[start_word_idx] &= !first_word_mask;
90            for i in (start_word_idx + 1)..end_word_idx {
91                self.data[i] = 0;
92            }
93            self.data[end_word_idx] &= !last_word_mask;
94        }
95    }
96
97    /// Returns the index of the first bit that is set to `false`
98    /// that occurs on or after the specified starting index.
99    pub fn next_clear_bit(&self, from_index: usize) -> usize {
100        let mut u = self.word_index(from_index);
101        if u >= self.data.len() {
102            return from_index;
103        }
104
105        let mut word = !self.data[u] & (u64::MAX.wrapping_shl(from_index.try_into().unwrap()));
106
107        loop {
108            if word != 0 {
109                return (u * 64) + word.trailing_zeros() as usize;
110            }
111            u += 1;
112            if u == self.data.len() {
113                return self.data.len() * 64;
114            }
115            word = !self.data[u];
116        }
117    }
118
119    #[inline]
120    fn word_index(&self, bit_index: usize) -> usize {
121        bit_index >> LOG2_BITS_PER_WORD
122    }
123
124    /// Sets the bit at the given index to true.
125    ///
126    /// # Panics
127    ///
128    /// Panics if the index is out of bounds. Check [`Self::len`] first
129    /// if you need to avoid this.
130    #[inline]
131    pub fn set(&mut self, bit_index: usize) {
132        self.data[bit_index / 64] |= 1u64 << (bit_index % 64);
133    }
134
135    /// Returns the indices of all bits that are set to `true`.
136    pub fn iter_ones(&self) -> impl Iterator<Item = usize> {
137        (0..self.len()).filter(|i| self.index(*i))
138    }
139
140    /// Returns the maximum number of items that could be in this `BitSet`.
141    ///
142    /// This will always be a multiple of 64.
143    #[inline]
144    pub fn len(&self) -> usize {
145        self.data.len() * 64
146    }
147
148    /// Returns true if the `BitSet` was created with a size of 0.
149    ///
150    /// Equivalent to `self.len() == 0`.
151    pub fn is_empty(&self) -> bool {
152        self.len() == 0
153    }
154}
155
156impl From<Vec<u64>> for BitSet {
157    fn from(data: Vec<u64>) -> Self {
158        BitSet { data: data.into() }
159    }
160}
161
162impl From<Vec<u8>> for BitSet {
163    fn from(data: Vec<u8>) -> Self {
164        let mut words = vec![0; data.len().div_ceil(8)];
165        for (i, byte) in data.iter().enumerate() {
166            words[i / 8] |= (*byte as u64) << ((i % 8) * 8);
167        }
168        BitSet { data: words.into() }
169    }
170}
171
172/// A compact fixed-size array of bits.
173///
174/// The `N` is the number of bits reserved for the bitset. You're encouraged to
175/// use it like `FixedBitSet<20>` if you need 20 bits.
176///
177/// Note that this is optimized for fast serialization and deserialization for
178/// Minecraft, and may not be as performant as it could be for other purposes.
179/// Consider using [`FastFixedBitSet`] if you don't need the `AzBuf`
180/// implementation.
181#[derive(Clone, Debug, Eq, Hash, PartialEq)]
182pub struct FixedBitSet<const N: usize>
183where
184    [u8; bits_to_bytes(N)]: Sized,
185{
186    data: [u8; bits_to_bytes(N)],
187}
188
189impl<const N: usize> FixedBitSet<N>
190where
191    [u8; bits_to_bytes(N)]: Sized,
192{
193    pub const fn new() -> Self {
194        FixedBitSet {
195            data: [0; bits_to_bytes(N)],
196        }
197    }
198
199    pub const fn new_with_data(data: [u8; bits_to_bytes(N)]) -> Self {
200        FixedBitSet { data }
201    }
202
203    #[inline]
204    pub fn index(&self, index: usize) -> bool {
205        (self.data[index / 8] & (1u8 << (index % 8))) != 0
206    }
207
208    #[inline]
209    pub fn set(&mut self, bit_index: usize) {
210        self.data[bit_index / 8] |= 1u8 << (bit_index % 8);
211    }
212}
213
214impl<const N: usize> AzBuf for FixedBitSet<N>
215where
216    [u8; bits_to_bytes(N)]: Sized,
217{
218    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
219        let mut data = [0; bits_to_bytes(N)];
220        for item in data.iter_mut().take(bits_to_bytes(N)) {
221            *item = u8::azalea_read(buf)?;
222        }
223        Ok(FixedBitSet { data })
224    }
225    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
226        for i in 0..bits_to_bytes(N) {
227            self.data[i].azalea_write(buf)?;
228        }
229        Ok(())
230    }
231}
232impl<const N: usize> Default for FixedBitSet<N>
233where
234    [u8; bits_to_bytes(N)]: Sized,
235{
236    fn default() -> Self {
237        Self::new()
238    }
239}
240
241pub const fn bits_to_bytes(n: usize) -> usize {
242    n.div_ceil(8)
243}
244
245/// A slightly faster compact fixed-size array of bits.
246///
247/// The `N` is the number of bits reserved for the bitset. You're encouraged to
248/// use it like `FastFixedBitSet<20>` if you need 20 bits.
249///
250/// This is almost identical to [`FixedBitSet`], but more efficient (~20% faster
251/// access) and doesn't implement `AzBuf`.
252#[derive(Clone, Debug, Eq, Hash, PartialEq)]
253pub struct FastFixedBitSet<const N: usize>
254where
255    [u64; bits_to_longs(N)]: Sized,
256{
257    data: [u64; bits_to_longs(N)],
258}
259impl<const N: usize> FastFixedBitSet<N>
260where
261    [u64; bits_to_longs(N)]: Sized,
262{
263    pub const fn new() -> Self {
264        FastFixedBitSet {
265            data: [0; bits_to_longs(N)],
266        }
267    }
268
269    #[inline]
270    pub fn index(&self, index: usize) -> bool {
271        (self.data[index / 64] & (1u64 << (index % 64))) != 0
272    }
273
274    #[inline]
275    pub fn set(&mut self, bit_index: usize) {
276        self.data[bit_index / 64] |= 1u64 << (bit_index % 64);
277    }
278}
279impl<const N: usize> Default for FastFixedBitSet<N>
280where
281    [u64; bits_to_longs(N)]: Sized,
282{
283    fn default() -> Self {
284        Self::new()
285    }
286}
287pub const fn bits_to_longs(n: usize) -> usize {
288    n.div_ceil(64)
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn test_bitset() {
297        let mut bitset = BitSet::new(64);
298        assert!(!bitset.index(0));
299        assert!(!bitset.index(1));
300        assert!(!bitset.index(2));
301        bitset.set(1);
302        assert!(!bitset.index(0));
303        assert!(bitset.index(1));
304        assert!(!bitset.index(2));
305    }
306
307    #[test]
308    fn test_clear() {
309        let mut bitset = BitSet::new(128);
310        bitset.set(62);
311        bitset.set(63);
312        bitset.set(64);
313        bitset.set(65);
314        bitset.set(66);
315
316        bitset.clear(63..65);
317
318        assert!(bitset.index(62));
319        assert!(!bitset.index(63));
320        assert!(!bitset.index(64));
321        assert!(bitset.index(65));
322        assert!(bitset.index(66));
323    }
324
325    #[test]
326    fn test_clear_2() {
327        let mut bitset = BitSet::new(128);
328        bitset.set(64);
329        bitset.set(65);
330        bitset.set(66);
331        bitset.set(67);
332        bitset.set(68);
333
334        bitset.clear(65..67);
335
336        assert!(bitset.index(64));
337        assert!(!bitset.index(65));
338        assert!(!bitset.index(66));
339        assert!(bitset.index(67));
340        assert!(bitset.index(68));
341    }
342}
azalea_core/bitset.rs

azalea_core/
bitset.rs
