Preimages tests (#263)

* preimages

* preimages test

Blockchain/Sources/Blockchain/RuntimeProtocols/Preimages.swift

@@ -0,0 +1,78 @@
+import Foundation
+import Utils
+
+public enum PreimagesError: Error {
+    case preimagesNotSorted
+    case duplicatedPreimage
+    case invalidServiceIndex
+}
+
+public struct PreimageUpdate: Sendable, Equatable {
+    public let serviceIndex: UInt32
+    public let hash: Data32
+    public let data: Data
+    public let length: UInt32
+    public let timeslot: TimeslotIndex
+
+    public init(serviceIndex: UInt32, hash: Data32, data: Data, length: UInt32, timeslot: TimeslotIndex) {
+        self.serviceIndex = serviceIndex
+        self.hash = hash
+        self.data = data
+        self.length = length
+        self.timeslot = timeslot
+    }
+}
+
+public struct PreimagesPostState: Sendable, Equatable {
+    public let updates: [PreimageUpdate]
+
+    public init(updates: [PreimageUpdate]) {
+        self.updates = updates
+    }
+}
+
+public protocol Preimages {
+    func get(serviceAccount index: ServiceIndex, preimageHash hash: Data32) async throws -> Data?
+    func get(serviceAccount index: ServiceIndex, preimageHash hash: Data32, length: UInt32) async throws
+        -> LimitedSizeArray<TimeslotIndex, ConstInt0, ConstInt3>?
+
+    mutating func mergeWith(postState: PreimagesPostState)
+}
+
+extension Preimages {
+    public func updatePreimages(
+        config _: ProtocolConfigRef,
+        timeslot: TimeslotIndex,
+        preimages: ExtrinsicPreimages
+    ) async throws(PreimagesError) -> PreimagesPostState {
+        let preimages = preimages.preimages
+        var updates: [PreimageUpdate] = []
+
+        guard preimages.isSortedAndUnique() else {
+            throw PreimagesError.preimagesNotSorted
+        }
+
+        for preimage in preimages {
+            let hash = preimage.data.blake2b256hash()
+
+            // check prior state
+            let prevPreimageData = try await Result {
+                try await get(serviceAccount: preimage.serviceIndex, preimageHash: hash)
+            }.mapError { _ in PreimagesError.invalidServiceIndex }.get()
+
+            guard prevPreimageData == nil else {
+                throw PreimagesError.duplicatedPreimage
+            }
+
+            updates.append(PreimageUpdate(
+                serviceIndex: preimage.serviceIndex,
+                hash: hash,
+                data: preimage.data,
+                length: UInt32(preimage.data.count),
+                timeslot: timeslot
+            ))
+        }
+
+        return PreimagesPostState(updates: updates)
+    }
+}

Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift

@@ -336,38 +336,9 @@ public final class Runtime {
     }
 
     public func updatePreimages(block: BlockRef, state newState: inout State, prevState: StateRef) async throws {
-        let preimages = block.extrinsic.preimages.preimages
-
-        guard preimages.isSortedAndUnique() else {
-            throw Error.preimagesNotSorted
-        }
-
-        for preimage in preimages {
-            let hash = preimage.data.blake2b256hash()
-
-            // check prior state
-            let prevPreimageData: Data? = try await prevState.value.get(serviceAccount: preimage.serviceIndex, preimageHash: hash)
-            let prevInfo = try await prevState.value.get(
-                serviceAccount: preimage.serviceIndex, preimageHash: hash, length: UInt32(preimage.data.count)
-            )
-            guard prevPreimageData == nil, prevInfo == nil else {
-                throw Error.duplicatedPreimage
-            }
-
-            // disregard no longer useful ones in new state
-            let preimageData: Data? = try await newState.get(serviceAccount: preimage.serviceIndex, preimageHash: hash)
-            let info = try await newState.get(
-                serviceAccount: preimage.serviceIndex, preimageHash: hash, length: UInt32(preimage.data.count)
-            )
-            if preimageData != nil || info != nil {
-                continue
-            }
-
-            // update state
-            newState[serviceAccount: preimage.serviceIndex, preimageHash: hash] = preimage.data
-            newState[
-                serviceAccount: preimage.serviceIndex, preimageHash: hash, length: UInt32(preimage.data.count)
-            ] = .init([newState.timeslot])
-        }
+        let res = try await prevState.value.updatePreimages(
+            config: config, timeslot: newState.timeslot, preimages: block.extrinsic.preimages
+        )
+        newState.mergeWith(postState: res)
     }
 }

Blockchain/Sources/Blockchain/State/State.swift

@@ -471,6 +471,16 @@ extension State: Authorization {
 
 extension State: ActivityStatistics {}
 
+extension State: Preimages {
+    public mutating func mergeWith(postState: PreimagesPostState) {
+        for update in postState.updates {
+            self[serviceAccount: update.serviceIndex, preimageHash: update.hash] = update.data
+            self[serviceAccount: update.serviceIndex, preimageHash: update.hash, length: update.length] =
+                LimitedSizeArray([update.timeslot])
+        }
+    }
+}
+
 struct DummyFunction: AccumulateFunction, OnTransferFunction {
     func invoke(
         config _: ProtocolConfigRef,

Blockchain/Sources/Blockchain/Types/HashAndLength.swift

@@ -1,13 +1,20 @@
 import Utils
 
-public struct HashAndLength: Sendable, Codable {
+public struct HashAndLength: Sendable, Codable, Comparable {
     public var hash: Data32
     public var length: DataLength
 
     public init(hash: Data32, length: DataLength) {
         self.hash = hash
         self.length = length
     }
+
+    public static func < (lhs: HashAndLength, rhs: HashAndLength) -> Bool {
+        if lhs.hash == rhs.hash {
+            return lhs.length < rhs.length
+        }
+        return lhs.hash < rhs.hash
+    }
 }
 
 extension HashAndLength: Hashable {
@@ -16,6 +23,6 @@ extension HashAndLength: Hashable {
         // and we know the output is 32 bytes
         // so we can just take the first 4 bytes and should be good enough
         // NOTE: we will never use the Hashable protocol for any critical operations
-        hasher.combine(hash.data[0 ..< 4])
+        hasher.combine(hash.data[hash.data.startIndex ..< hash.data.startIndex + 4])
     }
 }

JAMTests/Tests/JAMTests/PreimagesTests.swift

@@ -0,0 +1,122 @@
+import Blockchain
+import Codec
+import Foundation
+import Testing
+import Utils
+
+@testable import JAMTests
+
+struct PreimageInfo: Codable, Equatable, Hashable, Comparable {
+    var hash: Data32
+    var blob: Data
+
+    static func < (lhs: PreimageInfo, rhs: PreimageInfo) -> Bool {
+        lhs.hash < rhs.hash
+    }
+}
+
+struct HistoryEntry: Codable, Equatable {
+    var key: HashAndLength
+    var value: [TimeslotIndex]
+}
+
+struct AccountsMapEntry: Codable, Equatable {
+    var index: ServiceIndex
+    @CodingAs<SortedSet<PreimageInfo>> var preimages: Set<PreimageInfo>
+    @CodingAs<SortedKeyValues<HashAndLength, [TimeslotIndex]>> var history: [HashAndLength: [TimeslotIndex]]
+}
+
+struct PreimagesState: Equatable, Codable, Preimages {
+    var accounts: [AccountsMapEntry] = []
+
+    func get(serviceAccount index: ServiceIndex, preimageHash hash: Data32) async throws -> Data? {
+        for account in accounts where account.index == index {
+            for preimage in account.preimages where preimage.hash == hash {
+                return preimage.blob
+            }
+        }
+        return nil
+    }
+
+    func get(serviceAccount index: ServiceIndex, preimageHash hash: Data32,
+             length: UInt32) async throws -> LimitedSizeArray<TimeslotIndex, ConstInt0, ConstInt3>?
+    {
+        for account in accounts where account.index == index {
+            for history in account.history where history.key.hash == hash && history.key.length == length {
+                return .init(history.value)
+            }
+        }
+        return nil
+    }
+
+    mutating func mergeWith(postState: PreimagesPostState) {
+        for update in postState.updates {
+            let accountIndex = accounts.firstIndex { account in
+                account.index == update.serviceIndex
+            }
+            if let accountIndex {
+                var account = accounts[accountIndex]
+                account.preimages.insert(PreimageInfo(hash: update.hash, blob: update.data))
+                account.history[HashAndLength(hash: update.hash, length: update.length)] = [update.timeslot]
+                accounts[accountIndex] = account
+            }
+        }
+    }
+}
+
+struct PreimagesInput: Codable {
+    var preimages: ExtrinsicPreimages
+    var slot: TimeslotIndex
+}
+
+struct PreimagesTestcase: Codable {
+    var input: PreimagesInput
+    var preState: PreimagesState
+    var output: UInt8?
+    var postState: PreimagesState
+}
+
+struct PreimagesTests {
+    static func loadTests() throws -> [Testcase] {
+        try TestLoader.getTestcases(path: "preimages/data", extension: "bin")
+    }
+
+    func preimagesTests(_ testcase: Testcase, variant: TestVariants) async throws {
+        let config = variant.config
+        let decoder = JamDecoder(data: testcase.data, config: config)
+        let testcase = try decoder.decode(PreimagesTestcase.self)
+
+        var state = testcase.preState
+        let result = await Result {
+            try await state.updatePreimages(
+                config: config,
+                timeslot: testcase.input.slot,
+                preimages: testcase.input.preimages
+            )
+        }
+
+        switch result {
+        case let .success(postState):
+            switch testcase.output {
+            case .none:
+                state.mergeWith(postState: postState)
+                #expect(state == testcase.postState)
+            case .some:
+                Issue.record("Expected error, got \(result)")
+            }
+        case .failure:
+            switch testcase.output {
+            case .none:
+                Issue.record("Expected success, got \(result)")
+            case .some:
+                // ignore error code because it is unspecified
+                break
+            }
+        }
+    }
+
+    @Test(arguments: try PreimagesTests.loadTests())
+    func tests(_ testcase: Testcase) async throws {
+        try await preimagesTests(testcase, variant: .full)
+    }
+}