add new apis for BitSet

closes #241

3 files changed, 82 insertions, 45 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 909c39fc..e7f05bdf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ is breaking anyways, semantic versioning is not followed.
 - `ItemStack` now has a `get_component` function that supports default components.
 - `Client::nearest_entity_by`.
 - Sneaking/crouching.
+- `BitSet::len`, `BitSet::get`, `BitSet::iter_ones`.
 
 ### Changed
 
diff --git a/azalea-core/src/bitset.rs b/azalea-core/src/bitset.rs
index fbb12ee0..0be3fabd 100644
--- a/azalea-core/src/bitset.rs
+++ b/azalea-core/src/bitset.rs
@@ -1,4 +1,7 @@
-use std::io::{self, Cursor, Write};
+use std::{
+    io::{self, Cursor, Write},
+    ops::Range,
+};
 
 use azalea_buf::{AzBuf, AzaleaRead, AzaleaWrite, BufReadError};
 
@@ -8,79 +11,89 @@ pub struct BitSet {
     data: Vec<u64>,
 }
 
-const ADDRESS_BITS_PER_WORD: usize = 6;
+/// `log2(64)`.
+const LOG2_BITS_PER_WORD: usize = 6;
 
 // the Index trait requires us to return a reference, but we can't do that
 impl BitSet {
+    #[inline]
     pub fn new(num_bits: usize) -> Self {
         BitSet {
             data: vec![0; num_bits.div_ceil(64)],
         }
     }
 
+    /// Returns the bit at the given index.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds. Use [`Self::get`] for a
+    /// non-panicking version.
+    #[inline]
     pub fn index(&self, index: usize) -> bool {
-        (self.data[index / 64] & (1u64 << (index % 64))) != 0
+        self.get(index).unwrap_or_else(|| {
+            let len = self.len();
+            panic!("index out of bounds: the len is {len} but the index is {index}")
+        })
     }
 
-    fn check_range(&self, from_index: usize, to_index: usize) {
-        assert!(
-            from_index <= to_index,
-            "fromIndex: {from_index} > toIndex: {to_index}",
-        );
+    #[inline]
+    pub fn get(&self, index: usize) -> Option<bool> {
+        self.data
+            .get(index / 64)
+            .map(|word| (word & (1u64 << (index % 64))) != 0)
     }
 
-    fn word_index(&self, bit_index: usize) -> usize {
-        bit_index >> ADDRESS_BITS_PER_WORD
-    }
+    /// Zeros the bits in the given range.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the range is invalid (start > end).
+    pub fn clear(&mut self, range: Range<usize>) {
+        assert!(
+            range.start <= range.end,
+            "Range ends before it starts; {} must be greater than {}",
+            range.start,
+            range.end
+        );
 
-    pub fn clear(&mut self, from_index: usize, mut to_index: usize) {
-        self.check_range(from_index, to_index);
+        let from_idx = range.start;
+        let mut to_idx = range.end;
 
-        if from_index == to_index {
+        if from_idx == to_idx {
             return;
         }
 
-        let start_word_index = self.word_index(from_index);
-        if start_word_index >= self.data.len() {
+        let start_word_idx = self.word_index(from_idx);
+        if start_word_idx >= self.data.len() {
             return;
         }
 
-        let mut end_word_index = self.word_index(to_index - 1);
-        if end_word_index >= self.data.len() {
-            to_index = self.len();
-            end_word_index = self.data.len() - 1;
+        let mut end_word_idx = self.word_index(to_idx - 1);
+        if end_word_idx >= self.data.len() {
+            to_idx = self.len();
+            end_word_idx = self.data.len() - 1;
         }
 
         let first_word_mask = u64::MAX.wrapping_shl(
-            from_index
+            from_idx
                 .try_into()
                 .expect("from_index shouldn't be larger than u32"),
         );
-        let last_word_mask = u64::MAX.wrapping_shr((64 - (to_index % 64)) as u32);
-        if start_word_index == end_word_index {
-            // Case 1: One word
-            self.data[start_word_index] &= !(first_word_mask & last_word_mask);
+        let last_word_mask = u64::MAX.wrapping_shr((64 - (to_idx % 64)) as u32);
+        if start_word_idx == end_word_idx {
+            // one word
+            self.data[start_word_idx] &= !(first_word_mask & last_word_mask);
         } else {
-            // Case 2: Multiple words
-            // Handle first word
-            self.data[start_word_index] &= !first_word_mask;
-
-            // Handle intermediate words, if any
-            for i in start_word_index + 1..end_word_index {
+            // multiple words
+            self.data[start_word_idx] &= !first_word_mask;
+            for i in (start_word_idx + 1)..end_word_idx {
                 self.data[i] = 0;
             }
-
-            // Handle last word
-            self.data[end_word_index] &= !last_word_mask;
+            self.data[end_word_idx] &= !last_word_mask;
         }
     }
 
-    /// Returns the maximum potential items in the BitSet. This will be
-    /// divisible by 64.
-    fn len(&self) -> usize {
-        self.data.len() * 64
-    }
-
     /// Returns the index of the first bit that is set to `false`
     /// that occurs on or after the specified starting index.
     pub fn next_clear_bit(&self, from_index: usize) -> usize {
@@ -103,9 +116,34 @@ impl BitSet {
         }
     }
 
+    #[inline]
+    fn word_index(&self, bit_index: usize) -> usize {
+        bit_index >> LOG2_BITS_PER_WORD
+    }
+
+    /// Sets the bit at the given index to true.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds. Check [`Self::len`] first
+    /// if you need to avoid this.
+    #[inline]
     pub fn set(&mut self, bit_index: usize) {
         self.data[bit_index / 64] |= 1u64 << (bit_index % 64);
     }
+
+    /// Returns the indices of all bits that are set to `true`.
+    pub fn iter_ones(&self) -> impl Iterator<Item = usize> {
+        (0..self.len()).filter_map(|i| if self.index(i) { Some(i) } else { None })
+    }
+
+    /// Returns the maximum number of items that could be in this `BitSet`.
+    ///
+    /// This will always be a multiple of 64.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.data.len() * 64
+    }
 }
 
 impl From<Vec<u64>> for BitSet {
diff --git a/azalea-physics/src/collision/discrete_voxel_shape.rs b/azalea-physics/src/collision/discrete_voxel_shape.rs
index a56c8085..033d1225 100644
--- a/azalea-physics/src/collision/discrete_voxel_shape.rs
+++ b/azalea-physics/src/collision/discrete_voxel_shape.rs
@@ -294,10 +294,8 @@ impl BitSetDiscreteVoxelShape {
     }
 
     fn clear_z_strip(&mut self, var1: u32, var2: u32, var3: u32, var4: u32) {
-        self.storage.clear(
-            self.get_index(var3, var4, var1),
-            self.get_index(var3, var4, var2),
-        );
+        self.storage
+            .clear(self.get_index(var3, var4, var1)..self.get_index(var3, var4, var2));
     }
 }
 
@@ -316,7 +314,7 @@ impl BitSetDiscreteVoxelShape {
     }
 
     fn is_full(&self, x: u32, y: u32, z: u32) -> bool {
-        self.storage.index(self.get_index(x, y, z))
+        self.storage.get(self.get_index(x, y, z)).unwrap_or(false)
     }
 
     fn is_full_wide(&self, x: u32, y: u32, z: u32) -> bool {