SQUASH Use chunks_exact() to avoid if statements

cloudflare · Aug 22, 2024 · 8ab98b7 · 8ab98b7
1 parent 96914e6
commit 8ab98b7
Showing 1 changed file with 40 additions and 10 deletions.
diff --git a/crates/daphne/src/pine/flp.rs b/crates/daphne/src/pine/flp.rs
@@ -189,40 +189,70 @@ impl<F: FftFriendlyFieldElement> PineType<F> {
             .iter_mut()
             .zip(rand.chunks(rand_len_per_test))
         {
-            for (gradient_chunk, r) in gradient.chunks(4).zip(rand_per_test) {
+            for (gradient_chunk, r) in gradient.chunks_exact(4).zip(rand_per_test) {
                 match r & 3 /* 0b11 */ {
                     0 => *wr_test_result -= gradient_chunk[0],
                     3 => *wr_test_result += gradient_chunk[0],
                     _ => (),
                 };
 
-                if gradient_chunk.len() < 2 {
-                    continue;
-                }
                 match r & 12 /* 0b1100 */ {
                     0 => *wr_test_result -= gradient_chunk[1],
                     12 => *wr_test_result += gradient_chunk[1],
                     _ => (),
                 };
 
-                if gradient_chunk.len() < 3 {
-                    continue;
-                }
                 match r & 48 /* 0b110000 */ {
                     0 => *wr_test_result -= gradient_chunk[2],
                     48 => *wr_test_result += gradient_chunk[2],
                     _ => (),
                 };
 
-                if gradient_chunk.len() < 4 {
-                    continue;
-                }
                 match r & 192 /* 0b11000000 */ {
                     0 => *wr_test_result -= gradient_chunk[3],
                     192 => *wr_test_result += gradient_chunk[3],
                     _ => (),
                 };
             }
+
+            let gradient_chunk = gradient.chunks_exact(4).remainder();
+            let r = rand_per_test.last().unwrap();
+
+            if gradient_chunk.len() < 1 {
+                continue;
+            }
+            match r & 3 /* 0b11 */ {
+                    0 => *wr_test_result -= gradient_chunk[0],
+                    3 => *wr_test_result += gradient_chunk[0],
+                    _ => (),
+                };
+
+            if gradient_chunk.len() < 2 {
+                continue;
+            }
+            match r & 12 /* 0b1100 */ {
+                    0 => *wr_test_result -= gradient_chunk[1],
+                    12 => *wr_test_result += gradient_chunk[1],
+                    _ => (),
+                };
+
+            if gradient_chunk.len() < 3 {
+                continue;
+            }
+            match r & 48 /* 0b110000 */ {
+                    0 => *wr_test_result -= gradient_chunk[2],
+                    48 => *wr_test_result += gradient_chunk[2],
+                    _ => (),
+                };
+
+            if gradient_chunk.len() < 4 {
+                continue;
+            }
+            match r & 192 /* 0b11000000 */ {
+                    0 => *wr_test_result -= gradient_chunk[3],
+                    192 => *wr_test_result += gradient_chunk[3],
+                    _ => (),
+                };
         }
 
         wr_test_results