azalea_core/
bitset.rs

1use std::{
2    io::{self, Cursor, Write},
3    ops::Range,
4};
5
6use azalea_buf::{AzBuf, AzaleaRead, AzaleaWrite, BufReadError};
7
8/// Represents Java's BitSet, a list of bits.
9#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, AzBuf)]
10pub struct BitSet {
11    data: Vec<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)],
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 greater than {}",
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 }
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 }
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
180/// `AzaleaRead`/`AzaleaWrite` implementation.
181#[derive(Debug, Clone, PartialEq, Eq, Hash)]
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> AzaleaRead 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}
226impl<const N: usize> AzaleaWrite for FixedBitSet<N>
227where
228    [u8; bits_to_bytes(N)]: Sized,
229{
230    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
231        for i in 0..bits_to_bytes(N) {
232            self.data[i].azalea_write(buf)?;
233        }
234        Ok(())
235    }
236}
237impl<const N: usize> Default for FixedBitSet<N>
238where
239    [u8; bits_to_bytes(N)]: Sized,
240{
241    fn default() -> Self {
242        Self::new()
243    }
244}
245
246pub const fn bits_to_bytes(n: usize) -> usize {
247    n.div_ceil(8)
248}
249
250/// A slightly faster compact fixed-size array of bits.
251///
252/// The `N` is the number of bits reserved for the bitset. You're encouraged to
253/// use it like `FastFixedBitSet<20>` if you need 20 bits.
254///
255/// This is almost identical to [`FixedBitSet`], but more efficient (~20% faster
256/// access) and doesn't implement `AzaleaRead`/`AzaleaWrite`.
257#[derive(Debug, Clone, PartialEq, Eq, Hash)]
258pub struct FastFixedBitSet<const N: usize>
259where
260    [u64; bits_to_longs(N)]: Sized,
261{
262    data: [u64; bits_to_longs(N)],
263}
264impl<const N: usize> FastFixedBitSet<N>
265where
266    [u64; bits_to_longs(N)]: Sized,
267{
268    pub const fn new() -> Self {
269        FastFixedBitSet {
270            data: [0; bits_to_longs(N)],
271        }
272    }
273
274    #[inline]
275    pub fn index(&self, index: usize) -> bool {
276        (self.data[index / 64] & (1u64 << (index % 64))) != 0
277    }
278
279    #[inline]
280    pub fn set(&mut self, bit_index: usize) {
281        self.data[bit_index / 64] |= 1u64 << (bit_index % 64);
282    }
283}
284impl<const N: usize> Default for FastFixedBitSet<N>
285where
286    [u64; bits_to_longs(N)]: Sized,
287{
288    fn default() -> Self {
289        Self::new()
290    }
291}
292pub const fn bits_to_longs(n: usize) -> usize {
293    n.div_ceil(64)
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_bitset() {
302        let mut bitset = BitSet::new(64);
303        assert!(!bitset.index(0));
304        assert!(!bitset.index(1));
305        assert!(!bitset.index(2));
306        bitset.set(1);
307        assert!(!bitset.index(0));
308        assert!(bitset.index(1));
309        assert!(!bitset.index(2));
310    }
311
312    #[test]
313    fn test_clear() {
314        let mut bitset = BitSet::new(128);
315        bitset.set(62);
316        bitset.set(63);
317        bitset.set(64);
318        bitset.set(65);
319        bitset.set(66);
320
321        bitset.clear(63..65);
322
323        assert!(bitset.index(62));
324        assert!(!bitset.index(63));
325        assert!(!bitset.index(64));
326        assert!(bitset.index(65));
327        assert!(bitset.index(66));
328    }
329
330    #[test]
331    fn test_clear_2() {
332        let mut bitset = BitSet::new(128);
333        bitset.set(64);
334        bitset.set(65);
335        bitset.set(66);
336        bitset.set(67);
337        bitset.set(68);
338
339        bitset.clear(65..67);
340
341        assert!(bitset.index(64));
342        assert!(!bitset.index(65));
343        assert!(!bitset.index(66));
344        assert!(bitset.index(67));
345        assert!(bitset.index(68));
346    }
347}
azalea_core/bitset.rs

azalea_core/
bitset.rs
