config support; started more sharding

src/sharding.rs

@@ -1,5 +1,8 @@
 use sha1::{Digest, Sha1};
 
+// https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/include/catalog/partition.h#L20
+const PARTITION_HASH_SEED: u64 = 0x7A5B22367996DCFD;
+
 pub struct Sharder {
     shards: usize,
 }
@@ -9,6 +12,8 @@ impl Sharder {
         Sharder { shards: shards }
     }
 
+    /// Use SHA1 to pick a shard for the key. The key can be anything,
+    /// including an int or a string.
     pub fn sha1(&self, key: &[u8]) -> usize {
         let mut hasher = Sha1::new();
         hasher.update(key);
@@ -17,6 +22,81 @@ impl Sharder {
         let i = u32::from_le_bytes(result[result.len() - 4..result.len()].try_into().unwrap());
         i as usize % self.shards
     }
+
+    /// Hash function used by Postgres to determine which partition
+    /// to put the row in when using HASH(column) partitioning.
+    /// Source: https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/common/hashfn.c#L631
+    pub fn pg_bigint_hash(&self, key: i64) -> usize {
+        let mut lohalf = key as u32;
+        let hihalf = (key >> 32) as u32;
+        println!("{}, {}", lohalf, hihalf);
+        lohalf ^= if key >= 0 { hihalf } else { !hihalf };
+        println!("Low half: {}", lohalf);
+        Self::pg_u32_hash(lohalf) as usize % self.shards
+    }
+
+    fn rot(x: u32, k: u32) -> u32 {
+        ((x) << (k)) | ((x) >> (32 - (k)))
+    }
+
+    #[inline]
+    fn mix(mut a: u32, mut b: u32, mut c: u32) -> (u32, u32, u32) {
+        a = a.wrapping_sub(c);
+        a ^= Self::rot(c, 4);
+        c = c.wrapping_add(b);
+        b = b.wrapping_add(a);
+        b ^= Self::rot(a, 6);
+        a = a.wrapping_add(c);
+        c = c.wrapping_add(b);
+        c ^= Self::rot(b, 8);
+        b = b.wrapping_add(a);
+        a = a.wrapping_add(c);
+        a ^= Self::rot(c, 16);
+        c = c.wrapping_add(b);
+        b = b.wrapping_add(a);
+        b ^= Self::rot(a, 19);
+        a = a.wrapping_add(c);
+        c = c.wrapping_add(b);
+        c ^= Self::rot(b, 4);
+        b = b.wrapping_add(a);
+        (a, b, c)
+    }
+
+    #[inline]
+    fn _final(mut a: u32, mut b: u32, mut c: u32) -> (u32, u32, u32) {
+        c ^= b;
+        c = c.wrapping_add(Self::rot(b, 14));
+        a ^= c;
+        a = a.wrapping_add(Self::rot(c, 11));
+        b ^= a;
+        b = b.wrapping_add(Self::rot(a, 25));
+        c ^= b;
+        c = c.wrapping_add(Self::rot(b, 16));
+        a ^= c;
+        a = a.wrapping_add(Self::rot(c, 4));
+        b ^= a;
+        b = b.wrapping_add(Self::rot(a, 14));
+        c ^= b;
+        c = c.wrapping_add(Self::rot(b, 24));
+        (a, b, c)
+    }
+
+    fn pg_u32_hash(val: u32) -> u64 {
+        let mut a: u32 = 0x9e3779b9 + 4 + 3923095;
+        let mut b = a;
+        let c = a;
+        let seed = PARTITION_HASH_SEED;
+
+        a = a.wrapping_add((seed >> 32) as u32);
+        b = b.wrapping_add(seed as u32);
+        let (mut a, b, c) = Self::mix(a, b, c);
+
+        a = a.wrapping_add(val);
+
+        let (a, b, c) = Self::_final(a, b, c);
+
+        (b as u64) << 32 | c as u64
+    }
 }
 
 #[cfg(test)]
@@ -30,4 +110,12 @@ mod test {
         let shard = sharder.sha1(key);
         assert_eq!(shard, 1);
     }
+
+    #[test]
+    fn test_pg_bigint_hash() {
+        let sharder = Sharder::new(4);
+        let key = 23423423 as i64;
+        let shard = sharder.pg_bigint_hash(key);
+        assert_eq!(shard, 0);
+    }
 }