layout data in simd friendly way

Signed-off-by: eternal-flame-AD <yume@yumechi.jp>

Author: eternal-flame-AD

salsa20/src/backends/soft.rs

@@ -24,8 +24,9 @@ impl<R: Unsigned> StreamCipherBackend for Backend<'_, R> {
 
         self.0.set_block_pos(self.0.get_block_pos() + 1);
 
-        for (chunk, val) in block.chunks_exact_mut(4).zip(res.iter()) {
-            chunk.copy_from_slice(&val.to_le_bytes());
+        for i in 0..16 {
+            block[i * 4..(i + 1) * 4]
+                .copy_from_slice(&res[crate::DATA_LAYOUT_INVERSE[i]].to_le_bytes());
         }
     }
 }
@@ -39,6 +40,10 @@ pub(crate) fn quarter_round(
     d: usize,
     state: &mut [u32; STATE_WORDS],
 ) {
+    let a = crate::DATA_LAYOUT_INVERSE[a];
+    let b = crate::DATA_LAYOUT_INVERSE[b];
+    let c = crate::DATA_LAYOUT_INVERSE[c];
+    let d = crate::DATA_LAYOUT_INVERSE[d];
     state[b] ^= state[a].wrapping_add(state[d]).rotate_left(7);
     state[c] ^= state[b].wrapping_add(state[a]).rotate_left(9);
     state[d] ^= state[c].wrapping_add(state[b]).rotate_left(13);

salsa20/src/lib.rs

@@ -118,6 +118,25 @@ const STATE_WORDS: usize = 16;
 /// State initialization constant ("expand 32-byte k")
 const CONSTANTS: [u32; 4] = [0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574];
 
+const DATA_LAYOUT: [usize; 16] = [0, 5, 10, 15, 4, 9, 14, 3, 12, 1, 6, 11, 8, 13, 2, 7];
+
+const DATA_LAYOUT_INVERSE: [usize; 16] = {
+    let mut index = [0; 16];
+    let mut i = 0;
+    while i < 16 {
+        let mut inverse = 0;
+        while inverse < 16 {
+            if DATA_LAYOUT[inverse] == i {
+                index[i] = inverse;
+                break;
+            }
+            inverse += 1;
+        }
+        i += 1;
+    }
+    index
+};
+
 /// The Salsa20 core function.
 pub struct SalsaCore<R: Unsigned> {
     /// Internal state of the core function
@@ -133,7 +152,7 @@ impl<R: Unsigned> SalsaCore<R> {
     /// Other users generally should not use this method.
     pub fn from_raw_state(state: [u32; STATE_WORDS]) -> Self {
         Self {
-            state,
+            state: core::array::from_fn(|i| state[DATA_LAYOUT[i]]),
             rounds: PhantomData,
         }
     }
@@ -177,7 +196,7 @@ impl<R: Unsigned> KeyIvInit for SalsaCore<R> {
         state[15] = CONSTANTS[3];
 
         Self {
-            state,
+            state: core::array::from_fn(|i| state[DATA_LAYOUT[i]]),
             rounds: PhantomData,
         }
     }
@@ -199,13 +218,14 @@ impl<R: Unsigned> StreamCipherSeekCore for SalsaCore<R> {
 
     #[inline(always)]
     fn get_block_pos(&self) -> u64 {
-        (self.state[8] as u64) + ((self.state[9] as u64) << 32)
+        (self.state[DATA_LAYOUT_INVERSE[8]] as u64)
+            + ((self.state[DATA_LAYOUT_INVERSE[9]] as u64) << 32)
     }
 
     #[inline(always)]
     fn set_block_pos(&mut self, pos: u64) {
-        self.state[8] = (pos & 0xffff_ffff) as u32;
-        self.state[9] = ((pos >> 32) & 0xffff_ffff) as u32;
+        self.state[DATA_LAYOUT_INVERSE[8]] = (pos & 0xffff_ffff) as u32;
+        self.state[DATA_LAYOUT_INVERSE[9]] = ((pos >> 32) & 0xffff_ffff) as u32;
     }
 }
 

salsa20/src/xsalsa.rs

@@ -90,28 +90,32 @@ impl<R: Unsigned> ZeroizeOnDrop for XSalsaCore<R> {}
 ///
 /// It produces 256-bits of output suitable for use as a Salsa20 key
 pub fn hsalsa<R: Unsigned>(key: &Key, input: &Array<u8, U16>) -> Array<u8, U32> {
+    const KEY_IDX: [usize; 8] = [0, 5, 10, 15, 6, 7, 8, 9];
+
     #[inline(always)]
     fn to_u32(chunk: &[u8]) -> u32 {
         u32::from_le_bytes(chunk.try_into().unwrap())
     }
 
-    let mut state = [0u32; 16];
-    state[0] = CONSTANTS[0];
-    state[1..5]
+    let mut t = [0u32; 16];
+    t[0] = CONSTANTS[0];
+    t[1..5]
         .iter_mut()
         .zip(key[0..16].chunks_exact(4))
         .for_each(|(v, chunk)| *v = to_u32(chunk));
-    state[5] = CONSTANTS[1];
-    state[6..10]
+    t[5] = CONSTANTS[1];
+    t[6..10]
         .iter_mut()
         .zip(input.chunks_exact(4))
         .for_each(|(v, chunk)| *v = to_u32(chunk));
-    state[10] = CONSTANTS[2];
-    state[11..15]
+    t[10] = CONSTANTS[2];
+    t[11..15]
         .iter_mut()
         .zip(key[16..].chunks_exact(4))
         .for_each(|(v, chunk)| *v = to_u32(chunk));
-    state[15] = CONSTANTS[3];
+    t[15] = CONSTANTS[3];
+
+    let mut state = core::array::from_fn(|i| t[crate::DATA_LAYOUT[i]]);
 
     // 20 rounds consisting of 10 column rounds and 10 diagonal rounds
     for _ in 0..R::USIZE {
@@ -129,10 +133,10 @@ pub fn hsalsa<R: Unsigned>(key: &Key, input: &Array<u8, U16>) -> Array<u8, U32>
     }
 
     let mut output = Array::default();
-    let key_idx: [usize; 8] = [0, 5, 10, 15, 6, 7, 8, 9];
 
-    for (i, chunk) in output.chunks_exact_mut(4).enumerate() {
-        chunk.copy_from_slice(&state[key_idx[i]].to_le_bytes());
+    for i in 0..8 {
+        output[i * 4..(i + 1) * 4]
+            .copy_from_slice(&state[crate::DATA_LAYOUT_INVERSE[KEY_IDX[i]]].to_le_bytes());
     }
 
     output