azalea_core/
bitset.rs

1use std::io::{self, Cursor, Write};
2
3use azalea_buf::{AzBuf, AzaleaRead, AzaleaWrite, BufReadError};
4
5/// Represents Java's BitSet, a list of bits.
6#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, AzBuf)]
7pub struct BitSet {
8    data: Vec<u64>,
9}
10
11const ADDRESS_BITS_PER_WORD: usize = 6;
12
13// the Index trait requires us to return a reference, but we can't do that
14impl BitSet {
15    pub fn new(num_bits: usize) -> Self {
16        BitSet {
17            data: vec![0; num_bits.div_ceil(64)],
18        }
19    }
20
21    pub fn index(&self, index: usize) -> bool {
22        (self.data[index / 64] & (1u64 << (index % 64))) != 0
23    }
24
25    fn check_range(&self, from_index: usize, to_index: usize) {
26        assert!(
27            from_index <= to_index,
28            "fromIndex: {from_index} > toIndex: {to_index}",
29        );
30    }
31
32    fn word_index(&self, bit_index: usize) -> usize {
33        bit_index >> ADDRESS_BITS_PER_WORD
34    }
35
36    pub fn clear(&mut self, from_index: usize, mut to_index: usize) {
37        self.check_range(from_index, to_index);
38
39        if from_index == to_index {
40            return;
41        }
42
43        let start_word_index = self.word_index(from_index);
44        if start_word_index >= self.data.len() {
45            return;
46        }
47
48        let mut end_word_index = self.word_index(to_index - 1);
49        if end_word_index >= self.data.len() {
50            to_index = self.len();
51            end_word_index = self.data.len() - 1;
52        }
53
54        let first_word_mask = u64::MAX.wrapping_shl(
55            from_index
56                .try_into()
57                .expect("from_index shouldn't be larger than u32"),
58        );
59        let last_word_mask = u64::MAX.wrapping_shr((64 - (to_index % 64)) as u32);
60        if start_word_index == end_word_index {
61            // Case 1: One word
62            self.data[start_word_index] &= !(first_word_mask & last_word_mask);
63        } else {
64            // Case 2: Multiple words
65            // Handle first word
66            self.data[start_word_index] &= !first_word_mask;
67
68            // Handle intermediate words, if any
69            for i in start_word_index + 1..end_word_index {
70                self.data[i] = 0;
71            }
72
73            // Handle last word
74            self.data[end_word_index] &= !last_word_mask;
75        }
76    }
77
78    /// Returns the maximum potential items in the BitSet. This will be
79    /// divisible by 64.
80    fn len(&self) -> usize {
81        self.data.len() * 64
82    }
83
84    /// Returns the index of the first bit that is set to `false`
85    /// that occurs on or after the specified starting index.
86    pub fn next_clear_bit(&self, from_index: usize) -> usize {
87        let mut u = self.word_index(from_index);
88        if u >= self.data.len() {
89            return from_index;
90        }
91
92        let mut word = !self.data[u] & (u64::MAX.wrapping_shl(from_index.try_into().unwrap()));
93
94        loop {
95            if word != 0 {
96                return (u * 64) + word.trailing_zeros() as usize;
97            }
98            u += 1;
99            if u == self.data.len() {
100                return self.data.len() * 64;
101            }
102            word = !self.data[u];
103        }
104    }
105
106    pub fn set(&mut self, bit_index: usize) {
107        self.data[bit_index / 64] |= 1u64 << (bit_index % 64);
108    }
109}
110
111impl From<Vec<u64>> for BitSet {
112    fn from(data: Vec<u64>) -> Self {
113        BitSet { data }
114    }
115}
116
117impl From<Vec<u8>> for BitSet {
118    fn from(data: Vec<u8>) -> Self {
119        let mut words = vec![0; data.len().div_ceil(8)];
120        for (i, byte) in data.iter().enumerate() {
121            words[i / 8] |= (*byte as u64) << ((i % 8) * 8);
122        }
123        BitSet { data: words }
124    }
125}
126
127/// A compact fixed-size array of bits.
128///
129/// The `N` is the number of bits reserved for the bitset. You're encouraged to
130/// use it like `FixedBitSet<20>` if you need 20 bits.
131///
132/// Note that this is optimized for fast serialization and deserialization for
133/// Minecraft, and may not be as performant as it could be for other purposes.
134/// Notably, the internal representation is an array of `u8`s even though
135/// `usize` would be slightly faster.
136#[derive(Debug, Clone, PartialEq, Eq, Hash)]
137pub struct FixedBitSet<const N: usize>
138where
139    [u8; bits_to_bytes(N)]: Sized,
140{
141    data: [u8; bits_to_bytes(N)],
142}
143
144impl<const N: usize> FixedBitSet<N>
145where
146    [u8; bits_to_bytes(N)]: Sized,
147{
148    pub const fn new() -> Self {
149        FixedBitSet {
150            data: [0; bits_to_bytes(N)],
151        }
152    }
153
154    #[inline]
155    pub fn index(&self, index: usize) -> bool {
156        (self.data[index / 8] & (1u8 << (index % 8))) != 0
157    }
158
159    #[inline]
160    pub fn set(&mut self, bit_index: usize) {
161        self.data[bit_index / 8] |= 1u8 << (bit_index % 8);
162    }
163}
164
165impl<const N: usize> AzaleaRead for FixedBitSet<N>
166where
167    [u8; bits_to_bytes(N)]: Sized,
168{
169    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
170        let mut data = [0; bits_to_bytes(N)];
171        for item in data.iter_mut().take(bits_to_bytes(N)) {
172            *item = u8::azalea_read(buf)?;
173        }
174        Ok(FixedBitSet { data })
175    }
176}
177impl<const N: usize> AzaleaWrite for FixedBitSet<N>
178where
179    [u8; bits_to_bytes(N)]: Sized,
180{
181    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
182        for i in 0..bits_to_bytes(N) {
183            self.data[i].azalea_write(buf)?;
184        }
185        Ok(())
186    }
187}
188impl<const N: usize> Default for FixedBitSet<N>
189where
190    [u8; bits_to_bytes(N)]: Sized,
191{
192    fn default() -> Self {
193        Self::new()
194    }
195}
196
197pub const fn bits_to_bytes(n: usize) -> usize {
198    n.div_ceil(8)
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn test_bitset() {
207        let mut bitset = BitSet::new(64);
208        assert!(!bitset.index(0));
209        assert!(!bitset.index(1));
210        assert!(!bitset.index(2));
211        bitset.set(1);
212        assert!(!bitset.index(0));
213        assert!(bitset.index(1));
214        assert!(!bitset.index(2));
215    }
216
217    #[test]
218    fn test_clear() {
219        let mut bitset = BitSet::new(128);
220        bitset.set(62);
221        bitset.set(63);
222        bitset.set(64);
223        bitset.set(65);
224        bitset.set(66);
225
226        bitset.clear(63, 65);
227
228        assert!(bitset.index(62));
229        assert!(!bitset.index(63));
230        assert!(!bitset.index(64));
231        assert!(bitset.index(65));
232        assert!(bitset.index(66));
233    }
234
235    #[test]
236    fn test_clear_2() {
237        let mut bitset = BitSet::new(128);
238        bitset.set(64);
239        bitset.set(65);
240        bitset.set(66);
241        bitset.set(67);
242        bitset.set(68);
243
244        bitset.clear(65, 67);
245
246        assert!(bitset.index(64));
247        assert!(!bitset.index(65));
248        assert!(!bitset.index(66));
249        assert!(bitset.index(67));
250        assert!(bitset.index(68));
251    }
252}
azalea_core/bitset.rs

azalea_core/
bitset.rs
