perf: Improve TokenPair hashing for faster generation

crates/nailkov/src/distribution.rs

···
 //! [`TokenWeights`] are representations of how common [`Token`]s are, and are paired up with
 //! a [`TokenPair`](crate::token::TokenPair) in a [`NailKov`](crate::NailKov).
 
-use core::hash::BuildHasherDefault;
-
-use estr::IdentityHasher;
 use indexmap::IndexMap;
 use rand::Rng;
 use rand_distr::{Distribution, weighted::WeightedAliasIndex};
 
-use crate::{error::NailError, token::Token};
+use crate::{TokenHasher, error::NailError, token::Token};
 
 /// A distribution of choices and their likelihood.
 #[derive(Clone, Debug)]
···
 #[derive(Clone, Debug)]
 pub struct TokenWeightsBuilder {
     /// Counts how many times a token is likely to appear.
-    occurrences: IndexMap<Token, u32, BuildHasherDefault<IdentityHasher>>,
+    occurrences: IndexMap<Token, u32, TokenHasher>,
 }
 
 impl TokenWeightsBuilder {

crates/nailkov/src/lib.rs

···
 mod error;
 mod token;
 
+use core::hash::BuildHasherDefault;
+
 use crossbeam_utils::CachePadded;
 use error::NailError;
+use estr::IdentityHasher;
 use indexmap::IndexMap;
 use itertools::Itertools;
 use rand::{RngCore, seq::IteratorRandom};
 use rand_distr::Distribution;
 
 use distribution::{TokenWeights, TokenWeightsBuilder};
-use rapidhash::fast::RandomState;
 use token::{Token, TokenPair};
 use unicode_segmentation::UnicodeSegmentation;
 
+/// `nailkov` relies on `estr`'s precomputed hashes, so we avoid
+/// hashing ourselves and can just use the precomputed hashes instead.
+type TokenHasher = BuildHasherDefault<IdentityHasher>;
+
 #[derive(Clone, Debug)]
 pub struct NailKov {
-    chain: CachePadded<IndexMap<TokenPair, TokenWeights, RandomState>>,
+    chain: CachePadded<IndexMap<TokenPair, TokenWeights, TokenHasher>>,
 }
 
 pub struct NailKovIter<'a, R: RngCore> {
···
 
 impl NailKov {
     pub fn from_input(input: &str) -> Result<NailKov, NailError> {
-        NailBuilder::new(RandomState::new()).with_input(input)
+        NailBuilder::new(TokenHasher::new()).with_input(input)
     }
 }
 
 struct NailBuilder {
-    chain: IndexMap<TokenPair, TokenWeightsBuilder, RandomState>,
+    chain: IndexMap<TokenPair, TokenWeightsBuilder, TokenHasher>,
 }
 
 impl NailBuilder {
-    fn new(hasher: RandomState) -> Self {
+    fn new(hasher: TokenHasher) -> Self {
         Self {
             chain: IndexMap::with_hasher(hasher),
         }
···
             return Err(NailError::EmptyInput);
         }
 
-        let chain: IndexMap<TokenPair, TokenWeights, RandomState> = self
+        let chain: IndexMap<TokenPair, TokenWeights, TokenHasher> = self
             .chain
             .into_iter()
             .flat_map(|(pair, dist)| {
···
             return Err(NailError::EmptyInput);
         }
 
-        Ok(NailKov { chain: CachePadded::new(chain) })
+        Ok(NailKov {
+            chain: CachePadded::new(chain),
+        })
     }
 
     /// Add the occurrence of `next` following `prev`.

crates/nailkov/src/token.rs

···
 }
 
 /// An owned pair of [`Token`]s.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 // Alignment repr necessary to allow LLVM to better output
-// optimized codegen for `to_bits`, `PartialEq`
 // Prior art taken from my contribution to Bevy:
 // https://github.com/bevyengine/bevy/blob/main/crates/bevy_ecs/src/entity/mod.rs#L309
 #[repr(C, align(16))]
 pub struct TokenPair {
     // Do not reorder the fields here. The ordering is explicitly used by repr(C)
-    // to make this struct equivalent to a u64.
+    // to make this struct equivalent to a u128.
     #[cfg(target_endian = "little")]
     pub left: Token,
     pub right: Token,
···
     pub left: Token,
 }
 
-// By not short-circuiting in comparisons, we get better codegen.
-// See <https://github.com/rust-lang/rust/issues/117800>
-impl PartialEq for TokenPair {
-    #[inline(always)]
-    fn eq(&self, other: &TokenPair) -> bool {
-        // By using `to_bits`, the codegen can be optimized out even
-        // further potentially. Relies on the correct alignment/field
-        // order of `TokenPair`.
-        self.to_bits() == other.to_bits()
-    }
-}
-
-impl Eq for TokenPair {}
-
 impl core::hash::Hash for TokenPair {
-    #[inline(always)]
+    #[inline]
     fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
-        self.to_bits().hash(state);
+        // Use only with an IdentityHasher so that you don't rehash the hash
+        self.double_hash().hash(state);
     }
 }
 
···
         Self { left, right }
     }
 
+    /// Use the precomputed hashes to generate a secondary hash.
+    /// Method from [fastbloom](https://github.com/tomtomwombat/fastbloom/blob/main/src/hasher.rs#L190),
+    /// which was in turn adapted from <https://www.eecs.harvard.edu/~michaelm/postscripts/rsa2008.pdf>.
     #[inline(always)]
-    fn to_bits(self) -> u128 {
-        (self.left.0.digest().hash() as u128) | ((self.right.0.digest().hash() as u128) << 64)
+    fn double_hash(&self) -> u64 {
+        self.left
+            .0
+            .digest()
+            .hash()
+            .wrapping_add(self.right.0.digest().hash())
+            .rotate_left(5)
     }
 }