Add null hasher for Pubkey (anza-xyz#96)

* null hasher for Pubkey that allows 4x faster hashmaps
* truncate the counter in new_unique to 4 bytes and fill the rest with pseudorandom bytes to improve statistical properties of the resulting pubkeys
---------

Co-authored-by: Alex Pyattaev <[email protected]>

Author: alexpyattaev

pubkey/src/lib.rs

@@ -23,7 +23,9 @@ use {
     core::{
         array,
         convert::{Infallible, TryFrom},
-        fmt, mem,
+        fmt,
+        hash::{Hash, Hasher},
+        mem,
         str::{from_utf8, FromStr},
     },
     num_traits::{FromPrimitive, ToPrimitive},
@@ -158,10 +160,163 @@ impl From<u64> for PubkeyError {
 #[cfg_attr(all(feature = "borsh", feature = "std"), derive(BorshSchema))]
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
 #[cfg_attr(feature = "bytemuck", derive(Pod, Zeroable))]
-#[derive(Clone, Copy, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
+#[derive(Clone, Copy, Default, Eq, Ord, PartialEq, PartialOrd)]
 #[cfg_attr(feature = "dev-context-only-utils", derive(Arbitrary))]
 pub struct Pubkey(pub(crate) [u8; 32]);
 
+/// Custom impl of Hash for Pubkey
+/// allows us to skip hashing the length of the pubkey
+/// which is always the same anyway
+impl Hash for Pubkey {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        state.write(self.as_array());
+    }
+}
+
+#[cfg(all(feature = "rand", not(target_os = "solana")))]
+mod hasher {
+    use {
+        crate::PUBKEY_BYTES,
+        core::{
+            cell::Cell,
+            hash::{BuildHasher, Hasher},
+        },
+        rand::{thread_rng, Rng},
+    };
+
+    /// A faster, but less collision resistant hasher for pubkeys.
+    ///
+    /// Specialized hasher that uses a random 8 bytes subslice of the
+    /// pubkey as the hash value. Should not be used when collisions
+    /// might be used to mount DOS attacks.
+    ///
+    /// Using this results in about 4x faster lookups in a typical hashmap.
+    #[derive(Default)]
+    pub struct PubkeyHasher {
+        offset: usize,
+        state: u64,
+    }
+
+    impl Hasher for PubkeyHasher {
+        #[inline]
+        fn finish(&self) -> u64 {
+            self.state
+        }
+        #[inline]
+        fn write(&mut self, bytes: &[u8]) {
+            debug_assert_eq!(
+                bytes.len(),
+                PUBKEY_BYTES,
+                "This hasher is intended to be used with pubkeys and nothing else"
+            );
+            // This slice/unwrap can never panic since offset is < PUBKEY_BYTES - size_of::<u64>()
+            let chunk: &[u8; size_of::<u64>()] = bytes[self.offset..self.offset + size_of::<u64>()]
+                .try_into()
+                .unwrap();
+            self.state = u64::from_ne_bytes(*chunk);
+        }
+    }
+
+    /// A builder for faster, but less collision resistant hasher for pubkeys.
+    ///
+    /// Initializes `PubkeyHasher` instances that use an 8-byte
+    /// slice of the pubkey as the hash value. Should not be used when
+    /// collisions might be used to mount DOS attacks.
+    ///
+    /// Using this results in about 4x faster lookups in a typical hashmap.
+    #[derive(Clone)]
+    pub struct PubkeyHasherBuilder {
+        offset: usize,
+    }
+
+    impl Default for PubkeyHasherBuilder {
+        /// Default construct the PubkeyHasherBuilder.
+        ///
+        /// The position of the slice is determined initially
+        /// through random draw and then by incrementing a thread-local
+        /// This way each hashmap can be expected to use a slightly different
+        /// slice. This is essentially the same mechanism as what is used by
+        /// `RandomState`
+        fn default() -> Self {
+            std::thread_local!(static OFFSET: Cell<usize>  = {
+                let mut rng = thread_rng();
+                Cell::new(rng.gen_range(0..PUBKEY_BYTES - size_of::<u64>()))
+            });
+
+            let offset = OFFSET.with(|offset| {
+                let mut next_offset = offset.get() + 1;
+                if next_offset > PUBKEY_BYTES - size_of::<u64>() {
+                    next_offset = 0;
+                }
+                offset.set(next_offset);
+                next_offset
+            });
+            PubkeyHasherBuilder { offset }
+        }
+    }
+
+    impl BuildHasher for PubkeyHasherBuilder {
+        type Hasher = PubkeyHasher;
+        #[inline]
+        fn build_hasher(&self) -> Self::Hasher {
+            PubkeyHasher {
+                offset: self.offset,
+                state: 0,
+            }
+        }
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use {
+            super::PubkeyHasherBuilder,
+            crate::Pubkey,
+            core::hash::{BuildHasher, Hasher},
+        };
+        #[test]
+        fn test_pubkey_hasher_builder() {
+            let key = Pubkey::new_unique();
+            let builder = PubkeyHasherBuilder::default();
+            let mut hasher1 = builder.build_hasher();
+            let mut hasher2 = builder.build_hasher();
+            hasher1.write(key.as_array());
+            hasher2.write(key.as_array());
+            assert_eq!(
+                hasher1.finish(),
+                hasher2.finish(),
+                "Hashers made with same builder should be identical"
+            );
+            // Make sure that when we make new builders we get different slices
+            // chosen for hashing
+            let builder2 = PubkeyHasherBuilder::default();
+            for _ in 0..64 {
+                let mut hasher3 = builder2.build_hasher();
+                hasher3.write(key.as_array());
+                std::dbg!(hasher1.finish());
+                std::dbg!(hasher3.finish());
+                if hasher1.finish() != hasher3.finish() {
+                    return;
+                }
+            }
+            panic!("Hashers built with different builder should be different due to random offset");
+        }
+
+        #[test]
+        fn test_pubkey_hasher() {
+            let key1 = Pubkey::new_unique();
+            let key2 = Pubkey::new_unique();
+            let builder = PubkeyHasherBuilder::default();
+            let mut hasher1 = builder.build_hasher();
+            let mut hasher2 = builder.build_hasher();
+            hasher1.write(key1.as_array());
+            hasher2.write(key2.as_array());
+            assert_ne!(hasher1.finish(), hasher2.finish());
+        }
+    }
+}
+#[cfg(all(feature = "rand", not(target_os = "solana")))]
+pub use hasher::{PubkeyHasher, PubkeyHasherBuilder};
+
 impl solana_sanitize::Sanitize for Pubkey {}
 
 // Use strum when testing to ensure our FromPrimitive
@@ -322,12 +477,33 @@ impl Pubkey {
     pub fn new_unique() -> Self {
         use solana_atomic_u64::AtomicU64;
         static I: AtomicU64 = AtomicU64::new(1);
-
-        let mut b = [0u8; 32];
-        let i = I.fetch_add(1);
+        type T = u32;
+        const COUNTER_BYTES: usize = size_of::<T>();
+        let mut b = [0u8; PUBKEY_BYTES];
+        let mut i = I.fetch_add(1) as T;
         // use big endian representation to ensure that recent unique pubkeys
-        // are always greater than less recent unique pubkeys
-        b[0..8].copy_from_slice(&i.to_be_bytes());
+        // are always greater than less recent unique pubkeys.
+        b[0..COUNTER_BYTES].copy_from_slice(&i.to_be_bytes());
+        // fill the rest of the pubkey with pseudorandom numbers to make
+        // data statistically similar to real pubkeys.
+        #[cfg(any(feature = "std", target_arch = "wasm32"))]
+        {
+            extern crate std;
+            let mut hash = std::hash::DefaultHasher::new();
+            for slice in b[COUNTER_BYTES..].chunks_mut(COUNTER_BYTES) {
+                hash.write_u32(i);
+                i += 1;
+                slice.copy_from_slice(&hash.finish().to_ne_bytes()[0..COUNTER_BYTES]);
+            }
+        }
+        // if std is not available, just replicate last byte of the counter.
+        // this is not as good as a proper hash, but at least it is uniform
+        #[cfg(not(any(feature = "std", target_arch = "wasm32")))]
+        {
+            for b in b[COUNTER_BYTES..].iter_mut() {
+                *b = (i & 0xFF) as u8;
+            }
+        }
         Self::from(b)
     }