fix: use better types in HAMT hashing (#1865)

Using a u8 here means we can statically guarantee that no bitfield
access will panic due to a bounds error.

Author: Stebalien

ipld/hamt/src/bitfield.rs

@@ -69,20 +69,20 @@ impl Default for Bitfield {
 }
 
 impl Bitfield {
-    pub fn clear_bit(&mut self, idx: u32) {
+    pub fn clear_bit(&mut self, idx: u8) {
         let ai = idx / 64;
         let bi = idx % 64;
         self.0[ai as usize] &= u64::MAX - (1 << bi);
     }
 
-    pub fn test_bit(&self, idx: u32) -> bool {
+    pub fn test_bit(&self, idx: u8) -> bool {
         let ai = idx / 64;
         let bi = idx % 64;
 
         self.0[ai as usize] & (1 << bi) != 0
     }
 
-    pub fn set_bit(&mut self, idx: u32) {
+    pub fn set_bit(&mut self, idx: u8) {
         let ai = idx / 64;
         let bi = idx % 64;
 
@@ -106,14 +106,14 @@ impl Bitfield {
         Bitfield([0, 0, 0, 0])
     }
 
-    pub fn set_bits_le(self, bit: u32) -> Self {
+    pub fn set_bits_le(self, bit: u8) -> Self {
         if bit == 0 {
             return self;
         }
         self.set_bits_leq(bit - 1)
     }
 
-    pub fn set_bits_leq(mut self, bit: u32) -> Self {
+    pub fn set_bits_leq(mut self, bit: u8) -> Self {
         if bit < 64 {
             self.0[0] = set_bits_leq(self.0[0], bit);
         } else if bit < 128 {
@@ -135,7 +135,7 @@ impl Bitfield {
 }
 
 #[inline]
-fn set_bits_leq(v: u64, bit: u32) -> u64 {
+fn set_bits_leq(v: u64, bit: u8) -> u64 {
     (v as u128 | ((1u128 << (1 + bit)) - 1)) as u64
 }
 

ipld/hamt/src/hash_bits.rs

@@ -16,9 +16,10 @@ pub struct HashBits<'a> {
     pub consumed: u32,
 }
 
+// n must be less than 8
 #[inline]
-pub(crate) fn mkmask(n: u32) -> u32 {
-    ((1u64 << n) - 1) as u32
+pub(crate) fn mkmask(n: u32) -> u8 {
+    ((1u16 << n) - 1) as u8
 }
 
 impl<'a> HashBits<'a> {
@@ -36,7 +37,7 @@ impl<'a> HashBits<'a> {
 
     /// Returns next `i` bits of the hash and returns the value as an integer and returns
     /// Error when maximum depth is reached
-    pub fn next(&mut self, i: u32) -> Result<u32, Error> {
+    pub fn next(&mut self, i: u32) -> Result<u8, Error> {
         if i > 8 || i == 0 {
             return Err(Error::InvalidHashBitLen);
         }
@@ -49,11 +50,12 @@ impl<'a> HashBits<'a> {
         Ok(self.next_bits(std::cmp::min(i, maxi)))
     }
 
-    fn next_bits(&mut self, i: u32) -> u32 {
+    // `i` must be between 1 and 8, inclusive.
+    fn next_bits(&mut self, i: u32) -> u8 {
         let curbi = self.consumed / 8;
         let leftb = 8 - (self.consumed % 8);
 
-        let curb = self.b[curbi as usize] as u32;
+        let curb = self.b[curbi as usize];
         match i.cmp(&leftb) {
             Ordering::Equal => {
                 // bits to consume is equal to the bits remaining in the currently indexed byte
@@ -71,11 +73,11 @@ impl<'a> HashBits<'a> {
             }
             Ordering::Greater => {
                 // Consumes remaining bits and remaining bits from a recursive call
-                let mut out = (mkmask(leftb) & curb) as u64;
+                let mut out = mkmask(leftb) & curb;
                 out <<= i - leftb;
                 self.consumed += leftb;
-                out += self.next_bits(i - leftb) as u64;
-                out as u32
+                out += self.next_bits(i - leftb);
+                out
             }
         }
     }

ipld/hamt/src/node.rs

@@ -478,18 +478,18 @@ where
         Ok(())
     }
 
-    fn rm_child(&mut self, i: usize, idx: u32) -> Pointer<K, V, H, Ver> {
+    fn rm_child(&mut self, i: usize, idx: u8) -> Pointer<K, V, H, Ver> {
         self.bitfield.clear_bit(idx);
         self.pointers.remove(i)
     }
 
-    fn insert_child(&mut self, idx: u32, key: K, value: V) {
+    fn insert_child(&mut self, idx: u8, key: K, value: V) {
         let i = self.index_for_bit_pos(idx);
         self.bitfield.set_bit(idx);
         self.pointers.insert(i, Pointer::from_key_value(key, value))
     }
 
-    fn insert_child_dirty(&mut self, idx: u32, node: Box<Node<K, V, H, Ver>>) {
+    fn insert_child_dirty(&mut self, idx: u8, node: Box<Node<K, V, H, Ver>>) {
         let i = self.index_for_bit_pos(idx);
         self.bitfield.set_bit(idx);
         self.pointers.insert(i, Pointer::Dirty(node))
@@ -517,9 +517,9 @@ where
 }
 
 impl<K, V, H, Ver> Node<K, V, H, Ver> {
-    pub(crate) fn index_for_bit_pos(&self, bp: u32) -> usize {
+    pub(crate) fn index_for_bit_pos(&self, bp: u8) -> usize {
         let mask = Bitfield::zero().set_bits_le(bp);
-        assert_eq!(mask.count_ones(), bp as usize);
+        debug_assert_eq!(mask.count_ones(), bp as usize);
         mask.and(&self.bitfield).count_ones()
     }
 }