From 184f73c6048228be7cbc8a1733b37be34371d3bc Mon Sep 17 00:00:00 2001
From: Jakub Nabaglo <j@nab.gl>
Date: Tue, 2 Nov 2021 12:38:30 -0700
Subject: [PATCH] Expose optimized Poseidon routines to the Poseidon gate (ARM)
 (#331)

* Expose optimized Poseidon routines to the Poseidon gate

* Daniel PR comments
---
 .../arch/aarch64/poseidon_goldilocks_neon.rs  | 60 +++++++++++++------
 src/hash/poseidon_goldilocks.rs               | 16 +++++
 2 files changed, 58 insertions(+), 18 deletions(-)

diff --git a/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs b/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs
index f3225f76..f122e0ef 100644
--- a/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs
+++ b/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs
@@ -1006,8 +1006,8 @@ unsafe fn partial_rounds(
 }
 
 #[inline(always)]
-pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
-    let state = [
+fn unwrap_state(state: [GoldilocksField; 12]) -> [u64; 12] {
+    [
         state[0].0,
         state[1].0,
         state[2].0,
@@ -1020,22 +1020,11 @@ pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
         state[9].0,
         state[10].0,
         state[11].0,
-    ];
-    let state = const_layer_full(state, ALL_ROUND_CONSTANTS[0..WIDTH].try_into().unwrap());
-    let state = full_rounds(
-        state,
-        ALL_ROUND_CONSTANTS[WIDTH..WIDTH * (HALF_N_FULL_ROUNDS + 1)]
-            .try_into()
-            .unwrap(),
-    );
-    let state = partial_rounds(
-        state,
-        ALL_ROUND_CONSTANTS
-            [WIDTH * (HALF_N_FULL_ROUNDS + 1)..WIDTH * (HALF_N_FULL_ROUNDS + N_PARTIAL_ROUNDS + 1)]
-            .try_into()
-            .unwrap(),
-    );
-    let state = full_rounds(state, &FINAL_ROUND_CONSTANTS);
+    ]
+}
+
+#[inline(always)]
+fn wrap_state(state: [u64; 12]) -> [GoldilocksField; 12] {
     [
         GoldilocksField(state[0]),
         GoldilocksField(state[1]),
@@ -1051,3 +1040,38 @@ pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
         GoldilocksField(state[11]),
     ]
 }
+
+#[inline(always)]
+pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
+    let state = unwrap_state(state);
+    let state = const_layer_full(state, ALL_ROUND_CONSTANTS[0..WIDTH].try_into().unwrap());
+    let state = full_rounds(
+        state,
+        ALL_ROUND_CONSTANTS[WIDTH..WIDTH * (HALF_N_FULL_ROUNDS + 1)]
+            .try_into()
+            .unwrap(),
+    );
+    let state = partial_rounds(
+        state,
+        ALL_ROUND_CONSTANTS
+            [WIDTH * (HALF_N_FULL_ROUNDS + 1)..WIDTH * (HALF_N_FULL_ROUNDS + N_PARTIAL_ROUNDS + 1)]
+            .try_into()
+            .unwrap(),
+    );
+    let state = full_rounds(state, &FINAL_ROUND_CONSTANTS);
+    wrap_state(state)
+}
+
+#[inline(always)]
+pub unsafe fn sbox_layer(state: &mut [GoldilocksField; WIDTH]) {
+    *state = wrap_state(sbox_layer_full(unwrap_state(*state)));
+}
+
+#[inline(always)]
+pub unsafe fn mds_layer(state: &[GoldilocksField; WIDTH]) -> [GoldilocksField; WIDTH] {
+    let state = unwrap_state(*state);
+    // We want to do an MDS layer without the constant layer.
+    let round_consts = [0u64; WIDTH];
+    let state = mds_const_layers_full(state, &round_consts);
+    wrap_state(state)
+}
diff --git a/src/hash/poseidon_goldilocks.rs b/src/hash/poseidon_goldilocks.rs
index 32d5e237..0b1f9a49 100644
--- a/src/hash/poseidon_goldilocks.rs
+++ b/src/hash/poseidon_goldilocks.rs
@@ -389,6 +389,22 @@ impl Poseidon<12> for GoldilocksField {
             crate::hash::arch::aarch64::poseidon_goldilocks_neon::poseidon(input)
         }
     }
+
+    #[cfg(all(target_arch="aarch64", target_feature="neon"))]
+    #[inline(always)]
+    fn sbox_layer(state: &mut [Self; 12]) {
+        unsafe {
+            crate::hash::arch::aarch64::poseidon_goldilocks_neon::sbox_layer(state);
+        }
+    }
+
+    #[cfg(all(target_arch="aarch64", target_feature="neon"))]
+    #[inline(always)]
+    fn mds_layer(state: &[Self; 12]) -> [Self; 12] {
+        unsafe {
+            crate::hash::arch::aarch64::poseidon_goldilocks_neon::mds_layer(state)
+        }
+    }
 }
 
 #[cfg(test)]