From c14af2cecf4fde69602ac9913ee424ea6da6430b Mon Sep 17 00:00:00 2001
From: Qiwei Yang <yangqiwei97@gmail.com>
Date: Wed, 14 Aug 2024 18:16:01 +0800
Subject: [PATCH] state merkl and trie test vectors (#60)

* wip

* fix
---
 JAMTests/Tests/JAMTests/TrieTests.swift | 37 +++++++++++++
 Utils/Sources/Utils/Data+Utils.swift    |  7 ++-
 Utils/Sources/Utils/merklization.swift  | 71 +++++++++++++++++++++++++
 3 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 JAMTests/Tests/JAMTests/TrieTests.swift
 create mode 100644 Utils/Sources/Utils/merklization.swift

diff --git a/JAMTests/Tests/JAMTests/TrieTests.swift b/JAMTests/Tests/JAMTests/TrieTests.swift
new file mode 100644
index 00000000..c0097526
--- /dev/null
+++ b/JAMTests/Tests/JAMTests/TrieTests.swift
@@ -0,0 +1,37 @@
+import Foundation
+import Testing
+import Utils
+
+@testable import JAMTests
+
+struct TrieElement: Codable {
+    let input: [String: String]
+    let output: String
+}
+
+typealias TrieTestCase = [TrieElement]
+
+struct TrieTests {
+    static func loadTests() throws -> [TrieTestCase] {
+        let tests = try TestLoader.getTestFiles(path: "trie", extension: "json")
+        return try tests.map {
+            let data = try Data(contentsOf: URL(fileURLWithPath: $0.path))
+            let decoder = JSONDecoder()
+            return try decoder.decode(TrieTestCase.self, from: data)
+        }
+    }
+
+    @Test(arguments: try loadTests())
+    func trieTests(_ testcase: TrieTestCase) throws {
+        for element in testcase {
+            let kv = element.input.reduce(into: [Data32: Data]()) { result, entry in
+                let keyData = Data(fromHexString: entry.key)
+                let valueData = Data(fromHexString: entry.value)
+                result[Data32(keyData!)!] = valueData
+            }
+
+            let result = try stateMerklize(kv: kv)
+            #expect(result.data.toHexString() == element.output)
+        }
+    }
+}
diff --git a/Utils/Sources/Utils/Data+Utils.swift b/Utils/Sources/Utils/Data+Utils.swift
index 4508b6a6..6ecfaa83 100644
--- a/Utils/Sources/Utils/Data+Utils.swift
+++ b/Utils/Sources/Utils/Data+Utils.swift
@@ -3,7 +3,8 @@ import Foundation
 extension Data {
     public init?(fromHexString hexString: String) {
         guard !hexString.isEmpty else {
-            return nil
+            self.init()
+            return
         }
 
         var data = Data()
@@ -23,4 +24,8 @@ extension Data {
 
         self.init(data)
     }
+
+    public func toHexString() -> String {
+        map { String(format: "%02x", $0) }.joined()
+    }
 }
diff --git a/Utils/Sources/Utils/merklization.swift b/Utils/Sources/Utils/merklization.swift
new file mode 100644
index 00000000..3c16c855
--- /dev/null
+++ b/Utils/Sources/Utils/merklization.swift
@@ -0,0 +1,71 @@
+import Foundation
+
+public enum MerklizeError: Error {
+    case invalidIndex
+}
+
+/// State Merklization function from GP D.2
+///
+/// Input is serialized state defined in the GP D.1
+public func stateMerklize(kv: [Data32: Data], i: Int = 0) throws -> Data32 {
+    func branch(l: Data32, r: Data32) throws -> Data64 {
+        var data = l.data + r.data
+        data[0] = l.data[0] & 0xFE
+        return Data64(data)!
+    }
+
+    func embeddedLeaf(key: Data32, value: Data, size: UInt8) -> Data64 {
+        var data = Data()
+        data.reserveCapacity(64)
+        data.append(0b01 | (size << 2))
+        data += key.data[..<31]
+        data += value
+        data.append(contentsOf: repeatElement(0, count: 32 - Int(size)))
+        return Data64(data)!
+    }
+
+    func regularLeaf(key: Data32, value: Data) throws -> Data64 {
+        var data = Data()
+        data.reserveCapacity(64)
+        data.append(0b11)
+        data += key.data[..<31]
+        data += try blake2b256(value).data
+        return Data64(data)!
+    }
+
+    func leaf(key: Data32, value: Data) throws -> Data64 {
+        if value.count <= 32 {
+            return embeddedLeaf(key: key, value: value, size: UInt8(value.count))
+        } else {
+            return try regularLeaf(key: key, value: value)
+        }
+    }
+
+    /// bit at i, returns true if it is 1
+    func bit(_ data: Data, _ i: Int) throws -> Bool {
+        guard let byte = data[safe: i / 8] else {
+            throw MerklizeError.invalidIndex
+        }
+        return (byte & (1 << (i % 8))) != 0
+    }
+
+    if kv.isEmpty {
+        return Data32()
+    }
+
+    if kv.count == 1 {
+        return try blake2b256(leaf(key: kv.first!.key, value: kv.first!.value).data)
+    }
+
+    var l: [Data32: Data] = [:]
+    var r: [Data32: Data] = [:]
+    for (k, v) in kv {
+        if try bit(k.data, i) {
+            r[k] = v
+        } else {
+            l[k] = v
+        }
+    }
+
+    return try blake2b256(branch(l: stateMerklize(kv: l, i: i + 1), r: stateMerklize(kv: r, i: i + 1)).data)
+}