Extract node_storage as well

agb-hashmap/src/lib.rs

@@ -17,21 +17,22 @@
 
 extern crate alloc;
 
-use alloc::{alloc::Global, vec::Vec};
+use alloc::alloc::Global;
 use core::{
     alloc::Allocator,
     borrow::Borrow,
     hash::{BuildHasher, BuildHasherDefault, Hash, Hasher},
     iter::FromIterator,
-    mem,
     ops::Index,
 };
 
 use rustc_hash::FxHasher;
 
 mod node;
+mod node_storage;
 
 use node::Node;
+use node_storage::NodeStorage;
 
 type HashType = u32;
 
@@ -249,7 +250,7 @@ impl<K, V, ALLOCATOR: ClonableAllocator> HashMap<K, V, ALLOCATOR> {
 
     /// An iterator visiting all key-value pairs in an arbitrary order, with mutable references to the values
     pub fn iter_mut(&mut self) -> impl Iterator<Item = (&'_ K, &'_ mut V)> {
-        self.nodes.nodes.iter_mut().filter_map(Node::key_value_mut)
+        self.nodes.iter_mut().filter_map(Node::key_value_mut)
     }
 
     /// Retains only the elements specified by the predicate `f`.
@@ -285,11 +286,6 @@ impl<K, V> Default for HashMap<K, V> {
     }
 }
 
-const fn fast_mod(len: usize, hash: HashType) -> usize {
-    debug_assert!(len.is_power_of_two(), "Length must be a power of 2");
-    (hash as usize) & (len - 1)
-}
-
 impl<K, V, ALLOCATOR: ClonableAllocator> HashMap<K, V, ALLOCATOR>
 where
     K: Eq + Hash,
@@ -331,7 +327,7 @@ where
             self.nodes.insert_new(key, value, hash)
         };
 
-        self.nodes.nodes[location].value_mut().unwrap()
+        self.nodes.node_at_mut(location).value_mut().unwrap()
     }
 
     /// Returns `true` if the map contains a value for the specified key.
@@ -354,7 +350,7 @@ where
 
         self.nodes
             .location(key, hash)
-            .and_then(|location| self.nodes.nodes[location].key_value_ref())
+            .and_then(|location| self.nodes.node_at(location).key_value_ref())
     }
 
     /// Returns a reference to the value corresponding to the key. Returns [`None`] if there is
@@ -401,7 +397,7 @@ where
         let hash = self.hash(key);
 
         if let Some(location) = self.nodes.location(key, hash) {
-            self.nodes.nodes[location].value_mut()
+            self.nodes.node_at_mut(location).value_mut()
         } else {
             None
         }
@@ -466,7 +462,7 @@ impl<'a, K, V, ALLOCATOR: ClonableAllocator> Iterator for Iter<'a, K, V, ALLOCAT
                 return None;
             }
 
-            let node = &self.map.nodes.nodes[self.at];
+            let node = &self.map.nodes.node_at(self.at);
             self.at += 1;
 
             if node.has_value() {
@@ -516,7 +512,7 @@ impl<K, V, ALLOCATOR: ClonableAllocator> Iterator for IterOwned<K, V, ALLOCATOR>
                 return None;
             }
 
-            let maybe_kv = self.map.nodes.nodes[self.at].take_key_value();
+            let maybe_kv = self.map.nodes.node_at_mut(self.at).take_key_value();
             self.at += 1;
 
             if let Some((k, v, _)) = maybe_kv {
@@ -572,7 +568,7 @@ impl<'a, K: 'a, V: 'a, ALLOCATOR: ClonableAllocator> OccupiedEntry<'a, K, V, ALL
 
     /// Gets a reference to the value in the entry.
     pub fn get(&self) -> &V {
-        self.map.nodes.nodes[self.location].value_ref().unwrap()
+        self.map.nodes.node_at(self.location).value_ref().unwrap()
     }
 
     /// Gets a mutable reference to the value in the entry.
@@ -582,7 +578,11 @@ impl<'a, K: 'a, V: 'a, ALLOCATOR: ClonableAllocator> OccupiedEntry<'a, K, V, ALL
     ///
     /// [`into_mut`]: Self::into_mut
     pub fn get_mut(&mut self) -> &mut V {
-        self.map.nodes.nodes[self.location].value_mut().unwrap()
+        self.map
+            .nodes
+            .node_at_mut(self.location)
+            .value_mut()
+            .unwrap()
     }
 
     /// Converts the `OccupiedEntry` into a mutable reference to the value in the entry with
@@ -592,12 +592,19 @@ impl<'a, K: 'a, V: 'a, ALLOCATOR: ClonableAllocator> OccupiedEntry<'a, K, V, ALL
     ///
     /// [`get_mut`]: Self::get_mut
     pub fn into_mut(self) -> &'a mut V {
-        self.map.nodes.nodes[self.location].value_mut().unwrap()
+        self.map
+            .nodes
+            .node_at_mut(self.location)
+            .value_mut()
+            .unwrap()
     }
 
     /// Sets the value of the entry and returns the entry's old value.
     pub fn insert(&mut self, value: V) -> V {
-        self.map.nodes.nodes[self.location].replace_value(value)
+        self.map
+            .nodes
+            .node_at_mut(self.location)
+            .replace_value(value)
     }
 
     /// Takes the value out of the entry and returns it.
@@ -783,180 +790,12 @@ const fn number_before_resize(capacity: usize) -> usize {
     capacity * 85 / 100
 }
 
-struct NodeStorage<K, V, ALLOCATOR: Allocator = Global> {
-    nodes: Vec<Node<K, V>, ALLOCATOR>,
-    max_distance_to_initial_bucket: i32,
-
-    number_of_items: usize,
-    max_number_before_resize: usize,
-}
-
-impl<K, V, ALLOCATOR: ClonableAllocator> NodeStorage<K, V, ALLOCATOR> {
-    fn with_size_in(capacity: usize, alloc: ALLOCATOR) -> Self {
-        assert!(capacity.is_power_of_two(), "Capacity must be a power of 2");
-
-        let mut nodes = Vec::with_capacity_in(capacity, alloc);
-        for _ in 0..capacity {
-            nodes.push(Default::default());
-        }
-
-        Self {
-            nodes,
-            max_distance_to_initial_bucket: 0,
-            number_of_items: 0,
-            max_number_before_resize: number_before_resize(capacity),
-        }
-    }
-
-    fn allocator(&self) -> &ALLOCATOR {
-        self.nodes.allocator()
-    }
-
-    fn capacity(&self) -> usize {
-        self.max_number_before_resize
-    }
-
-    fn backing_vec_size(&self) -> usize {
-        self.nodes.len()
-    }
-
-    fn len(&self) -> usize {
-        self.number_of_items
-    }
-
-    fn insert_new(&mut self, key: K, value: V, hash: HashType) -> usize {
-        debug_assert!(
-            self.capacity() > self.len(),
-            "Do not have space to insert into len {} with {}",
-            self.backing_vec_size(),
-            self.len()
-        );
-
-        let mut new_node = Node::new_with(key, value, hash);
-        let mut inserted_location = usize::MAX;
-
-        loop {
-            let location = fast_mod(
-                self.backing_vec_size(),
-                new_node.hash() + new_node.distance() as HashType,
-            );
-            let current_node = &mut self.nodes[location];
-
-            if current_node.has_value() {
-                if current_node.distance() <= new_node.distance() {
-                    mem::swap(&mut new_node, current_node);
-
-                    if inserted_location == usize::MAX {
-                        inserted_location = location;
-                    }
-                }
-            } else {
-                self.nodes[location] = new_node;
-                if inserted_location == usize::MAX {
-                    inserted_location = location;
-                }
-                break;
-            }
-
-            new_node.increment_distance();
-            self.max_distance_to_initial_bucket =
-                new_node.distance().max(self.max_distance_to_initial_bucket);
-        }
-
-        self.number_of_items += 1;
-        inserted_location
-    }
-
-    fn retain<F>(&mut self, mut f: F)
-    where
-        F: FnMut(&K, &mut V) -> bool,
-    {
-        let num_nodes = self.nodes.len();
-        let mut i = 0;
-
-        while i < num_nodes {
-            let node = &mut self.nodes[i];
-
-            if let Some((k, v)) = node.key_value_mut() {
-                if !f(k, v) {
-                    self.remove_from_location(i);
-
-                    // Need to continue before adding 1 to i because remove from location could
-                    // put the element which was next into the ith location in the nodes array,
-                    // so we need to check if that one needs removing too.
-                    continue;
-                }
-            }
-
-            i += 1;
-        }
-    }
-
-    fn remove_from_location(&mut self, location: usize) -> V {
-        let mut current_location = location;
-        self.number_of_items -= 1;
-
-        loop {
-            let next_location =
-                fast_mod(self.backing_vec_size(), (current_location + 1) as HashType);
-
-            // if the next node is empty, or the next location has 0 distance to initial bucket then
-            // we can clear the current node
-            if !self.nodes[next_location].has_value() || self.nodes[next_location].distance() == 0 {
-                return self.nodes[current_location].take_key_value().unwrap().1;
-            }
-
-            self.nodes.swap(current_location, next_location);
-            self.nodes[current_location].decrement_distance();
-            current_location = next_location;
-        }
-    }
-
-    fn location<Q>(&self, key: &Q, hash: HashType) -> Option<usize>
-    where
-        K: Borrow<Q>,
-        Q: Eq + ?Sized,
-    {
-        for distance_to_initial_bucket in 0..(self.max_distance_to_initial_bucket + 1) {
-            let location = fast_mod(
-                self.nodes.len(),
-                hash + distance_to_initial_bucket as HashType,
-            );
-
-            let node = &self.nodes[location];
-            if let Some(node_key_ref) = node.key_ref() {
-                if node_key_ref.borrow() == key {
-                    return Some(location);
-                }
-            } else {
-                return None;
-            }
-        }
-
-        None
-    }
-
-    fn resized_to(&mut self, new_size: usize) -> Self {
-        let mut new_node_storage = Self::with_size_in(new_size, self.allocator().clone());
-
-        for mut node in self.nodes.drain(..) {
-            if let Some((key, value, hash)) = node.take_key_value() {
-                new_node_storage.insert_new(key, value, hash);
-            }
-        }
-
-        new_node_storage
-    }
-
-    fn replace_at_location(&mut self, location: usize, key: K, value: V) -> V {
-        self.nodes[location].replace(key, value).1
-    }
-}
-
 #[cfg(test)]
 mod test {
     use core::cell::RefCell;
 
+    use alloc::vec::Vec;
+
     use super::*;
 
     #[test]

agb-hashmap/src/node_storage.rs

@@ -0,0 +1,192 @@
+use core::{alloc::Allocator, borrow::Borrow, mem};
+
+use alloc::{alloc::Global, vec::Vec};
+
+use crate::{node::Node, number_before_resize, ClonableAllocator, HashType};
+
+pub(crate) struct NodeStorage<K, V, ALLOCATOR: Allocator = Global> {
+    nodes: Vec<Node<K, V>, ALLOCATOR>,
+    max_distance_to_initial_bucket: i32,
+
+    number_of_items: usize,
+    max_number_before_resize: usize,
+}
+
+impl<K, V, ALLOCATOR: ClonableAllocator> NodeStorage<K, V, ALLOCATOR> {
+    pub(crate) fn with_size_in(capacity: usize, alloc: ALLOCATOR) -> Self {
+        assert!(capacity.is_power_of_two(), "Capacity must be a power of 2");
+
+        let mut nodes = Vec::with_capacity_in(capacity, alloc);
+        for _ in 0..capacity {
+            nodes.push(Default::default());
+        }
+
+        Self {
+            nodes,
+            max_distance_to_initial_bucket: 0,
+            number_of_items: 0,
+            max_number_before_resize: number_before_resize(capacity),
+        }
+    }
+
+    pub(crate) fn allocator(&self) -> &ALLOCATOR {
+        self.nodes.allocator()
+    }
+
+    pub(crate) fn capacity(&self) -> usize {
+        self.max_number_before_resize
+    }
+
+    pub(crate) fn backing_vec_size(&self) -> usize {
+        self.nodes.len()
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.number_of_items
+    }
+
+    pub(crate) fn insert_new(&mut self, key: K, value: V, hash: HashType) -> usize {
+        debug_assert!(
+            self.capacity() > self.len(),
+            "Do not have space to insert into len {} with {}",
+            self.backing_vec_size(),
+            self.len()
+        );
+
+        let mut new_node = Node::new_with(key, value, hash);
+        let mut inserted_location = usize::MAX;
+
+        loop {
+            let location = fast_mod(
+                self.backing_vec_size(),
+                new_node.hash() + new_node.distance() as HashType,
+            );
+            let current_node = &mut self.nodes[location];
+
+            if current_node.has_value() {
+                if current_node.distance() <= new_node.distance() {
+                    mem::swap(&mut new_node, current_node);
+
+                    if inserted_location == usize::MAX {
+                        inserted_location = location;
+                    }
+                }
+            } else {
+                self.nodes[location] = new_node;
+                if inserted_location == usize::MAX {
+                    inserted_location = location;
+                }
+                break;
+            }
+
+            new_node.increment_distance();
+            self.max_distance_to_initial_bucket =
+                new_node.distance().max(self.max_distance_to_initial_bucket);
+        }
+
+        self.number_of_items += 1;
+        inserted_location
+    }
+
+    pub(crate) fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        let num_nodes = self.nodes.len();
+        let mut i = 0;
+
+        while i < num_nodes {
+            let node = &mut self.nodes[i];
+
+            if let Some((k, v)) = node.key_value_mut() {
+                if !f(k, v) {
+                    self.remove_from_location(i);
+
+                    // Need to continue before adding 1 to i because remove from location could
+                    // put the element which was next into the ith location in the nodes array,
+                    // so we need to check if that one needs removing too.
+                    continue;
+                }
+            }
+
+            i += 1;
+        }
+    }
+
+    pub(crate) fn remove_from_location(&mut self, location: usize) -> V {
+        let mut current_location = location;
+        self.number_of_items -= 1;
+
+        loop {
+            let next_location =
+                fast_mod(self.backing_vec_size(), (current_location + 1) as HashType);
+
+            // if the next node is empty, or the next location has 0 distance to initial bucket then
+            // we can clear the current node
+            if !self.nodes[next_location].has_value() || self.nodes[next_location].distance() == 0 {
+                return self.nodes[current_location].take_key_value().unwrap().1;
+            }
+
+            self.nodes.swap(current_location, next_location);
+            self.nodes[current_location].decrement_distance();
+            current_location = next_location;
+        }
+    }
+
+    pub(crate) fn location<Q>(&self, key: &Q, hash: HashType) -> Option<usize>
+    where
+        K: Borrow<Q>,
+        Q: Eq + ?Sized,
+    {
+        for distance_to_initial_bucket in 0..(self.max_distance_to_initial_bucket + 1) {
+            let location = fast_mod(
+                self.nodes.len(),
+                hash + distance_to_initial_bucket as HashType,
+            );
+
+            let node = &self.nodes[location];
+            if let Some(node_key_ref) = node.key_ref() {
+                if node_key_ref.borrow() == key {
+                    return Some(location);
+                }
+            } else {
+                return None;
+            }
+        }
+
+        None
+    }
+
+    pub(crate) fn resized_to(&mut self, new_size: usize) -> Self {
+        let mut new_node_storage = Self::with_size_in(new_size, self.allocator().clone());
+
+        for mut node in self.nodes.drain(..) {
+            if let Some((key, value, hash)) = node.take_key_value() {
+                new_node_storage.insert_new(key, value, hash);
+            }
+        }
+
+        new_node_storage
+    }
+
+    pub(crate) fn replace_at_location(&mut self, location: usize, key: K, value: V) -> V {
+        self.nodes[location].replace(key, value).1
+    }
+
+    pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = &mut Node<K, V>> {
+        self.nodes.iter_mut()
+    }
+
+    pub(crate) fn node_at(&self, at: usize) -> &Node<K, V> {
+        &self.nodes[at]
+    }
+
+    pub(crate) fn node_at_mut(&mut self, at: usize) -> &mut Node<K, V> {
+        &mut self.nodes[at]
+    }
+}
+
+const fn fast_mod(len: usize, hash: HashType) -> usize {
+    debug_assert!(len.is_power_of_two(), "Length must be a power of 2");
+    (hash as usize) & (len - 1)
+}