azalea_core/
bitset.rs

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

azalea_core/
bitset.rs
