jump_back_hash/

lib.rs

//! Implementation of the JumpBackHash consistent hash algorithm
//! (Ertl, 2024, <https://arxiv.org/abs/2403.18682>).
//!
//! Distributed systems need efficient ways to map keys to buckets. The naive modulo approach
//! is fast but unstable: changing the bucket count remaps most keys, causing traffic spikes.
//! Consistent hash algorithms minimise remappings, but existing ones are either slow, require
//! floating-point arithmetic, or depend on hash-function families not found in standard
//! libraries. JumpBackHash avoids all of these: it borrows the concept of *active indices*
//! from consistent weighted sampling to guarantee consistency, generates those indices in
//! reverse order to eliminate floating-point operations, and minimises the number of random
//! values consumed — making it both simple and fast with an expected constant runtime.
//!
//! This crate implements only the bucketing algorithm itself. The caller is responsible for
//! choosing and seeding the [`rand_core::RngCore`] passed to [`bucket`].

use rand_core::RngCore;

/// Maps an RNG state to a bucket index in `[0, num_buckets)` using JumpBackHash
/// (Ertl, 2024, <https://arxiv.org/abs/2403.18682>).
///
/// The caller seeds `rng` with the input hash (or any value) before calling this function.
/// Different [`RngCore`] implementations will produce different bucket assignments.
pub fn bucket<R: RngCore>(rng: &mut R, num_buckets: u32) -> u32 {
    if num_buckets <= 1 {
        return 0;
    }

    let r0 = rng.next_u64();

    let x_masked_init =
        ((r0 ^ (r0 >> 32)) as u32) & (0xFFFF_FFFFu32 >> (num_buckets - 1).leading_zeros());

    let mut x_masked = x_masked_init;

    loop {
        if x_masked == 0 {
            return 0;
        }
        let bucket_range_min: u32 = 1u32 << (31 - x_masked.leading_zeros());
        let bit_count = x_masked.count_ones();
        let bucket_idx =
            bucket_range_min | (r0.wrapping_shr(bit_count * 32) as u32 & (bucket_range_min - 1));

        if bucket_idx < num_buckets {
            return bucket_idx;
        }

        let bucket_range_max = (bucket_range_min << 1) - 1;

        loop {
            let r1 = rng.next_u64();

            let idx1 = (r1 as u32) & bucket_range_max;
            if idx1 < bucket_range_min {
                break;
            }
            if idx1 < num_buckets {
                return idx1;
            }

            let idx2 = ((r1 >> 32) as u32) & bucket_range_max;
            if idx2 < bucket_range_min {
                break;
            }
            if idx2 < num_buckets {
                return idx2;
            }
        }

        x_masked ^= bucket_range_min;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand_core::RngCore;

    /// Trivial additive RNG: state += 19_561_023 each step.
    struct AddRng(u64);

    impl RngCore for AddRng {
        fn next_u32(&mut self) -> u32 {
            self.next_u64() as u32
        }
        fn next_u64(&mut self) -> u64 {
            self.0 = self.0.wrapping_add(19_561_023);
            self.0
        }
        fn fill_bytes(&mut self, dest: &mut [u8]) {
            rand_core::impls::fill_bytes_via_next(self, dest);
        }
        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand_core::Error> {
            self.fill_bytes(dest);
            Ok(())
        }
    }

    fn add_rng(seed: u64) -> AddRng {
        AddRng(seed)
    }

    /// SplitMix64 RNG used only in the consistency test, where output quality matters.
    /// Algorithm: <https://prng.di.unimi.it/splitmix64.c>
    struct SplitMix64(u64);

    impl SplitMix64 {
        fn new(seed: u64) -> Self {
            Self(seed)
        }
    }

    impl RngCore for SplitMix64 {
        fn next_u32(&mut self) -> u32 {
            self.next_u64() as u32
        }
        fn next_u64(&mut self) -> u64 {
            self.0 = self.0.wrapping_add(0x9e3779b97f4a7c15);
            let mut z = self.0;
            z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
            z ^ (z >> 31)
        }
        fn fill_bytes(&mut self, dest: &mut [u8]) {
            rand_core::impls::fill_bytes_via_next(self, dest);
        }
        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand_core::Error> {
            self.fill_bytes(dest);
            Ok(())
        }
    }

    #[test]
    fn single_bucket_always_zero() {
        for seed in [0u64, 1, u64::MAX, 0xdeadbeef] {
            assert_eq!(bucket(&mut add_rng(seed), 1), 0);
        }
    }

    #[test]
    fn zero_buckets_returns_zero() {
        assert_eq!(bucket(&mut add_rng(42), 0), 0);
    }

    #[test]
    fn result_in_range() {
        for n in 2u32..=64 {
            for seed in [0u64, 1, 12345, 0xdeadbeefcafe, u64::MAX] {
                let b = bucket(&mut add_rng(seed), n);
                assert!(b < n, "bucket(seed={seed}, n={n}) = {b} out of range");
            }
        }
    }

    #[test]
    fn deterministic() {
        for n in [2u32, 3, 10, 100] {
            for seed in [0u64, 1, 999, u64::MAX / 3] {
                assert_eq!(bucket(&mut add_rng(seed), n), bucket(&mut add_rng(seed), n));
            }
        }
    }

    /// Statistical consistency test (Ertl, 2024, <https://arxiv.org/abs/2403.18682>).
    ///
    /// A consistent hash must satisfy: when the bucket count grows from `n` to `n+1`,
    /// exactly the keys previously assigned to bucket `n` are redistributed — so only
    /// a fraction of `1/(n+1)` of all keys change their assignment. This test verifies
    /// that property empirically over 10 000 random keys for several values of `n`.
    ///
    /// The tolerance is ±4 standard deviations of the binomial distribution
    /// (p = 1/(n+1), trials = KEY_COUNT), giving a false-failure probability well
    /// below 1 in 10 000 per (n, key) combination.
    #[test]
    fn consistent_remapping_fraction() {
        use rand_core::OsRng;

        const KEY_COUNT: u64 = 10_000;

        for n in [2u32, 3, 5, 10, 20, 50, 100] {
            let mut moved: u64 = 0;

            for _ in 0..KEY_COUNT {
                let key = OsRng.next_u64();
                let b_n = bucket(&mut SplitMix64::new(key), n);
                let b_n1 = bucket(&mut SplitMix64::new(key), n + 1);
                if b_n != b_n1 {
                    moved += 1;
                }
            }

            // Expected fraction moved: 1/(n+1)
            let p = 1.0 / (n + 1) as f64;
            let expected = KEY_COUNT as f64 * p;
            let std_dev = (KEY_COUNT as f64 * p * (1.0 - p)).sqrt();
            let tolerance = 4.0 * std_dev;

            assert!(
                (moved as f64 - expected).abs() <= tolerance,
                "n={n}: moved {moved}/{KEY_COUNT} keys, expected {expected:.1} ± {tolerance:.1}"
            );
        }
    }
}