Add interning arena

hendrikvanantwerpen · Mar 13, 2023 · Mar 13, 2023 · Mar 13, 2023 · Mar 13, 2023 · Mar 13, 2023
commit 5861201a1f26c4f34ef63d2f452f6cc995e30651
@@ -30,6 +30,8 @@
 //! [`StackGraph`]: ../graph/struct.StackGraph.html
 
 use std::cell::Cell;
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::hash::Hasher;
@@ -195,7 +197,7 @@ impl<T> Arena<T> {
     pub fn get(&self, handle: Handle<T>) -> &T {
         unsafe { std::mem::transmute(&self.items[handle.as_usize()]) }
     }
-    ///
+
     /// Dereferences a handle to an instance owned by this arena, returning a mutable reference to
     /// it.
     pub fn get_mut(&mut self, handle: Handle<T>) -> &mut T {
@@ -222,6 +224,66 @@ impl<T> Arena<T> {
     }
 }
 
+/// An interning arena allocation.  Equal handles means equal elements.
+pub struct InterningArena<T>
+where
+    T: Clone + Eq + Hash,
+{
+    arena: Arena<T>,
+    index: HashMap<T, Handle<T>>,
+}
+
+impl<T> InterningArena<T>
+where
+    T: Clone + Eq + Hash,
+{
+    /// Creates a new interning arena.
+    pub fn new() -> Self {
+        Self {
+            arena: Arena::new(),
+            index: HashMap::new(),
+        }
+    }
+
+    /// Adds a new instance to this arena, returning a stable handle to it.
+    ///
+    /// Note that we deduplicate instances of `T`.  If you add two instances that
+    /// have the same content, you will get identical handles for each one.
+    pub fn add(&mut self, item: T) -> Handle<T> {
+        match self.index.entry(item) {
+            Entry::Occupied(entry) => *entry.get(),
+            Entry::Vacant(entry) => {
+                let handle = self.arena.add(entry.key().clone());
+                entry.insert(handle);
+                handle
+            }
+        }
+    }
+
+    /// Dereferences a handle to an instance owned by this arena, returning a reference to it.
+    pub fn get(&self, handle: Handle<T>) -> &T {
+        self.arena.get(handle)
+    }
+
+    /// Dereferences a handle to an instance owned by this arena, returning a mutable reference to
+    /// it.
+    pub fn get_mut(&mut self, handle: Handle<T>) -> &mut T {
+        self.arena.get_mut(handle)
+    }
+
+    /// Returns an iterator of all of the handles in this arena.  (Note that this iterator does not
+    /// retain a reference to the arena!)
+    pub fn iter_handles(&self) -> impl Iterator<Item = Handle<T>> {
+        self.arena.iter_handles()
+    }
+
+    /// Returns the number of instances stored in this arena.
+    #[inline(always)]
+    pub fn len(&self) -> usize {
+        self.arena.len()
+    }
+}
+
 //-------------------------------------------------------------------------------------------------
 // Supplemental arenas
 

@@ -8,6 +8,7 @@
 use stack_graphs::arena::Arena;
 use stack_graphs::arena::Deque;
 use stack_graphs::arena::DequeArena;
+use stack_graphs::arena::InterningArena;
 use stack_graphs::arena::List;
 use stack_graphs::arena::ListArena;
 use stack_graphs::arena::ReversibleList;
@@ -229,3 +230,17 @@ fn can_compare_deques() {
     deque10.ensure_backwards(&mut arena);
     assert_eq!(deque1.cmp(&mut arena, deque10), Ordering::Less);
 }
+
+#[test]
+fn can_allocate_in_interning_arena() {
+    let mut arena = InterningArena::new();
+    let hello1 = arena.add("hello".to_string());
+    let hello2 = arena.add("hello".to_string());
+    let there = arena.add("there".to_string());
+    assert_eq!(hello1, hello2);
+    assert_ne!(hello1, there);
+    assert_ne!(hello2, there);
+    assert_eq!(arena.get(hello1), arena.get(hello2));
+    assert_ne!(arena.get(hello1), arena.get(there));
+    assert_ne!(arena.get(hello2), arena.get(there));
+}