Pass through u128 hashes to the hashmap

src/accelerate.rs

@@ -1,8 +1,9 @@
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-use rustc_hash::FxHashMap;
 use parking_lot::{MappedRwLockReadGuard, Mutex, RwLock, RwLockReadGuard};
 
+use crate::passthroughhasher::PassthroughHashMap;
+
 /// The global list of currently alive accelerators.
 static ACCELERATORS: RwLock<(usize, Vec<Accelerator>)> = RwLock::new((0, Vec::new()));
 
@@ -12,7 +13,7 @@ static ID: AtomicUsize = AtomicUsize::new(0);
 /// The type of each individual accelerator.
 ///
 /// Maps from call hashes to return hashes.
-type Accelerator = Mutex<FxHashMap<u128, u128>>;
+type Accelerator = Mutex<PassthroughHashMap<u128, u128>>;
 
 /// Generate a new accelerator.
 pub fn id() -> usize {
@@ -58,6 +59,6 @@ pub fn get(id: usize) -> Option<MappedRwLockReadGuard<'static, Accelerator>> {
 fn resize(len: usize) {
     let mut pair = ACCELERATORS.write();
     if len > pair.1.len() {
-        pair.1.resize_with(len, || Mutex::new(FxHashMap::default()));
+        pair.1.resize_with(len, || Mutex::new(PassthroughHashMap::default()));
     }
 }

src/constraint.rs

@@ -2,9 +2,9 @@ use std::collections::hash_map::Entry;
 use std::hash::Hash;
 
 use parking_lot::Mutex;
-use rustc_hash::FxHashMap;
 
 use crate::Track;
+use crate::passthroughhasher::PassthroughHashMap;
 use crate::track::{Call, Sink};
 
 /// Records calls performed on a trackable type.
@@ -82,7 +82,7 @@ pub struct CallSequence<C> {
     /// The raw calls. In order, but deduplicated via the `map`.
     vec: Vec<Option<(C, u128)>>,
     /// A map from hashes of calls to the indices in the vector.
-    map: FxHashMap<u128, usize>,
+    map: PassthroughHashMap<u128, usize>,
     /// A cursor for iteration in `Self::next`.
     cursor: usize,
 }
@@ -92,7 +92,7 @@ impl<C> CallSequence<C> {
     pub fn new() -> Self {
         Self {
             vec: Vec::new(),
-            map: FxHashMap::default(),
+            map: PassthroughHashMap::default(),
             cursor: 0,
         }
     }

src/lib.rs

@@ -93,6 +93,7 @@ mod constraint;
 mod hash;
 mod input;
 mod memoize;
+mod passthroughhasher;
 mod track;
 mod tree;
 

src/passthroughhasher.rs

@@ -0,0 +1,40 @@
+use std::collections::HashMap;
+use std::hash::{BuildHasher, Hasher};
+
+/// Hash Map that re-uses the u128 as the hash value
+pub(crate) type PassthroughHashMap<Key, Value> =
+    HashMap<Key, Value, BuildPassthroughHasher>;
+
+#[derive(Copy, Clone, Default)]
+pub(crate) struct BuildPassthroughHasher;
+
+#[derive(Default)]
+pub(crate) struct PassthroughHasher {
+    value: u64,
+}
+
+impl Hasher for PassthroughHasher {
+    #[inline(always)]
+    fn finish(&self) -> u64 {
+        self.value
+    }
+
+    #[inline]
+    fn write(&mut self, _bytes: &[u8]) {
+        unimplemented!("Unsupported operation")
+    }
+
+    #[inline]
+    fn write_u128(&mut self, i: u128) {
+        // truncating conversion
+        self.value = i as u64;
+    }
+}
+
+impl BuildHasher for BuildPassthroughHasher {
+    type Hasher = PassthroughHasher;
+    #[inline]
+    fn build_hasher(&self) -> PassthroughHasher {
+        PassthroughHasher::default()
+    }
+}

src/tree.rs

@@ -5,6 +5,7 @@ use rustc_hash::FxHashMap;
 use slab::Slab;
 
 use crate::constraint::CallSequence;
+use crate::passthroughhasher::PassthroughHashMap;
 
 /// A tree data structure that associates a value with a key hash and a sequence
 /// of (call, return hash) pairs.
@@ -18,7 +19,7 @@ pub struct CallTree<C, T> {
     /// Leaf nodes, directly storing outputs.
     leaves: Slab<LeafNode<T>>,
     /// The initial node for the given key hash.
-    start: FxHashMap<u128, NodeId>,
+    start: PassthroughHashMap<u128, NodeId>,
     /// Maps from parent nodes to child nodes. The key is a pair of an inner
     /// node ID and a return hash for that call. The value is the node to
     /// transition to.
@@ -50,8 +51,8 @@ impl<C, T> CallTree<C, T> {
         Self {
             inner: Slab::new(),
             leaves: Slab::new(),
+            start: PassthroughHashMap::default(),
             edges: FxHashMap::default(),
-            start: FxHashMap::default(),
         }
     }
 }