[storage/index] Optimize Translator Usage (#843)

Author: patrick-ogrady

storage/src/index/mod.rs

@@ -15,16 +15,20 @@ mod storage;
 pub use storage::{Index, RemoveValueIterator, UpdateValueIterator, ValueIterator};
 pub mod translator;
 
-use std::hash::Hash;
+use std::hash::{BuildHasher, Hash};
 
 /// Translate keys into an internal representation used by `Index`.
 ///
 /// # Warning
 ///
-/// If invoking `transform` on keys results in many conflicts, the performance of `Index` will
-/// degrade substantially.
-pub trait Translator: Clone {
-    type Key: Eq + Hash + Send + Sync + Clone;
+/// The output of `transform` is used as the key in a hash table. If the output is not uniformly
+/// distributed, the performance of [Index] will degrade substantially.
+pub trait Translator: Clone + BuildHasher {
+    /// The type of the internal representation of keys.
+    ///
+    /// Although `Translator` is a [BuildHasher], the `Key` type must still implement [Hash] for compatibility
+    /// with the [std::collections::HashMap] used internally by [Index].
+    type Key: Eq + Hash;
 
     /// Transform a key into its internal representation.
     fn transform(&self, key: &[u8]) -> Self::Key;

storage/src/index/storage.rs

@@ -9,6 +9,11 @@ use std::{
     mem::swap,
 };
 
+/// The initial capacity of the hashmap. This is a guess at the number of unique keys we will
+/// encounter. The hashmap will grow as needed, but this is a good starting point (covering
+/// the entire [super::translator::OneCap] range).
+const INITIAL_CAPACITY: usize = 256;
+
 /// Each key is mapped to a `Record` that contains a linked list of potential values for the key.
 ///
 /// In the common case of a single value associated with a key, the value is stored within the
@@ -205,7 +210,7 @@ impl<V> Record<V> {
 /// An index that maps translated keys to values.
 pub struct Index<T: Translator, V> {
     translator: T,
-    map: HashMap<T::Key, Record<V>>,
+    map: HashMap<T::Key, Record<V>, T>,
 
     collisions: Counter,
     keys_pruned: Counter,
@@ -215,8 +220,8 @@ impl<T: Translator, V> Index<T, V> {
     /// Create a new index.
     pub fn init(context: impl Metrics, translator: T) -> Self {
         let s = Self {
-            translator,
-            map: HashMap::new(),
+            translator: translator.clone(),
+            map: HashMap::with_capacity_and_hasher(INITIAL_CAPACITY, translator),
             collisions: Counter::default(),
             keys_pruned: Counter::default(),
         };

storage/src/index/translator.rs

@@ -1,4 +1,57 @@
+//! Primitive implementations of [Translator].
+
 use super::Translator;
+use std::hash::{BuildHasher, Hasher};
+
+/// A “do-nothing” hasher for `uint`.
+///
+/// [super::Index] typically stores keys that are **already hashed** (shortened by the [Translator]).
+/// Re-hashing them with SipHash (by [std::collections::HashMap]) would waste CPU, so we give `HashMap`
+/// this identity hasher instead:
+///
+/// * `write_u8`, `write_u16`, `write_u32`, `write_u64` copies the input into an internal field;
+/// * `finish()` returns that value unchanged.
+///
+/// # Warning
+///
+/// This hasher is not suitable for general use. If the hasher is called over some type that is not
+/// `u8`, `u16`, `u32` or `u64`, it will panic.
+#[derive(Default, Clone)]
+pub struct UintIdentity {
+    value: u64,
+}
+
+impl Hasher for UintIdentity {
+    #[inline]
+    fn write(&mut self, _: &[u8]) {
+        unimplemented!("we should only ever call type-specific write methods");
+    }
+
+    #[inline]
+    fn write_u8(&mut self, i: u8) {
+        self.value = i as u64;
+    }
+
+    #[inline]
+    fn write_u16(&mut self, i: u16) {
+        self.value = i as u64;
+    }
+
+    #[inline]
+    fn write_u32(&mut self, i: u32) {
+        self.value = i as u64;
+    }
+
+    #[inline]
+    fn write_u64(&mut self, i: u64) {
+        self.value = i;
+    }
+
+    #[inline]
+    fn finish(&self) -> u64 {
+        self.value
+    }
+}
 
 /// Cap the key to a fixed length.
 ///
@@ -13,60 +66,96 @@ fn cap<const N: usize>(key: &[u8]) -> [u8; N] {
     capped[..len].copy_from_slice(&key[..len]);
     capped
 }
+
 macro_rules! define_cap_translator {
-    ($name:ident, $size:expr) => {
-        #[doc = concat!("A translator that caps the key to ", stringify!($size), " bytes.")]
-        #[derive(Clone)]
+    ($name:ident, $size:expr, $int:ty) => {
+        #[doc = concat!("Translator that caps the key to ", stringify!($size), " byte(s) and returns it packed in a ", stringify!($int), ".")]
+        #[derive(Clone, Default)]
         pub struct $name;
 
         impl Translator for $name {
-            type Key = [u8; $size];
+            // Minimal uint size for the key.
+            type Key = $int;
 
+            #[inline]
             fn transform(&self, key: &[u8]) -> Self::Key {
-                cap(key)
+                let capped = cap::<$size>(key);
+                <$int>::from_le_bytes(capped)
+            }
+        }
+
+        // Implement the `BuildHasher` trait for `IdentityHasher`.
+        impl BuildHasher for $name {
+            type Hasher = UintIdentity;
+
+            #[inline]
+            fn build_hasher(&self) -> Self::Hasher {
+                UintIdentity::default()
             }
         }
     };
 }
 
-define_cap_translator!(TwoCap, 2);
-define_cap_translator!(FourCap, 4);
-define_cap_translator!(EightCap, 8);
+// Define translators for different sizes.
+define_cap_translator!(OneCap, 1, u8);
+define_cap_translator!(TwoCap, 2, u16);
+define_cap_translator!(FourCap, 4, u32);
+define_cap_translator!(EightCap, 8, u64);
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::hash::Hasher;
+
+    #[test]
+    fn test_one_cap() {
+        let t = OneCap;
+        assert_eq!(t.transform(b"").to_le_bytes(), [0]);
+        assert_eq!(t.transform(b"a").to_le_bytes(), [b'a']);
+        assert_eq!(t.transform(b"ab").to_le_bytes(), [b'a']);
+        assert_eq!(t.transform(b"abc").to_le_bytes(), [b'a']);
+    }
 
     #[test]
     fn test_two_cap() {
-        let translator = TwoCap;
-        assert_eq!(translator.transform(b""), [0, 0]);
-        assert_eq!(translator.transform(b"a"), [b'a', 0]);
-        assert_eq!(translator.transform(b"ab"), [b'a', b'b']);
-        assert_eq!(translator.transform(b"abc"), [b'a', b'b']);
+        let t = TwoCap;
+        assert_eq!(t.transform(b"").to_le_bytes(), [0, 0]);
+        assert_eq!(t.transform(b"a").to_le_bytes(), [b'a', 0]);
+        assert_eq!(t.transform(b"ab").to_le_bytes(), [b'a', b'b']);
+        assert_eq!(t.transform(b"abc").to_le_bytes(), [b'a', b'b']);
     }
 
     #[test]
     fn test_four_cap() {
-        let translator = FourCap;
-        assert_eq!(translator.transform(b""), [0, 0, 0, 0]);
-        assert_eq!(translator.transform(b"a"), [b'a', 0, 0, 0]);
-        assert_eq!(translator.transform(b"abcd"), [b'a', b'b', b'c', b'd']);
-        assert_eq!(translator.transform(b"abcdef"), [b'a', b'b', b'c', b'd']);
+        let t = FourCap;
+        assert_eq!(t.transform(b"").to_le_bytes(), [0, 0, 0, 0]);
+        assert_eq!(t.transform(b"a").to_le_bytes(), [b'a', 0, 0, 0]);
+        assert_eq!(t.transform(b"abcd").to_le_bytes(), [b'a', b'b', b'c', b'd']);
+        assert_eq!(
+            t.transform(b"abcdef").to_le_bytes(),
+            [b'a', b'b', b'c', b'd']
+        );
     }
 
     #[test]
     fn test_eight_cap() {
-        let translator = EightCap;
-        assert_eq!(translator.transform(b""), [0, 0, 0, 0, 0, 0, 0, 0]);
-        assert_eq!(translator.transform(b"a"), [b'a', 0, 0, 0, 0, 0, 0, 0]);
+        let t = EightCap;
+        assert_eq!(t.transform(b"").to_le_bytes(), [0, 0, 0, 0, 0, 0, 0, 0]);
+        assert_eq!(t.transform(b"a").to_le_bytes(), [b'a', 0, 0, 0, 0, 0, 0, 0]);
         assert_eq!(
-            translator.transform(b"abcdefgh"),
+            t.transform(b"abcdefgh").to_le_bytes(),
             [b'a', b'b', b'c', b'd', b'e', b'f', b'g', b'h']
         );
         assert_eq!(
-            translator.transform(b"abcdefghijk"),
+            t.transform(b"abcdefghijk").to_le_bytes(),
             [b'a', b'b', b'c', b'd', b'e', b'f', b'g', b'h']
         );
     }
+
+    #[test]
+    #[should_panic(expected = "we should only ever call type-specific write methods")]
+    fn identity_hasher_panics_on_write_slice() {
+        let mut h = UintIdentity::default();
+        h.write(b"not an int");
+    }
 }