Add bit storage

1 files changed, 48 insertions, 5 deletions

diff --git a/azalea-world/src/bit_storage.rs b/azalea-world/src/bit_storage.rs
index b2202e48..5ced67fe 100644
--- a/azalea-world/src/bit_storage.rs
+++ b/azalea-world/src/bit_storage.rs
@@ -1,4 +1,4 @@
-use std::{error::Error, fmt};
+use std::{error::Error, fmt, ops::Index};
 
 // this is from minecraft's code
 // yeah idk either
@@ -72,9 +72,9 @@ const MAGIC: [(i32, i32, i32); 64] = [
 /// A compact list of integers with the given number of bits per entry.
 pub struct BitStorage {
     data: Vec<u64>,
-    bits: u32,
+    bits: usize,
     mask: u64,
-    size: u32,
+    size: usize,
     values_per_long: u8,
     divide_mul: i32,
     divide_add: i32,
@@ -101,7 +101,7 @@ impl Error for BitStorageError {}
 impl BitStorage {
     /// Create a new BitStorage with the given number of bits per entry.
     /// `size` is the number of entries in the BitStorage.
-    pub fn new(bits: u32, size: u32, data: Option<Vec<u64>>) -> Result<Self, BitStorageError> {
+    pub fn new(bits: usize, size: usize, data: Option<Vec<u64>>) -> Result<Self, BitStorageError> {
         let values_per_long = 64 / bits;
         let magic_index = values_per_long - 1;
         let (divide_mul, divide_add, divide_shift) = MAGIC[magic_index as usize];
@@ -132,6 +132,45 @@ impl BitStorage {
             divide_shift,
         })
     }
+
+    pub fn cell_index(&self, index: u64) -> usize {
+        let first = self.divide_mul as u64;
+        let second = self.divide_mul as u64;
+
+        (index * first + second >> 32 >> self.divide_shift)
+            .try_into()
+            .unwrap()
+    }
+
+    pub fn get(&self, index: usize) -> u64 {
+        // Validate.inclusiveBetween(0L, (long)(this.size - 1), (long)var1);
+        // int var2 = this.cellIndex(var1);
+        // long var3 = this.data[var2];
+        // int var5 = (var1 - var2 * this.valuesPerLong) * this.bits;
+        // return (int)(var3 >> var5 & this.mask);
+
+        assert!(index <= self.size - 1);
+        let cell_index = self.cell_index(index as u64);
+        let cell = &self.data[cell_index as usize];
+        let bit_index = (index - cell_index * self.values_per_long as usize) * self.bits;
+        cell >> bit_index & self.mask
+    }
+
+    pub fn set(&mut self, index: usize, value: u64) {
+        // Validate.inclusiveBetween(0L, (long)(this.size - 1), (long)var1);
+        // Validate.inclusiveBetween(0L, this.mask, (long)var2);
+        // int var3 = this.cellIndex(var1);
+        // long var4 = this.data[var3];
+        // int var6 = (var1 - var3 * this.valuesPerLong) * this.bits;
+        // this.data[var3] = var4 & ~(this.mask << var6) | ((long)var2 & this.mask) << var6;
+
+        assert!(index <= self.size - 1);
+        assert!(value <= self.mask);
+        let cell_index = self.cell_index(index as u64);
+        let cell = &mut self.data[cell_index as usize];
+        let bit_index = (index - cell_index * self.values_per_long as usize) * self.bits;
+        *cell = *cell & !(self.mask << bit_index) | (value & self.mask) << bit_index;
+    }
 }
 
 #[cfg(test)]
@@ -144,6 +183,10 @@ mod tests {
             1, 2, 2, 3, 4, 4, 5, 6, 6, 4, 8, 0, 7, 4, 3, 13, 15, 16, 9, 14, 10, 12, 0, 2,
         ];
         let expected_compact: [u64; 2] = [0x0020863148418841, 0x01018A7260F68C87];
-        let storage = BitStorage::new(5, 10, None).unwrap();
+        let storage = BitStorage::new(5, data.len(), Some(expected_compact.to_vec())).unwrap();
+
+        for (i, expected) in data.iter().enumerate() {
+            assert_eq!(storage.get(i), *expected);
+        }
     }
 }