From 7119192b674f632fc7e255cd2d05841b8a631084 Mon Sep 17 00:00:00 2001 From: Qiwei Yang Date: Wed, 14 Aug 2024 08:40:57 +0800 Subject: [PATCH 1/2] wip --- JAMTests/Tests/JAMTests/TrieTests.swift | 37 +++++++++++++ Utils/Sources/Utils/merklization.swift | 74 +++++++++++++++++++++++++ 2 files changed, 111 insertions(+) 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..90faf47c --- /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 == Data(fromHexString: element.output)) + } + } +} diff --git a/Utils/Sources/Utils/merklization.swift b/Utils/Sources/Utils/merklization.swift new file mode 100644 index 00000000..6862cefd --- /dev/null +++ b/Utils/Sources/Utils/merklization.swift @@ -0,0 +1,74 @@ +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]) throws -> Data32 { + func branch(l: Data32, r: Data32) -> Data64 { + var data = l.data + r.data + data[0] = l.data[0] & 0xFE + assert(data.count == 64, "branch data should be 64 bytes") + return Data64(data)! + } + + func embeddedLeaf(key: Data32, value: Data, size: UInt8) -> Data64 { + var data = Data() + data.reserveCapacity(64) + data[0] = 0b01 | (size << 2) + data += key.data[..<31] + data += value + data.append(contentsOf: repeatElement(0, count: 32 - Int(size))) + assert(data.count == 64, "embeddedLeaf data should be 64 bytes") + return Data64(data)! + } + + func regularLeaf(key: Data32, value: Data) throws -> Data64 { + var data = Data() + data.reserveCapacity(64) + data[0] = 0b11 + data += key.data[..<31] + data += try blake2b256(value).data + assert(data.count == 64, "regularLeaf data should be 64 bytes") + return Data64(data)! + } + + func leaf(key: Data32, value: Data) throws -> Data64 { + if value.count <= 32 { + embeddedLeaf(key: key, value: value, size: UInt8(value.count)) + } else { + try regularLeaf(key: key, value: value) + } + } + + /// bit at index i, returns true if it is 1 + func bit(_ bits: Data, _ i: Int) throws -> Bool { + guard let byte = bits[safe: i / 8] else { + throw MerklizeError.invalidIndex + } + return (byte & (1 << (7 - i % 8))) == 1 + } + + 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, 0) { + r[k] = v + } else { + l[k] = v + } + } + + return try blake2b256(branch(l: stateMerklize(kv: l), r: stateMerklize(kv: r)).data) +} From f3c66f3201067d5b8a98cec2e293e867d47c5582 Mon Sep 17 00:00:00 2001 From: Qiwei Yang Date: Wed, 14 Aug 2024 16:06:42 +0800 Subject: [PATCH 2/2] fix --- JAMTests/Tests/JAMTests/TrieTests.swift | 2 +- Utils/Sources/Utils/Data+Utils.swift | 7 ++++++- Utils/Sources/Utils/merklization.swift | 27 +++++++++++-------------- 3 files changed, 19 insertions(+), 17 deletions(-) diff --git a/JAMTests/Tests/JAMTests/TrieTests.swift b/JAMTests/Tests/JAMTests/TrieTests.swift index 90faf47c..c0097526 100644 --- a/JAMTests/Tests/JAMTests/TrieTests.swift +++ b/JAMTests/Tests/JAMTests/TrieTests.swift @@ -31,7 +31,7 @@ struct TrieTests { } let result = try stateMerklize(kv: kv) - #expect(result.data == Data(fromHexString: element.output)) + #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 index 6862cefd..3c16c855 100644 --- a/Utils/Sources/Utils/merklization.swift +++ b/Utils/Sources/Utils/merklization.swift @@ -7,49 +7,46 @@ public enum MerklizeError: Error { /// State Merklization function from GP D.2 /// /// Input is serialized state defined in the GP D.1 -public func stateMerklize(kv: [Data32: Data]) throws -> Data32 { - func branch(l: Data32, r: Data32) -> Data64 { +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 - assert(data.count == 64, "branch data should be 64 bytes") return Data64(data)! } func embeddedLeaf(key: Data32, value: Data, size: UInt8) -> Data64 { var data = Data() data.reserveCapacity(64) - data[0] = 0b01 | (size << 2) + data.append(0b01 | (size << 2)) data += key.data[..<31] data += value data.append(contentsOf: repeatElement(0, count: 32 - Int(size))) - assert(data.count == 64, "embeddedLeaf data should be 64 bytes") return Data64(data)! } func regularLeaf(key: Data32, value: Data) throws -> Data64 { var data = Data() data.reserveCapacity(64) - data[0] = 0b11 + data.append(0b11) data += key.data[..<31] data += try blake2b256(value).data - assert(data.count == 64, "regularLeaf data should be 64 bytes") return Data64(data)! } func leaf(key: Data32, value: Data) throws -> Data64 { if value.count <= 32 { - embeddedLeaf(key: key, value: value, size: UInt8(value.count)) + return embeddedLeaf(key: key, value: value, size: UInt8(value.count)) } else { - try regularLeaf(key: key, value: value) + return try regularLeaf(key: key, value: value) } } - /// bit at index i, returns true if it is 1 - func bit(_ bits: Data, _ i: Int) throws -> Bool { - guard let byte = bits[safe: i / 8] else { + /// 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 << (7 - i % 8))) == 1 + return (byte & (1 << (i % 8))) != 0 } if kv.isEmpty { @@ -63,12 +60,12 @@ public func stateMerklize(kv: [Data32: Data]) throws -> Data32 { var l: [Data32: Data] = [:] var r: [Data32: Data] = [:] for (k, v) in kv { - if try bit(k.data, 0) { + if try bit(k.data, i) { r[k] = v } else { l[k] = v } } - return try blake2b256(branch(l: stateMerklize(kv: l), r: stateMerklize(kv: r)).data) + return try blake2b256(branch(l: stateMerklize(kv: l, i: i + 1), r: stateMerklize(kv: r, i: i + 1)).data) }