Improve codegen for mixing in length (#112)

stepantubanov · web-flow · commit fbd7485ba1c9 · 2022-02-27T14:55:13.000-08:00
diff --git a/src/aes_hash.rs b/src/aes_hash.rs
@@ -82,14 +82,6 @@ impl AHasher {
         }
     }
 
-    #[inline(always)]
-    fn add_in_length(&mut self, length: u64) {
-        //This will be scrambled by the next AES round.
-        let mut enc: [u64; 2] = self.enc.convert();
-        enc[0] = enc[0].wrapping_add(length);
-        self.enc = enc.convert();
-    }
-
     #[inline(always)]
     fn hash_in(&mut self, new_value: u128) {
         self.enc = aesenc(self.enc, new_value);
@@ -159,7 +151,8 @@ impl Hasher for AHasher {
     fn write(&mut self, input: &[u8]) {
         let mut data = input;
         let length = data.len();
-        self.add_in_length(length as u64);
+        add_in_length(&mut self.enc, length as u64);
+
         //A 'binary search' on sizes reduces the number of comparisons.
         if data.len() <= 8 {
             let value = read_small(data);
@@ -337,7 +330,8 @@ impl Hasher for AHasherStr {
             self.0.enc = aesdec(self.0.sum, self.0.enc);
             self.0.enc = aesenc(aesenc(self.0.enc, self.0.key), self.0.enc);
         } else {
-            self.0.add_in_length(bytes.len() as u64);
+            add_in_length(&mut self.0.enc, bytes.len() as u64);
+
             let value = read_small(bytes).convert();
             self.0.sum = shuffle_and_add(self.0.sum, value);
             self.0.enc = aesdec(self.0.sum, self.0.enc);
diff --git a/src/operations.rs b/src/operations.rs
@@ -146,6 +146,29 @@ pub(crate) fn aesdec(value: u128, xor: u128) -> u128 {
     }
 }
 
+#[inline(always)]
+pub(crate) fn add_in_length(enc: &mut u128, len: u64) {
+    #[cfg(all(target_arch = "x86_64", target_feature = "sse2", not(miri)))]
+    {
+        #[cfg(target_arch = "x86_64")]
+        use core::arch::x86_64::*;
+
+        unsafe {
+            let enc = enc as *mut u128;
+            let len = _mm_cvtsi64_si128(len as i64);
+            let data = _mm_loadu_si128(enc.cast());
+            let sum = _mm_add_epi64(data, len);
+            _mm_storeu_si128(enc.cast(), sum);
+        }
+    }
+    #[cfg(not(all(target_arch = "x86_64", target_feature = "sse2", not(miri))))]
+    {
+        let mut t: [u64; 2] = enc.convert();
+        t[0] = t[0].wrapping_add(len);
+        *enc = t.convert();
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -327,4 +350,12 @@ mod test {
             shuffled = shuffle(shuffled);
         }
     }
+
+    #[test]
+    fn test_add_length() {
+        let mut enc = (u64::MAX as u128) << 64 | 50;
+        add_in_length(&mut enc, u64::MAX);
+        assert_eq!(enc >> 64, u64::MAX as u128);
+        assert_eq!(enc as u64, 49);
+    }
 }
