Place all nodes in an arena of sorts

src/lib.rs

@@ -53,18 +53,23 @@
 
 pub use h3ron;
 use h3ron::{H3Cell, Index};
-use std::{iter::FromIterator, mem::size_of, ops::Deref, ops::DerefMut};
+use std::{
+    iter::FromIterator,
+    mem::size_of,
+    ops::{Deref, DerefMut},
+};
 
 /// An efficient way to represent any portion(s) of Earth as a set of
 /// `H3` hexagons.
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone)]
 #[cfg_attr(
     feature = "serde-support",
     derive(serde::Serialize, serde::Deserialize)
 )]
 pub struct HexSet {
     /// All h3 0 base cell indices in the tree
-    nodes: Box<[Option<Node>]>,
+    root: Box<[Option<usize>]>,
+    arena: NodeArena,
 }
 
 impl HexSet {
@@ -74,7 +79,19 @@ impl HexSet {
     /// H3 cells.
     pub fn new() -> Self {
         Self {
-            nodes: vec![None; 122].into_boxed_slice(),
+            root: vec![None; 122].into_boxed_slice(),
+            arena: NodeArena::new(),
+        }
+    }
+
+    /// Constructs a new, empty `HexSet`.
+    ///
+    /// Incurs a single heap allocation to store all 122 resolution-0
+    /// H3 cells.
+    pub fn with_capacity(cap: usize) -> Self {
+        Self {
+            root: vec![None; 122].into_boxed_slice(),
+            arena: NodeArena::with_capacity(cap),
         }
     }
 
@@ -85,7 +102,11 @@ impl HexSet {
     /// significantly smaller than the number of source cells used to
     /// create the set.
     pub fn len(&self) -> usize {
-        self.nodes.iter().flatten().map(|node| node.len()).sum()
+        self.root
+            .iter()
+            .flatten()
+            .map(|arena_idx| self.arena[*arena_idx].len(&self.arena))
+            .sum()
     }
 
     /// Returns `true` if the set contains no cells.
@@ -97,12 +118,17 @@ impl HexSet {
     pub fn insert(&mut self, hex: H3Cell) {
         let base_cell = base(&hex);
         let digits = Digits::new(hex);
-        match self.nodes[base_cell as usize].as_mut() {
-            Some(node) => node.insert(digits),
+        match self.root[base_cell as usize].as_mut() {
+            Some(arena_idx) => {
+                let mut node = self.arena[*arena_idx].clone();
+                node.insert(&mut self.arena, digits);
+                self.arena[*arena_idx] = node;
+            }
             None => {
                 let mut node = Node::new();
-                node.insert(digits);
-                self.nodes[base_cell as usize] = Some(node);
+                node.insert(&mut self.arena, digits);
+                let arena_idx = self.arena.alloc(node);
+                self.root[base_cell as usize] = Some(arena_idx);
             }
         }
     }
@@ -120,26 +146,31 @@ impl HexSet {
     ///    hex due to 1 or 2.
     pub fn contains(&self, hex: &H3Cell) -> bool {
         let base_cell = base(hex);
-        match self.nodes[base_cell as usize].as_ref() {
-            Some(node) => {
+        match self.root[base_cell as usize] {
+            Some(arena_idx) => {
                 let digits = Digits::new(*hex);
-                node.contains(digits)
+                self.arena[arena_idx].contains(&self.arena, digits)
             }
             None => false,
         }
     }
 
+    /// Returns an iterator over all cells in this set.
+    pub fn cells(&self) -> impl IntoIterator<Item = H3Cell> {
+        [H3Cell::new(577164439745200127)]
+    }
+
     /// Returns the current memory use of this set.
     ///
     /// Note: The actual total may be higher than reported due to
     ///       memory alignment.
     pub fn mem_size(&self) -> usize {
         size_of::<Self>()
             + self
-                .nodes
+                .root
                 .iter()
                 .flatten()
-                .map(|n| n.mem_size())
+                .map(|arena_idx| self.arena[*arena_idx].mem_size(&self.arena))
                 .sum::<usize>()
     }
 }
@@ -170,56 +201,76 @@ impl<'a> FromIterator<&'a H3Cell> for HexSet {
     }
 }
 
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone)]
 #[cfg_attr(
     feature = "serde-support",
     derive(serde::Serialize, serde::Deserialize)
 )]
-struct Node(Box<[Option<Node>; 7]>);
+struct Node([Option<usize>; 7]);
 
 impl Node {
-    fn mem_size(&self) -> usize {
-        size_of::<Self>() + self.iter().flatten().map(|n| n.mem_size()).sum::<usize>()
+    fn mem_size(&self, arena: &NodeArena) -> usize {
+        size_of::<Self>()
+            + self
+                .iter()
+                .flatten()
+                .map(|arena_idx| arena[*arena_idx].mem_size(arena))
+                .sum::<usize>()
     }
 
     fn new() -> Self {
-        Self(Box::new([None, None, None, None, None, None, None]))
+        Self([None, None, None, None, None, None, None])
     }
 
-    fn len(&self) -> usize {
+    fn len(&self, arena: &NodeArena) -> usize {
         if self.is_full() {
             1
         } else {
-            self.iter().flatten().map(|child| child.len()).sum()
+            self.iter()
+                .flatten()
+                .map(|arena_idx| arena[*arena_idx].len(arena))
+                .sum()
         }
     }
 
-    fn insert(&mut self, mut digits: Digits) {
+    fn insert(&mut self, arena: &mut NodeArena, mut digits: Digits) {
         match digits.next() {
             Some(digit) => match self[digit as usize].as_mut() {
-                Some(node) => node.insert(digits),
+                Some(arena_idx) => {
+                    let mut node = arena[*arena_idx].clone();
+                    node.insert(arena, digits);
+                    arena[*arena_idx] = node;
+                }
                 None => {
                     let mut node = Node::new();
-                    node.insert(digits);
-                    self[digit as usize] = Some(node);
+                    node.insert(arena, digits);
+                    let arena_idx = arena.alloc(node);
+                    self[digit as usize] = Some(arena_idx);
                 }
             },
-            None => *self.0 = [None, None, None, None, None, None, None],
+            None => self.0 = [None, None, None, None, None, None, None],
         };
-        self.coalesce();
+        self.coalesce(arena);
     }
 
-    fn coalesce(&mut self) {
-        if let [Some(n0), Some(n1), Some(n2), Some(n3), Some(n4), Some(n5), Some(n6)] = &*self.0 {
-            if n0.is_full()
-                && n1.is_full()
-                && n2.is_full()
-                && n3.is_full()
-                && n4.is_full()
-                && n5.is_full()
-                && n6.is_full()
+    fn coalesce(&mut self, arena: &mut NodeArena) {
+        if let [Some(n0), Some(n1), Some(n2), Some(n3), Some(n4), Some(n5), Some(n6)] = self.0 {
+            if arena[n0].is_full()
+                && arena[n1].is_full()
+                && arena[n2].is_full()
+                && arena[n3].is_full()
+                && arena[n4].is_full()
+                && arena[n5].is_full()
+                && arena[n6].is_full()
             {
-                *self.0 = [None, None, None, None, None, None, None]
+                arena.free(n0);
+                arena.free(n1);
+                arena.free(n2);
+                arena.free(n3);
+                arena.free(n4);
+                arena.free(n5);
+                arena.free(n6);
+                self.0 = [None, None, None, None, None, None, None]
             }
         };
     }
@@ -228,7 +279,7 @@ impl Node {
         self.iter().all(|c| c.is_none())
     }
 
-    fn contains(&self, mut digits: Digits) -> bool {
+    fn contains(&self, arena: &NodeArena, mut digits: Digits) -> bool {
         if self.is_full() {
             return true;
         }
@@ -237,7 +288,7 @@ impl Node {
             Some(digit) => {
                 // TODO check if this node is "full"
                 match &self[digit as usize] {
-                    Some(node) => node.contains(digits),
+                    Some(arena_idx) => arena[*arena_idx].contains(arena, digits),
                     None => false,
                 }
             }
@@ -249,7 +300,7 @@ impl Node {
 }
 
 impl Deref for Node {
-    type Target = [Option<Node>];
+    type Target = [Option<usize>];
 
     fn deref(&self) -> &Self::Target {
         &self.0[..]
@@ -262,6 +313,67 @@ impl DerefMut for Node {
     }
 }
 
+#[derive(Clone)]
+#[cfg_attr(
+    feature = "serde-support",
+    derive(serde::Serialize, serde::Deserialize)
+)]
+struct NodeArena {
+    nodes: Vec<Node>,
+    free: Vec<usize>,
+}
+
+impl NodeArena {
+    fn new() -> Self {
+        Self {
+            nodes: Vec::new(),
+            free: Vec::new(),
+        }
+    }
+
+    fn with_capacity(cap: usize) -> Self {
+        Self {
+            nodes: Vec::with_capacity(cap),
+            free: Vec::new(),
+        }
+    }
+
+    fn alloc(&mut self, node: Node) -> usize {
+        match self.free.pop() {
+            Some(idx) => {
+                self.nodes[idx] = node;
+                idx
+            }
+            None => {
+                self.nodes.push(node);
+                self.nodes.len() - 1
+            }
+        }
+    }
+
+    fn free(&mut self, node_idx: usize) {
+        for node in self.nodes[node_idx].0 {
+            if let Some(idx) = node {
+                self.free(idx)
+            }
+        }
+        self.free.push(node_idx);
+    }
+}
+
+impl std::ops::Index<usize> for NodeArena {
+    type Output = Node;
+    fn index(&self, index: usize) -> &Node {
+        &self.nodes[index]
+    }
+}
+
+impl std::ops::IndexMut<usize> for NodeArena {
+    fn index_mut(&mut self, index: usize) -> &mut Node {
+        &mut self.nodes[index]
+    }
+}
+
 struct Digits {
     digits: u64,
     remaining: u8,
@@ -328,16 +440,4 @@ mod tests {
             assert_eq!(&&digits, ref_digits);
         }
     }
-
-    #[test]
-    fn test_mem_size() {
-        // Sanity check that `Option<Node>` behaves the same as
-        // `Option<Box<[Option<Node>; 7]>>` in that it uses `NULL` to
-        // represent the `None` variant.
-        assert_eq!(size_of::<Option<Node>>(), size_of::<*const ()>());
-        assert_eq!(
-            size_of::<Option<Node>>(),
-            size_of::<Option<Box<[Option<Node>; 7]>>>()
-        );
-    }
 }

tests/tests.rs

@@ -129,7 +129,10 @@ fn test_compaction() {
         from_indicies(regions::nocompact::US915);
     let gulf_of_mexico = H3Cell::from_coordinate(&coord! {x: -83.101920, y: 28.128096}, 0).unwrap();
     assert_eq!(us915_tree.len(), us915_nocompact_tree.len());
-    assert!(us915_tree == us915_nocompact_tree);
+    assert!(us915_tree
+        .cells()
+        .into_iter()
+        .eq(us915_nocompact_tree.cells().into_iter()));
     assert!(us915_nocompact_tree.len() < us915_nocompact_cells.len());
     assert!(us915_nocompact_cells
         .iter()