From 2b601832b5c52b6c25fab1c7c83947547d42ea0d Mon Sep 17 00:00:00 2001
From: Maksim Zdobnikau <maksim.zdobnikau@swmansion.com>
Date: Tue, 2 Jan 2024 14:58:24 +0100
Subject: [PATCH] Refactor Poseidon

---
 Sources/CryptoToolkit/CryptoPoseidon.swift | 61 +++-------------------
 Sources/Starknet/Crypto/Poseidon.swift     | 50 ++++++++++++++++--
 2 files changed, 53 insertions(+), 58 deletions(-)

diff --git a/Sources/CryptoToolkit/CryptoPoseidon.swift b/Sources/CryptoToolkit/CryptoPoseidon.swift
index a810479de..f1b8796fa 100644
--- a/Sources/CryptoToolkit/CryptoPoseidon.swift
+++ b/Sources/CryptoToolkit/CryptoPoseidon.swift
@@ -1,59 +1,12 @@
-import BigInt
 import CFrameworkWrapper
 import Foundation
 
-public class CryptoPoseidon {
-    public class func poseidonHash(x: BigUInt, y: BigUInt) -> BigUInt {
-        unsplitBigUInt(
-            hades(
-                [
-                    splitBigUInt(x),
-                    splitBigUInt(y),
-                    (2, 0, 0, 0),
-                ]
-            )[0]
-        )
-    }
-
-    private static func hades(_ values: [(UInt64, UInt64, UInt64, UInt64)]) -> [(UInt64, UInt64, UInt64, UInt64)] {
-        permutation3(values)
-    }
-
-    private static func permutation3(_ state: [(UInt64, UInt64, UInt64, UInt64)]) -> [(UInt64, UInt64, UInt64, UInt64)] {
-        var mutableState = state
-        permutation_3(&mutableState)
-        return mutableState
-    }
-
-    private static func splitBigUInt(_ value: BigUInt) -> (UInt64, UInt64, UInt64, UInt64) {
-        var result: [UInt64] = [0, 0, 0, 0]
-
-        // if the input is zero, return the array of zeros
-        if value != BigUInt(0) {
-            // mask has all bits set to 1 except the least significant one
-            let mask = BigUInt(2).power(64) - 1
-
-            // loop through the 64-bit chunks of the BigUInt, shift them and store in the LongArray
-            for i in 0 ..< 4 {
-                result[i] = UInt64(value >> (i * 64) & mask)
-            }
-        }
-
-        return (result[0], result[1], result[2], result[3])
-    }
-
-    /// <#Description#>
-    /// - Parameter values: <#values description#>
-    /// - Returns: <#description#>
-    private static func unsplitBigUInt(_ values: (UInt64, UInt64, UInt64, UInt64)) -> BigUInt {
-        var arr: [UInt64] = [values.0, values.1, values.2, values.3]
-        let powersOfTwo = [BigUInt(2).power(0), BigUInt(2).power(64), BigUInt(2).power(128), BigUInt(2).power(192)]
-
-        // w * 2**0 + x * 2**64 + y * 2**128 + z * 2**192
-        var result = BigUInt(0)
-        for (b, p) in zip(arr, powersOfTwo) {
-            result += BigUInt(b) * p
-        }
-        return result
+public enum CryptoPoseidon {
+    public static func hades(
+        _ values: [(UInt64, UInt64, UInt64, UInt64)]
+    ) -> [(UInt64, UInt64, UInt64, UInt64)] {
+        var state = values
+        permutation_3(&state)
+        return state
     }
 }
diff --git a/Sources/Starknet/Crypto/Poseidon.swift b/Sources/Starknet/Crypto/Poseidon.swift
index b6d9d92ea..0ab148334 100644
--- a/Sources/Starknet/Crypto/Poseidon.swift
+++ b/Sources/Starknet/Crypto/Poseidon.swift
@@ -2,10 +2,52 @@ import BigInt
 import CryptoToolkit
 import Foundation
 
-public enum Poseidon {
-    public static func poseidonHash(x: Felt, y: Felt) -> Felt {
-        let result = CryptoPoseidon.poseidonHash(x: x.value, y: y.value)
+public class Poseidon {
+    /// Compute poseidon hash on input values.
+    ///
+    /// - Returns: Poseidon hash of the two values as Felt.
+    public class func poseidonHash(x: Felt, y: Felt) -> Felt {
+        let state = [
+            splitBigUInt(x.value),
+            splitBigUInt(y.value),
+            (2, 0, 0, 0),
+        ]
+        return Felt(clamping: combineToBigUInt(
+            CryptoPoseidon.hades(state)[0]
+        ))
+    }
+
+    private static func splitBigUInt(_ value: BigUInt) -> (UInt64, UInt64, UInt64, UInt64) {
+        var result: [UInt64] = [0, 0, 0, 0]
+
+        // if the input is zero, return the array of zeros
+        if value != BigUInt(0) {
+            // mask has all bits set to 1 except the least significant one
+            let mask = BigUInt(2).power(64) - 1
+
+            // loop through the 64-bit chunks of the BigUInt, shift them and store in the result array
+            for i in 0 ..< 4 {
+                result[i] = UInt64(value >> (i * 64) & mask)
+            }
+        }
+
+        return (result[0], result[1], result[2], result[3])
+    }
+
+    private static func combineToBigUInt(_ values: (UInt64, UInt64, UInt64, UInt64)) -> BigUInt {
+        let arr: [UInt64] = [values.0, values.1, values.2, values.3]
+        let powersOfTwo = [
+            BigUInt(2).power(0),
+            BigUInt(2).power(64),
+            BigUInt(2).power(128),
+            BigUInt(2).power(192),
+        ]
 
-        return Felt(result)!
+        // w * 2**0 + x * 2**64 + y * 2**128 + z * 2**192
+        var result = BigUInt(0)
+        for (b, p) in zip(arr, powersOfTwo) {
+            result += BigUInt(b) * p
+        }
+        return result
     }
 }