From a0f5c3757737bafc7199e3910f64609fb4c26f5e Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Wed, 6 Nov 2024 21:24:01 +1300 Subject: [PATCH 01/10] implements state trie --- .../Blockchain/State/InMemoryBackend.swift | 76 +++-- .../Sources/Blockchain/State/State.swift | 13 +- .../Blockchain/State/StateBackend.swift | 44 ++- .../State/StateBackendProtocol.swift | 18 ++ .../Sources/Blockchain/State/StateTrie.swift | 260 ++++++++++++++++++ Database/Sources/Database/RocksDB.swift | 3 +- Node/Package.resolved | 6 +- Node/Sources/Node/Node.swift | 3 +- Node/Tests/NodeTests/ChainSpecTests.swift | 3 +- .../Merklization/StateMerklization.swift | 2 +- 10 files changed, 376 insertions(+), 52 deletions(-) create mode 100644 Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift create mode 100644 Blockchain/Sources/Blockchain/State/StateTrie.swift diff --git a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift index 7716d086..bbd4b67c 100644 --- a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift +++ b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift @@ -2,41 +2,73 @@ import Codec import Foundation import Utils -public actor InMemoryBackend: StateBackend { - private let config: ProtocolConfigRef - private var store: [Data32: Data] +private struct KVPair: Comparable, Sendable { + var key: Data + var value: Data - public init(config: ProtocolConfigRef, store: [Data32: Data] = [:]) { - self.config = config - self.store = store + public static func < (lhs: KVPair, rhs: KVPair) -> Bool { + lhs.key.lexicographicallyPrecedes(rhs.key) } +} + +public actor InMemoryBackend: StateBackendProtocol { + private var store: SortedArray + private var refCounts: [Data: Int] + + public init(store: [Data32: Data] = [:]) { + self.store = .init(store.map { KVPair(key: $0.key.data, value: $0.value) }) + refCounts = [:] - public func readImpl(_ key: any StateKey) async throws -> (Codable & Sendable)? { - guard let value = store[key.encode()] else { - return nil + for key in store.keys { + refCounts[key.data] = 1 } - return try JamDecoder.decode(key.decodeType(), from: value, withConfig: config) } - public func batchRead(_ keys: [any StateKey]) async throws -> [(key: any StateKey, value: Codable & Sendable)] { - try keys.map { - let data = try store[$0.encode()].unwrap() - return try ($0, JamDecoder.decode($0.decodeType(), from: data, withConfig: config)) + public func read(key: Data) async throws -> Data? { + let idx = store.insertIndex(KVPair(key: key, value: Data())) + let item = store.array[safe: idx] + if item?.key == key { + return item?.value } + return nil } - public func batchWrite(_ changes: [(key: any StateKey, value: Codable & Sendable)]) async throws { - for (key, value) in changes { - store[key.encode()] = try JamEncoder.encode(value) + public func readAll(prefix: Data, startKey: Data?, limit: UInt32?) async throws -> [(key: Data, value: Data)] { + var resp = [(key: Data, value: Data)]() + let startKey = startKey ?? prefix + let startIndex = store.insertIndex(KVPair(key: startKey, value: Data())) + for i in startIndex ..< store.array.count { + let item = store.array[i] + if item.key.starts(with: prefix) { + resp.append((item.key, item.value)) + } else { + break + } + if let limit, resp.count == limit { + break + } } + return resp } - public func readAll() async throws -> [Data32: Data] { - store + public func batchRead(keys: [Data]) async throws -> [(key: Data, value: Data?)] { + var resp = [(key: Data, value: Data?)]() + for key in keys { + let value = try await read(key: key) + resp.append((key, value)) + } + return resp } - public func stateRoot() async throws -> Data32 { - // TODO: store intermediate state so we can calculate the root efficiently - try stateMerklize(kv: store) + public func batchUpdate(_: [StateBackendOperation]) async throws {} + + public func gc() async throws { + // check ref counts and remove keys with 0 ref count + for (key, count) in refCounts where count == 0 { + let idx = store.insertIndex(KVPair(key: key, value: Data())) + if store.array[safe: idx]?.key == key { + store.remove(at: idx) + } + } } } diff --git a/Blockchain/Sources/Blockchain/State/State.swift b/Blockchain/Sources/Blockchain/State/State.swift index dab62028..4baba00e 100644 --- a/Blockchain/Sources/Blockchain/State/State.swift +++ b/Blockchain/Sources/Blockchain/State/State.swift @@ -197,9 +197,8 @@ public struct State: Sendable { } } - public func stateRoot() -> Data32 { - // TODO: incorporate layer changes and calculate state root - Data32() + public var stateRoot: Data32 { + backend.rootHash } } @@ -324,11 +323,9 @@ extension State: Dummy { for (key, value) in kv { store[key.encode()] = try! JamEncoder.encode(value) } + let rootHash = try! stateMerklize(kv: store) - let backend = InMemoryBackend( - config: config, - store: store - ) + let backend = StateBackend(InMemoryBackend(store: store), config: config, rootHash: rootHash) let layer = StateLayer(changes: kv) @@ -484,6 +481,6 @@ public class StateRef: Ref, @unchecked Sendable { } public var stateRoot: Data32 { - value.stateRoot() + value.stateRoot } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackend.swift b/Blockchain/Sources/Blockchain/State/StateBackend.swift index 8c75f329..77ec799b 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackend.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackend.swift @@ -1,28 +1,42 @@ +import Codec import Foundation import Utils public enum StateBackendError: Error { case missingState + case invalidData } -public protocol StateBackend: Sendable { - func readImpl(_ key: any StateKey) async throws -> (Codable & Sendable)? +public final class StateBackend: Sendable { + private let impl: StateBackendProtocol + private let config: ProtocolConfigRef + public let rootHash: Data32 - func batchRead(_ keys: [any StateKey]) async throws -> [(key: any StateKey, value: Codable & Sendable)] - mutating func batchWrite(_ changes: [(key: any StateKey, value: Codable & Sendable)]) async throws - - func readAll() async throws -> [Data32: Data] - - func stateRoot() async throws -> Data32 - - // TODO: aux store for full key and intermidate merkle root -} + public init(_ impl: StateBackendProtocol, config: ProtocolConfigRef, rootHash: Data32) { + self.impl = impl + self.config = config + self.rootHash = rootHash + } -extension StateBackend { public func read(_ key: Key) async throws -> Key.Value.ValueType { - guard let ret = try await readImpl(key) as? Key.Value.ValueType else { - throw StateBackendError.missingState + let encodedKey = key.encode().data + if let ret = try await impl.read(key: encodedKey) { + guard let ret = try JamDecoder.decode(key.decodeType(), from: ret, withConfig: config) as? Key.Value.ValueType else { + throw StateBackendError.invalidData + } + return ret } - return ret + if Key.Value.optional { + return Key.Value.DecodeType?.none as! Key.Value.ValueType + } + throw StateBackendError.missingState + } + + func batchRead(_: [any StateKey]) async throws -> [(key: any StateKey, value: Codable & Sendable)] { + [] + } + + func readAll() async throws -> [Data32: Data] { + [:] } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift new file mode 100644 index 00000000..630b2003 --- /dev/null +++ b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift @@ -0,0 +1,18 @@ +import Foundation +import Utils + +public enum StateBackendOperation: Sendable { + case write(key: Data, value: Data) + case refIncrement(key: Data) + case refDecrement(key: Data) +} + +public protocol StateBackendProtocol: Sendable { + func read(key: Data) async throws -> Data? + func readAll(prefix: Data, startKey: Data?, limit: UInt32?) async throws -> [(key: Data, value: Data)] + func batchRead(keys: [Data]) async throws -> [(key: Data, value: Data?)] + func batchUpdate(_ ops: [StateBackendOperation]) async throws + + // remove entries with zero ref count + func gc() async throws +} diff --git a/Blockchain/Sources/Blockchain/State/StateTrie.swift b/Blockchain/Sources/Blockchain/State/StateTrie.swift new file mode 100644 index 00000000..6165accf --- /dev/null +++ b/Blockchain/Sources/Blockchain/State/StateTrie.swift @@ -0,0 +1,260 @@ +import Foundation +import Utils + +private enum TrieNodeType { + case branch + case embeddedLeaf + case regularLeaf +} + +private struct TrieNode { + var hash: Data32 + var left: Data32 + var right: Data32 + var type: TrieNodeType + var isNew: Bool + + init(hash: Data32, data: Data64, isNew: Bool = false) { + self.hash = hash + left = Data32(data.data.prefix(32))! + right = Data32(data.data.suffix(32))! + self.isNew = isNew + switch data.data[0] & 0b1100_0000 { + case 0b1000_0000: + type = .embeddedLeaf + case 0b1100_0000: + type = .regularLeaf + default: + type = .branch + } + } + + private init(left: Data32, right: Data32, type: TrieNodeType, isNew: Bool = false) { + hash = Blake2b256.hash(left.data, right.data) + self.left = left + self.right = right + self.type = type + self.isNew = isNew + } + + var encodedData: Data64 { + Data64(left.data + right.data)! + } + + var isBranch: Bool { + type == .branch + } + + static func leaf(key: Data32, value: Data) -> TrieNode { + var newKey = Data(capacity: 32) + if value.count <= 32 { + newKey.append(0b1000_0000 | UInt8(value.count)) + newKey += key.data.prefix(31) + let newValue = value + Data(repeating: 0, count: 32 - value.count) + return .init(left: Data32(newKey)!, right: Data32(newValue)!, type: .embeddedLeaf, isNew: true) + } + newKey.append(0b1100_0000) + newKey += key.data.prefix(31) + return .init(left: Data32(newKey)!, right: value.blake2b256hash(), type: .regularLeaf, isNew: true) + } + + static func branch(left: Data32, right: Data32) -> TrieNode { + .init(left: left, right: right, type: .branch, isNew: true) + } +} + +public enum StateTrieError: Error { + case invalidData + case invalidParent +} + +public class StateTrie { + private let backend: StateBackendProtocol + public private(set) var rootHash: Data32 + private var nodes: [Data32: TrieNode] = [:] + private var deleted: Set = [] + + public init(rootHash: Data32, backend: StateBackendProtocol) { + self.rootHash = rootHash + self.backend = backend + } + + private func get(hash: Data32) async throws -> TrieNode? { + if deleted.contains(hash) { + return nil + } + if let node = nodes[hash] { + return node + } + guard let data = try await backend.read(key: hash.data) else { + return nil + } + + guard let data64 = Data64(data) else { + throw StateTrieError.invalidData + } + + let node = TrieNode(hash: hash, data: data64) + saveNode(node: node) + return node + } + + public func update(_ updates: [(key: Data32, value: Data?)]) async throws { + // TODO: somehow improve the efficiency of this + for (key, value) in updates { + if let value { + rootHash = try await insert(hash: rootHash, key: key, value: value, depth: 0) + } else { + rootHash = try await delete(hash: rootHash, key: key, depth: 0) + } + } + } + + public func save() async throws { + var ops = [StateBackendOperation]() + var refChanges = [Data: Int]() + + // process deleted nodes + for hash in deleted { + let node = try await get(hash: hash).unwrap() + if node.isBranch { + // assign -1 to not worry about duplicates + refChanges[node.hash.data] = -1 + refChanges[node.left.data] = -1 + refChanges[node.right.data] = -1 + } + nodes.removeValue(forKey: hash) + } + deleted.removeAll() + + for node in nodes.values where node.isNew { + ops.append(.write(key: node.hash.data, value: node.encodedData.data)) + if node.isBranch { + refChanges[node.left.data] = (refChanges[node.left.data] ?? 0) + 1 + refChanges[node.right.data] = (refChanges[node.right.data] ?? 0) + 1 + } + } + + // pin root node + refChanges[rootHash.data] = (refChanges[rootHash.data] ?? 0) + 1 + + nodes.removeAll() + + for (key, value) in refChanges { + if value > 0 { + ops.append(.refIncrement(key: key)) + } else if value < 0 { + ops.append(.refDecrement(key: key)) + } + } + + try await backend.batchUpdate(ops) + } + + private func insert( + hash: Data32, key: Data32, value: Data, depth: UInt8 + ) async throws -> Data32 { + let parent = try await get(hash: hash).unwrap(orError: StateTrieError.invalidParent) + removeNode(hash: hash) + + if parent.isBranch { + let bitValue = bitAt(key.data, position: depth) + var left = parent.left + var right = parent.right + if bitValue { + right = try await insert(hash: parent.right, key: key, value: value, depth: depth + 1) + } else { + left = try await insert(hash: parent.left, key: key, value: value, depth: depth + 1) + } + let newBranch = TrieNode.branch(left: left, right: right) + saveNode(node: newBranch) + return newBranch.hash + } else { + // leaf + return try await insertLeafNode(existing: parent, newKey: key, newValue: value, depth: depth) + } + } + + private func insertLeafNode(existing: TrieNode, newKey: Data32, newValue: Data, depth: UInt8) async throws -> Data32 { + if existing.left == newKey { + // update existing leaf + let newLeaf = TrieNode.leaf(key: newKey, value: newValue) + nodes[newLeaf.hash] = newLeaf + return newLeaf.hash + } + + let existingKeyBit = bitAt(existing.left.data, position: depth) + let newKeyBit = bitAt(newKey.data, position: depth) + + if existingKeyBit == newKeyBit { + // need to go deeper + let childNodeHash = try await insertLeafNode( + existing: existing, newKey: newKey, newValue: newValue, depth: depth + 1 + ) + let newBranch = if existingKeyBit { + TrieNode.branch(left: Data32(), right: childNodeHash) + } else { + TrieNode.branch(left: childNodeHash, right: Data32()) + } + saveNode(node: newBranch) + return newBranch.hash + } else { + let newLeaf = TrieNode.leaf(key: newKey, value: newValue) + saveNode(node: newLeaf) + let newBranch = if existingKeyBit { + TrieNode.branch(left: newLeaf.hash, right: existing.hash) + } else { + TrieNode.branch(left: existing.hash, right: newLeaf.hash) + } + saveNode(node: newBranch) + return newBranch.hash + } + } + + private func delete(hash: Data32, key: Data32, depth: UInt8) async throws -> Data32 { + let node = try await get(hash: hash).unwrap(orError: StateTrieError.invalidParent) + removeNode(hash: hash) + + if node.isBranch { + let bitValue = bitAt(key.data, position: depth) + var left = node.left + var right = node.right + + if bitValue { + right = try await delete(hash: node.right, key: key, depth: depth + 1) + } else { + left = try await delete(hash: node.left, key: key, depth: depth + 1) + } + + if left == Data32(), right == Data32() { + // this branch is empty + return Data32() + } + + let newBranch = TrieNode.branch(left: left, right: right) + saveNode(node: newBranch) + return newBranch.hash + } else { + // leaf + return Data32() + } + } + + private func removeNode(hash: Data32) { + deleted.insert(hash) + nodes.removeValue(forKey: hash) + } + + private func saveNode(node: TrieNode) { + nodes[node.hash] = node + deleted.remove(node.hash) // TODO: maybe this is not needed + } +} + +/// bit at i, returns true if it is 1 +private func bitAt(_ data: Data, position: UInt8) -> Bool { + let byteIndex = position / 8 + let bitIndex = 7 - (position % 8) + let byte = data[safeRelative: Int(byteIndex)] ?? 0 + return (byte & (1 << bitIndex)) != 0 +} diff --git a/Database/Sources/Database/RocksDB.swift b/Database/Sources/Database/RocksDB.swift index 51d1c9f5..dd73f824 100644 --- a/Database/Sources/Database/RocksDB.swift +++ b/Database/Sources/Database/RocksDB.swift @@ -28,7 +28,8 @@ public final class RocksDB { // Optimize rocksdb rocksdb_options_increase_parallelism(dbOptions, Int32(cpus)) - rocksdb_options_optimize_level_style_compaction(dbOptions, 0) // TODO: check this + let memtable_memory_budget: UInt64 = 512 * 1024 * 1024 // 512 MB + rocksdb_options_optimize_level_style_compaction(dbOptions, memtable_memory_budget) // create the DB if it's not already present rocksdb_options_set_create_if_missing(dbOptions, 1) diff --git a/Node/Package.resolved b/Node/Package.resolved index f1dd53fc..21e00dc6 100644 --- a/Node/Package.resolved +++ b/Node/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "33177d9b5bd122b13d37a99e2337556cfc5a1b7e466229e7bdb19fd9423a117a", + "originHash" : "e6fc7ac1513fbfe8e482dc8b89ef5388fcc4ffb7a67e2def60644a806029f64e", "pins" : [ { "identity" : "async-channels", @@ -105,8 +105,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/apple/swift-collections.git", "state" : { - "revision" : "3d2dc41a01f9e49d84f0a3925fb858bed64f702d", - "version" : "1.1.2" + "revision" : "671108c96644956dddcd89dd59c203dcdb36cec7", + "version" : "1.1.4" } }, { diff --git a/Node/Sources/Node/Node.swift b/Node/Sources/Node/Node.swift index 503b63a5..09a93907 100644 --- a/Node/Sources/Node/Node.swift +++ b/Node/Sources/Node/Node.swift @@ -44,7 +44,8 @@ public class Node { let chainspec = try await genesis.load() let genesisBlock = try chainspec.getBlock() let genesisStateData = try chainspec.getState() - let backend = try InMemoryBackend(config: chainspec.getConfig(), store: genesisStateData) + let rootHash = try stateMerklize(kv: genesisStateData) + let backend = try StateBackend(InMemoryBackend(store: genesisStateData), config: chainspec.getConfig(), rootHash: rootHash) let genesisState = try await State(backend: backend) let genesisStateRef = StateRef(genesisState) let protocolConfig = try chainspec.getConfig() diff --git a/Node/Tests/NodeTests/ChainSpecTests.swift b/Node/Tests/NodeTests/ChainSpecTests.swift index 52c14bd1..10a04055 100644 --- a/Node/Tests/NodeTests/ChainSpecTests.swift +++ b/Node/Tests/NodeTests/ChainSpecTests.swift @@ -1,6 +1,7 @@ import Blockchain import Foundation import Testing +import Utils @testable import Node @@ -17,7 +18,7 @@ struct ChainSpecTests { for preset in GenesisPreset.allCases { let genesis = Genesis.preset(preset) let chainspec = try await genesis.load() - let backend = try InMemoryBackend(config: chainspec.getConfig(), store: chainspec.getState()) + let backend = try StateBackend(InMemoryBackend(store: chainspec.getState()), config: chainspec.getConfig(), rootHash: Data32()) let block = try chainspec.getBlock() let config = try chainspec.getConfig() diff --git a/Utils/Sources/Utils/Merklization/StateMerklization.swift b/Utils/Sources/Utils/Merklization/StateMerklization.swift index 1c852b25..b212e6d5 100644 --- a/Utils/Sources/Utils/Merklization/StateMerklization.swift +++ b/Utils/Sources/Utils/Merklization/StateMerklization.swift @@ -10,7 +10,7 @@ public enum MerklizeError: Error { public func stateMerklize(kv: [Data32: Data], i: Int = 0) throws(MerklizeError) -> Data32 { func branch(l: Data32, r: Data32) -> Data64 { var data = l.data + r.data - data[0] = l.data[0] & 0x7F + data[0] = l.data[0] & 0x7F // clear the highest bit return Data64(data)! } From 12438455043ba6cf53f1bfb31b3a88231cd4925b Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Thu, 7 Nov 2024 13:48:53 +1300 Subject: [PATCH 02/10] wip --- .../Blockchain/State/ServiceAccounts.swift | 10 +- .../Sources/Blockchain/State/State.swift | 106 ++++++------ .../Blockchain/State/StateBackend.swift | 49 +++++- .../State/StateBackendProtocol.swift | 8 +- .../Sources/Blockchain/State/StateKeys.swift | 123 +++++++------ .../Sources/Blockchain/State/StateLayer.swift | 162 +++++++++--------- .../Sources/Blockchain/State/StateTrie.swift | 8 +- .../Blockchain/Validator/BlockAuthor.swift | 2 +- 8 files changed, 253 insertions(+), 215 deletions(-) diff --git a/Blockchain/Sources/Blockchain/State/ServiceAccounts.swift b/Blockchain/Sources/Blockchain/State/ServiceAccounts.swift index 879f89bd..624e4323 100644 --- a/Blockchain/Sources/Blockchain/State/ServiceAccounts.swift +++ b/Blockchain/Sources/Blockchain/State/ServiceAccounts.swift @@ -7,15 +7,15 @@ public protocol ServiceAccounts { func get(serviceAccount index: ServiceIndex, preimageHash hash: Data32) async throws -> Data? func get( serviceAccount index: ServiceIndex, preimageHash hash: Data32, length: UInt32 - ) async throws -> StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType + ) async throws -> StateKeys.ServiceAccountPreimageInfoKey.Value? - mutating func set(serviceAccount index: ServiceIndex, account: ServiceAccountDetails) - mutating func set(serviceAccount index: ServiceIndex, storageKey key: Data32, value: Data) - mutating func set(serviceAccount index: ServiceIndex, preimageHash hash: Data32, value: Data) + mutating func set(serviceAccount index: ServiceIndex, account: ServiceAccountDetails?) + mutating func set(serviceAccount index: ServiceIndex, storageKey key: Data32, value: Data?) + mutating func set(serviceAccount index: ServiceIndex, preimageHash hash: Data32, value: Data?) mutating func set( serviceAccount index: ServiceIndex, preimageHash hash: Data32, length: UInt32, - value: StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType + value: StateKeys.ServiceAccountPreimageInfoKey.Value? ) } diff --git a/Blockchain/Sources/Blockchain/State/State.swift b/Blockchain/Sources/Blockchain/State/State.swift index 4baba00e..d5ab6a62 100644 --- a/Blockchain/Sources/Blockchain/State/State.swift +++ b/Blockchain/Sources/Blockchain/State/State.swift @@ -17,7 +17,7 @@ public struct State: Sendable { } // α: The core αuthorizations pool. - public var coreAuthorizationPool: StateKeys.CoreAuthorizationPoolKey.Value.ValueType { + public var coreAuthorizationPool: StateKeys.CoreAuthorizationPoolKey.Value { get { layer.coreAuthorizationPool } @@ -27,7 +27,7 @@ public struct State: Sendable { } // φ: The authorization queue. - public var authorizationQueue: StateKeys.AuthorizationQueueKey.Value.ValueType { + public var authorizationQueue: StateKeys.AuthorizationQueueKey.Value { get { layer.authorizationQueue } @@ -37,7 +37,7 @@ public struct State: Sendable { } // β: Information on the most recent βlocks. - public var recentHistory: StateKeys.RecentHistoryKey.Value.ValueType { + public var recentHistory: StateKeys.RecentHistoryKey.Value { get { layer.recentHistory } @@ -47,7 +47,7 @@ public struct State: Sendable { } // γ: State concerning Safrole. - public var safroleState: StateKeys.SafroleStateKey.Value.ValueType { + public var safroleState: StateKeys.SafroleStateKey.Value { get { layer.safroleState } @@ -57,7 +57,7 @@ public struct State: Sendable { } // ψ: past judgements - public var judgements: StateKeys.JudgementsKey.Value.ValueType { + public var judgements: StateKeys.JudgementsKey.Value { get { layer.judgements } @@ -67,7 +67,7 @@ public struct State: Sendable { } // η: The eηtropy accumulator and epochal raηdomness. - public var entropyPool: StateKeys.EntropyPoolKey.Value.ValueType { + public var entropyPool: StateKeys.EntropyPoolKey.Value { get { layer.entropyPool } @@ -77,7 +77,7 @@ public struct State: Sendable { } // ι: The validator keys and metadata to be drawn from next. - public var validatorQueue: StateKeys.ValidatorQueueKey.Value.ValueType { + public var validatorQueue: StateKeys.ValidatorQueueKey.Value { get { layer.validatorQueue } @@ -87,7 +87,7 @@ public struct State: Sendable { } // κ: The validator κeys and metadata currently active. - public var currentValidators: StateKeys.CurrentValidatorsKey.Value.ValueType { + public var currentValidators: StateKeys.CurrentValidatorsKey.Value { get { layer.currentValidators } @@ -97,7 +97,7 @@ public struct State: Sendable { } // λ: The validator keys and metadata which were active in the prior epoch. - public var previousValidators: StateKeys.PreviousValidatorsKey.Value.ValueType { + public var previousValidators: StateKeys.PreviousValidatorsKey.Value { get { layer.previousValidators } @@ -107,7 +107,7 @@ public struct State: Sendable { } // ρ: The ρending reports, per core, which are being made available prior to accumulation. - public var reports: StateKeys.ReportsKey.Value.ValueType { + public var reports: StateKeys.ReportsKey.Value { get { layer.reports } @@ -117,7 +117,7 @@ public struct State: Sendable { } // τ: The most recent block’s τimeslot. - public var timeslot: StateKeys.TimeslotKey.Value.ValueType { + public var timeslot: StateKeys.TimeslotKey.Value { get { layer.timeslot } @@ -127,7 +127,7 @@ public struct State: Sendable { } // χ: The privileged service indices. - public var privilegedServices: StateKeys.PrivilegedServicesKey.Value.ValueType { + public var privilegedServices: StateKeys.PrivilegedServicesKey.Value { get { layer.privilegedServices } @@ -137,7 +137,7 @@ public struct State: Sendable { } // π: The activity statistics for the validators. - public var activityStatistics: StateKeys.ActivityStatisticsKey.Value.ValueType { + public var activityStatistics: StateKeys.ActivityStatisticsKey.Value { get { layer.activityStatistics } @@ -147,7 +147,7 @@ public struct State: Sendable { } // δ: The (prior) state of the service accounts. - public subscript(serviceAccount index: ServiceIndex) -> StateKeys.ServiceAccountKey.Value.ValueType? { + public subscript(serviceAccount index: ServiceIndex) -> StateKeys.ServiceAccountKey.Value? { get { layer[serviceAccount: index] } @@ -157,7 +157,7 @@ public struct State: Sendable { } // s - public subscript(serviceAccount index: ServiceIndex, storageKey key: Data32) -> StateKeys.ServiceAccountStorageKey.Value.ValueType? { + public subscript(serviceAccount index: ServiceIndex, storageKey key: Data32) -> StateKeys.ServiceAccountStorageKey.Value? { get { layer[serviceAccount: index, storageKey: key] } @@ -169,7 +169,7 @@ public struct State: Sendable { // p public subscript( serviceAccount index: ServiceIndex, preimageHash hash: Data32 - ) -> StateKeys.ServiceAccountPreimagesKey.Value.ValueType? { + ) -> StateKeys.ServiceAccountPreimagesKey.Value? { get { layer[serviceAccount: index, preimageHash: hash] } @@ -181,7 +181,7 @@ public struct State: Sendable { // l public subscript( serviceAccount index: ServiceIndex, preimageHash hash: Data32, length length: UInt32 - ) -> StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType? { + ) -> StateKeys.ServiceAccountPreimageInfoKey.Value? { get { layer[serviceAccount: index, preimageHash: hash, length: length] } @@ -197,8 +197,12 @@ public struct State: Sendable { } } - public var stateRoot: Data32 { - backend.rootHash + public func save() async throws -> State { + let changes = layer.toKV() + let trie = StateTrie(rootHash: backend.rootHash, backend: backend) + try await trie.update(updates: changes) + let newRoot = trie.rootHash + return State(backend: backend.newBackend(rootHash: newRoot), layer: layer) } } @@ -211,13 +215,13 @@ extension State { typealias Element = (key: Data32, value: Data) let seq: any Sequence<(key: Data32, value: Data)> - let layer: [Data32: Data] + let layer: [Data32: Data?] init(state: State) async throws { seq = try await state.backend.readAll() - var layer = [Data32: Data]() + var layer = [Data32: Data?]() for (key, value) in state.layer.toKV() { - layer[key.encode()] = try JamEncoder.encode(value) + layer[key.encode()] = try value.map { try JamEncoder.encode($0) } } self.layer = layer } @@ -230,26 +234,32 @@ extension State { typealias Element = (key: Data32, value: Data) var iter: any IteratorProtocol - var layerIterator: (any IteratorProtocol)? - let layer: [Data32: Data] + var layerIterator: [Data32: Data?].Iterator? + let layer: [Data32: Data?] - init(iter: any IteratorProtocol, layer: [Data32: Data]) { + init(iter: any IteratorProtocol, layer: [Data32: Data?]) { self.iter = iter self.layer = layer } mutating func next() -> KVSequence.Iterator.Element? { if layerIterator != nil { - return layerIterator?.next() + if let (key, value) = layerIterator?.next() { + if let value { + return (key, value) + } + return next() // skip this one + } + return nil } if let (key, value) = iter.next() { - if layer[key] != nil { + if layer.keys.contains(key) { return next() // skip this one } return (key, value) } layerIterator = layer.makeIterator() - return layerIterator?.next() + return next() } } } @@ -271,9 +281,9 @@ extension State: Dummy { } public static func dummy(config: Config, block: BlockRef?) -> State { - let coreAuthorizationPool: StateKeys.CoreAuthorizationPoolKey.Value.ValueType = + let coreAuthorizationPool: StateKeys.CoreAuthorizationPoolKey.Value = try! ConfigFixedSizeArray(config: config, defaultValue: ConfigLimitedSizeArray(config: config)) - var recentHistory: StateKeys.RecentHistoryKey.Value.ValueType = RecentHistory.dummy(config: config) + var recentHistory: StateKeys.RecentHistoryKey.Value = RecentHistory.dummy(config: config) if let block { recentHistory.items.safeAppend(RecentHistory.HistoryItem( headerHash: block.hash, @@ -282,26 +292,26 @@ extension State: Dummy { workReportHashes: try! ConfigLimitedSizeArray(config: config) )) } - let safroleState: StateKeys.SafroleStateKey.Value.ValueType = SafroleState.dummy(config: config) - let entropyPool: StateKeys.EntropyPoolKey.Value.ValueType = EntropyPool((Data32(), Data32(), Data32(), Data32())) - let validatorQueue: StateKeys.ValidatorQueueKey.Value.ValueType = + let safroleState: StateKeys.SafroleStateKey.Value = SafroleState.dummy(config: config) + let entropyPool: StateKeys.EntropyPoolKey.Value = EntropyPool((Data32(), Data32(), Data32(), Data32())) + let validatorQueue: StateKeys.ValidatorQueueKey.Value = try! ConfigFixedSizeArray(config: config, defaultValue: ValidatorKey.dummy(config: config)) - let currentValidators: StateKeys.CurrentValidatorsKey.Value.ValueType = + let currentValidators: StateKeys.CurrentValidatorsKey.Value = try! ConfigFixedSizeArray(config: config, defaultValue: ValidatorKey.dummy(config: config)) - let previousValidators: StateKeys.PreviousValidatorsKey.Value.ValueType = + let previousValidators: StateKeys.PreviousValidatorsKey.Value = try! ConfigFixedSizeArray(config: config, defaultValue: ValidatorKey.dummy(config: config)) - let reports: StateKeys.ReportsKey.Value.ValueType = try! ConfigFixedSizeArray(config: config, defaultValue: nil) - let timeslot: StateKeys.TimeslotKey.Value.ValueType = block?.header.timeslot ?? 0 - let authorizationQueue: StateKeys.AuthorizationQueueKey.Value.ValueType = + let reports: StateKeys.ReportsKey.Value = try! ConfigFixedSizeArray(config: config, defaultValue: nil) + let timeslot: StateKeys.TimeslotKey.Value = block?.header.timeslot ?? 0 + let authorizationQueue: StateKeys.AuthorizationQueueKey.Value = try! ConfigFixedSizeArray(config: config, defaultValue: ConfigFixedSizeArray(config: config, defaultValue: Data32())) - let privilegedServices: StateKeys.PrivilegedServicesKey.Value.ValueType = PrivilegedServices( + let privilegedServices: StateKeys.PrivilegedServicesKey.Value = PrivilegedServices( empower: ServiceIndex(), assign: ServiceIndex(), designate: ServiceIndex(), basicGas: [:] ) - let judgements: StateKeys.JudgementsKey.Value.ValueType = JudgementsState.dummy(config: config) - let activityStatistics: StateKeys.ActivityStatisticsKey.Value.ValueType = ValidatorActivityStatistics.dummy(config: config) + let judgements: StateKeys.JudgementsKey.Value = JudgementsState.dummy(config: config) + let activityStatistics: StateKeys.ActivityStatisticsKey.Value = ValidatorActivityStatistics.dummy(config: config) let kv: [(any StateKey, Codable & Sendable)] = [ (StateKeys.CoreAuthorizationPoolKey(), coreAuthorizationPool), @@ -357,22 +367,22 @@ extension State: ServiceAccounts { public func get( serviceAccount index: ServiceIndex, preimageHash hash: Data32, length: UInt32 - ) async throws -> StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType { + ) async throws -> StateKeys.ServiceAccountPreimageInfoKey.Value? { if let res = layer[serviceAccount: index, preimageHash: hash, length: length] { return res } return try await backend.read(StateKeys.ServiceAccountPreimageInfoKey(index: index, hash: hash, length: length)) } - public mutating func set(serviceAccount index: ServiceIndex, account: ServiceAccountDetails) { + public mutating func set(serviceAccount index: ServiceIndex, account: ServiceAccountDetails?) { layer[serviceAccount: index] = account } - public mutating func set(serviceAccount index: ServiceIndex, storageKey key: Data32, value: Data) { + public mutating func set(serviceAccount index: ServiceIndex, storageKey key: Data32, value: Data?) { layer[serviceAccount: index, storageKey: key] = value } - public mutating func set(serviceAccount index: ServiceIndex, preimageHash hash: Data32, value: Data) { + public mutating func set(serviceAccount index: ServiceIndex, preimageHash hash: Data32, value: Data?) { layer[serviceAccount: index, preimageHash: hash] = value } @@ -380,7 +390,7 @@ extension State: ServiceAccounts { serviceAccount index: ServiceIndex, preimageHash hash: Data32, length: UInt32, - value: StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType + value: StateKeys.ServiceAccountPreimageInfoKey.Value? ) { layer[serviceAccount: index, preimageHash: hash, length: length] = value } @@ -479,8 +489,4 @@ public class StateRef: Ref, @unchecked Sendable { public static func dummy(config: ProtocolConfigRef, block: BlockRef?) -> StateRef { StateRef(State.dummy(config: config, block: block)) } - - public var stateRoot: Data32 { - value.stateRoot - } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackend.swift b/Blockchain/Sources/Blockchain/State/StateBackend.swift index 77ec799b..55c5a336 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackend.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackend.swift @@ -18,25 +18,58 @@ public final class StateBackend: Sendable { self.rootHash = rootHash } - public func read(_ key: Key) async throws -> Key.Value.ValueType { + public func read(_ key: Key) async throws -> Key.Value { let encodedKey = key.encode().data if let ret = try await impl.read(key: encodedKey) { - guard let ret = try JamDecoder.decode(key.decodeType(), from: ret, withConfig: config) as? Key.Value.ValueType else { + guard let ret = try JamDecoder.decode(key.decodeType(), from: ret, withConfig: config) as? Key.Value else { throw StateBackendError.invalidData } return ret } - if Key.Value.optional { - return Key.Value.DecodeType?.none as! Key.Value.ValueType + if Key.optional { + return Key.Value?.none as! Key.Value } throw StateBackendError.missingState } - func batchRead(_: [any StateKey]) async throws -> [(key: any StateKey, value: Codable & Sendable)] { - [] + public func batchRead(_ keys: [any StateKey]) async throws -> [(key: any StateKey, value: (Codable & Sendable)?)] { + let encodedKeys = keys.map { $0.encode().data } + let result = try await impl.batchRead(keys: encodedKeys) + return try zip(result, keys).map { data, key in + guard let rawValue = data.value else { + return (key: key, value: nil) + } + let value = try JamDecoder.decode(key.decodeType(), from: rawValue, withConfig: config) + return (key: key, value: value) + } + } + + public func readAll() async throws -> [Data32: Data] { + let all = try await impl.readAll(prefix: Data(), startKey: nil, limit: nil) + var result = [Data32: Data]() + for (key, value) in all { + try result[Data32(key).unwrap()] = value + } + return result + } + + public func newBackend(rootHash: Data32) -> StateBackend { + StateBackend(impl: impl, config: config, rootHash: rootHash) + } +} + +// MARK: - TrieNode + +extension StateBackend { + public func readTrieNode(_ key: Data) async throws -> Data? { + try await impl.read(key: key) + } + + public func batchUpdateTrieNodes(_ ops: [StateBackendOperation]) async throws { + try await impl.batchUpdate(ops) } - func readAll() async throws -> [Data32: Data] { - [:] + public func gc() async throws { + try await impl.gc() } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift index 630b2003..d38e5612 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift @@ -1,17 +1,19 @@ import Foundation import Utils -public enum StateBackendOperation: Sendable { +public enum StateTrieBackendOperation: Sendable { case write(key: Data, value: Data) case refIncrement(key: Data) case refDecrement(key: Data) } -public protocol StateBackendProtocol: Sendable { +// key: trie node hash (32 bytes) +// value: trie node data (64 bytes) +public protocol StateTrieBackend: Sendable { func read(key: Data) async throws -> Data? func readAll(prefix: Data, startKey: Data?, limit: UInt32?) async throws -> [(key: Data, value: Data)] func batchRead(keys: [Data]) async throws -> [(key: Data, value: Data?)] - func batchUpdate(_ ops: [StateBackendOperation]) async throws + func batchUpdate(_ ops: [StateTrieBackendOperation]) async throws // remove entries with zero ref count func gc() async throws diff --git a/Blockchain/Sources/Blockchain/State/StateKeys.swift b/Blockchain/Sources/Blockchain/State/StateKeys.swift index 422c474a..c1f1f630 100644 --- a/Blockchain/Sources/Blockchain/State/StateKeys.swift +++ b/Blockchain/Sources/Blockchain/State/StateKeys.swift @@ -2,34 +2,19 @@ import Foundation import Utils public protocol StateKey: Hashable, Sendable { - associatedtype Value: StateValueProtocol + associatedtype Value: Codable & Sendable func encode() -> Data32 + static var optional: Bool { get } } extension StateKey { public func decodeType() -> (Sendable & Codable).Type { - Value.DecodeType.self + Value.self } -} -public protocol StateValueProtocol { - associatedtype ValueType: Codable & Sendable - associatedtype DecodeType: Codable & Sendable - static var optional: Bool { get } -} - -public struct StateValue: StateValueProtocol { - public typealias ValueType = T - public typealias DecodeType = T public static var optional: Bool { false } } -public struct StateOptionalValue: StateValueProtocol { - public typealias ValueType = T? - public typealias DecodeType = T - public static var optional: Bool { true } -} - private func constructKey(_ idx: UInt8) -> Data32 { var data = Data(repeating: 0, count: 32) data[0] = idx @@ -68,16 +53,30 @@ private func constructKey(_ service: ServiceIndex, _ val: UInt32, _: Data) -> Da } public enum StateKeys { + public static let prefetchKeys: [any StateKey] = [ + CoreAuthorizationPoolKey(), + AuthorizationQueueKey(), + RecentHistoryKey(), + SafroleStateKey(), + JudgementsKey(), + EntropyPoolKey(), + ValidatorQueueKey(), + CurrentValidatorsKey(), + PreviousValidatorsKey(), + ReportsKey(), + TimeslotKey(), + PrivilegedServicesKey(), + ActivityStatisticsKey(), + ] + public struct CoreAuthorizationPoolKey: StateKey { - public typealias Value = StateValue< - ConfigFixedSizeArray< - ConfigLimitedSizeArray< - Data32, - ProtocolConfig.Int0, - ProtocolConfig.MaxAuthorizationsPoolItems - >, - ProtocolConfig.TotalNumberOfCores - > + public typealias Value = ConfigFixedSizeArray< + ConfigLimitedSizeArray< + Data32, + ProtocolConfig.Int0, + ProtocolConfig.MaxAuthorizationsPoolItems + >, + ProtocolConfig.TotalNumberOfCores > public init() {} @@ -88,14 +87,12 @@ public enum StateKeys { } public struct AuthorizationQueueKey: StateKey { - public typealias Value = StateValue< + public typealias Value = ConfigFixedSizeArray< ConfigFixedSizeArray< - ConfigFixedSizeArray< - Data32, - ProtocolConfig.MaxAuthorizationsQueueItems - >, - ProtocolConfig.TotalNumberOfCores - > + Data32, + ProtocolConfig.MaxAuthorizationsQueueItems + >, + ProtocolConfig.TotalNumberOfCores > public init() {} @@ -106,7 +103,7 @@ public enum StateKeys { } public struct RecentHistoryKey: StateKey { - public typealias Value = StateValue + public typealias Value = RecentHistory public init() {} @@ -116,7 +113,7 @@ public enum StateKeys { } public struct SafroleStateKey: StateKey { - public typealias Value = StateValue + public typealias Value = SafroleState public init() {} @@ -126,7 +123,7 @@ public enum StateKeys { } public struct JudgementsKey: StateKey { - public typealias Value = StateValue + public typealias Value = JudgementsState public init() {} @@ -136,7 +133,7 @@ public enum StateKeys { } public struct EntropyPoolKey: StateKey { - public typealias Value = StateValue + public typealias Value = EntropyPool public init() {} @@ -146,11 +143,9 @@ public enum StateKeys { } public struct ValidatorQueueKey: StateKey { - public typealias Value = StateValue< - ConfigFixedSizeArray< - ValidatorKey, - ProtocolConfig.TotalNumberOfValidators - > + public typealias Value = ConfigFixedSizeArray< + ValidatorKey, + ProtocolConfig.TotalNumberOfValidators > public init() {} @@ -161,11 +156,9 @@ public enum StateKeys { } public struct CurrentValidatorsKey: StateKey { - public typealias Value = StateValue< - ConfigFixedSizeArray< - ValidatorKey, - ProtocolConfig.TotalNumberOfValidators - > + public typealias Value = ConfigFixedSizeArray< + ValidatorKey, + ProtocolConfig.TotalNumberOfValidators > public init() {} @@ -176,11 +169,9 @@ public enum StateKeys { } public struct PreviousValidatorsKey: StateKey { - public typealias Value = StateValue< - ConfigFixedSizeArray< - ValidatorKey, - ProtocolConfig.TotalNumberOfValidators - > + public typealias Value = ConfigFixedSizeArray< + ValidatorKey, + ProtocolConfig.TotalNumberOfValidators > public init() {} @@ -191,11 +182,9 @@ public enum StateKeys { } public struct ReportsKey: StateKey { - public typealias Value = StateValue< - ConfigFixedSizeArray< - ReportItem?, - ProtocolConfig.TotalNumberOfCores - > + public typealias Value = ConfigFixedSizeArray< + ReportItem?, + ProtocolConfig.TotalNumberOfCores > public init() {} @@ -206,7 +195,7 @@ public enum StateKeys { } public struct TimeslotKey: StateKey { - public typealias Value = StateValue + public typealias Value = TimeslotIndex public init() {} @@ -216,7 +205,7 @@ public enum StateKeys { } public struct PrivilegedServicesKey: StateKey { - public typealias Value = StateValue + public typealias Value = PrivilegedServices public init() {} @@ -226,7 +215,7 @@ public enum StateKeys { } public struct ActivityStatisticsKey: StateKey { - public typealias Value = StateValue + public typealias Value = ValidatorActivityStatistics public init() {} @@ -236,7 +225,8 @@ public enum StateKeys { } public struct ServiceAccountKey: StateKey { - public typealias Value = StateOptionalValue + public typealias Value = ServiceAccountDetails + public static var optional: Bool { true } public var index: ServiceIndex @@ -250,7 +240,8 @@ public enum StateKeys { } public struct ServiceAccountStorageKey: StateKey { - public typealias Value = StateOptionalValue + public typealias Value = Data + public static var optional: Bool { true } public var index: ServiceIndex public var key: Data32 @@ -266,7 +257,8 @@ public enum StateKeys { } public struct ServiceAccountPreimagesKey: StateKey { - public typealias Value = StateOptionalValue + public typealias Value = Data + public static var optional: Bool { true } public var index: ServiceIndex public var hash: Data32 @@ -282,7 +274,8 @@ public enum StateKeys { } public struct ServiceAccountPreimageInfoKey: StateKey { - public typealias Value = StateOptionalValue> + public typealias Value = LimitedSizeArray + public static var optional: Bool { true } public var index: ServiceIndex public var hash: Data32 diff --git a/Blockchain/Sources/Blockchain/State/StateLayer.swift b/Blockchain/Sources/Blockchain/State/StateLayer.swift index e866eb1b..6742dd45 100644 --- a/Blockchain/Sources/Blockchain/State/StateLayer.swift +++ b/Blockchain/Sources/Blockchain/State/StateLayer.swift @@ -1,238 +1,242 @@ import Foundation import Utils +private enum StateLayerValue: Sendable { + case value(Codable & Sendable) + case deleted + + init(_ value: (Codable & Sendable)?) { + if let value { + self = .value(value) + } else { + self = .deleted + } + } + + func value() -> T? { + if case let .value(value) = self { + return value as? T + } + return nil + } +} + // @unchecked because AnyHashable is not Sendable public struct StateLayer: @unchecked Sendable { - private var changes: [AnyHashable: Codable & Sendable] = [:] + private var changes: [AnyHashable: StateLayerValue] = [:] public init(backend: StateBackend) async throws { - let keys: [any StateKey] = [ - StateKeys.CoreAuthorizationPoolKey(), - StateKeys.AuthorizationQueueKey(), - StateKeys.RecentHistoryKey(), - StateKeys.SafroleStateKey(), - StateKeys.JudgementsKey(), - StateKeys.EntropyPoolKey(), - StateKeys.ValidatorQueueKey(), - StateKeys.CurrentValidatorsKey(), - StateKeys.PreviousValidatorsKey(), - StateKeys.ReportsKey(), - StateKeys.TimeslotKey(), - StateKeys.PrivilegedServicesKey(), - StateKeys.ActivityStatisticsKey(), - ] - - let results = try await backend.batchRead(keys) + let results = try await backend.batchRead(StateKeys.prefetchKeys) for (key, value) in results { - changes[AnyHashable(key)] = value + changes[AnyHashable(key)] = try .init(value.unwrap()) } } public init(changes: [(key: any StateKey, value: Codable & Sendable)]) { for (key, value) in changes { - self.changes[AnyHashable(key)] = value + self.changes[AnyHashable(key)] = .value(value) } } // α: The core αuthorizations pool. - public var coreAuthorizationPool: StateKeys.CoreAuthorizationPoolKey.Value.ValueType { + public var coreAuthorizationPool: StateKeys.CoreAuthorizationPoolKey.Value { get { - changes[StateKeys.CoreAuthorizationPoolKey()] as! StateKeys.CoreAuthorizationPoolKey.Value.ValueType + changes[StateKeys.CoreAuthorizationPoolKey()]!.value()! } set { - changes[StateKeys.CoreAuthorizationPoolKey()] = newValue + changes[StateKeys.CoreAuthorizationPoolKey()] = .init(newValue) } } // φ: The authorization queue. - public var authorizationQueue: StateKeys.AuthorizationQueueKey.Value.ValueType { + public var authorizationQueue: StateKeys.AuthorizationQueueKey.Value { get { - changes[StateKeys.AuthorizationQueueKey()] as! StateKeys.AuthorizationQueueKey.Value.ValueType + changes[StateKeys.AuthorizationQueueKey()]!.value()! } set { - changes[StateKeys.AuthorizationQueueKey()] = newValue + changes[StateKeys.AuthorizationQueueKey()] = .init(newValue) } } // β: Information on the most recent βlocks. - public var recentHistory: StateKeys.RecentHistoryKey.Value.ValueType { + public var recentHistory: StateKeys.RecentHistoryKey.Value { get { - changes[StateKeys.RecentHistoryKey()] as! StateKeys.RecentHistoryKey.Value.ValueType + changes[StateKeys.RecentHistoryKey()]!.value()! } set { - changes[StateKeys.RecentHistoryKey()] = newValue + changes[StateKeys.RecentHistoryKey()] = .init(newValue) } } // γ: State concerning Safrole. - public var safroleState: StateKeys.SafroleStateKey.Value.ValueType { + public var safroleState: StateKeys.SafroleStateKey.Value { get { - changes[StateKeys.SafroleStateKey()] as! StateKeys.SafroleStateKey.Value.ValueType + changes[StateKeys.SafroleStateKey()]!.value()! } set { - changes[StateKeys.SafroleStateKey()] = newValue + changes[StateKeys.SafroleStateKey()] = .init(newValue) } } // ψ: past judgements - public var judgements: StateKeys.JudgementsKey.Value.ValueType { + public var judgements: StateKeys.JudgementsKey.Value { get { - changes[StateKeys.JudgementsKey()] as! StateKeys.JudgementsKey.Value.ValueType + changes[StateKeys.JudgementsKey()]!.value()! } set { - changes[StateKeys.JudgementsKey()] = newValue + changes[StateKeys.JudgementsKey()] = .init(newValue) } } // η: The eηtropy accumulator and epochal raηdomness. - public var entropyPool: StateKeys.EntropyPoolKey.Value.ValueType { + public var entropyPool: StateKeys.EntropyPoolKey.Value { get { - changes[StateKeys.EntropyPoolKey()] as! StateKeys.EntropyPoolKey.Value.ValueType + changes[StateKeys.EntropyPoolKey()]!.value()! } set { - changes[StateKeys.EntropyPoolKey()] = newValue + changes[StateKeys.EntropyPoolKey()] = .init(newValue) } } // ι: The validator keys and metadata to be drawn from next. - public var validatorQueue: StateKeys.ValidatorQueueKey.Value.ValueType { + public var validatorQueue: StateKeys.ValidatorQueueKey.Value { get { - changes[StateKeys.ValidatorQueueKey()] as! StateKeys.ValidatorQueueKey.Value.ValueType + changes[StateKeys.ValidatorQueueKey()]!.value()! } set { - changes[StateKeys.ValidatorQueueKey()] = newValue + changes[StateKeys.ValidatorQueueKey()] = .init(newValue) } } // κ: The validator κeys and metadata currently active. - public var currentValidators: StateKeys.CurrentValidatorsKey.Value.ValueType { + public var currentValidators: StateKeys.CurrentValidatorsKey.Value { get { - changes[StateKeys.CurrentValidatorsKey()] as! StateKeys.CurrentValidatorsKey.Value.ValueType + changes[StateKeys.CurrentValidatorsKey()]!.value()! } set { - changes[StateKeys.CurrentValidatorsKey()] = newValue + changes[StateKeys.CurrentValidatorsKey()] = .init(newValue) } } // λ: The validator keys and metadata which were active in the prior epoch. - public var previousValidators: StateKeys.PreviousValidatorsKey.Value.ValueType { + public var previousValidators: StateKeys.PreviousValidatorsKey.Value { get { - changes[StateKeys.PreviousValidatorsKey()] as! StateKeys.PreviousValidatorsKey.Value.ValueType + changes[StateKeys.PreviousValidatorsKey()]!.value()! } set { - changes[StateKeys.PreviousValidatorsKey()] = newValue + changes[StateKeys.PreviousValidatorsKey()] = .init(newValue) } } // ρ: The ρending reports, per core, which are being made available prior to accumulation. - public var reports: StateKeys.ReportsKey.Value.ValueType { + public var reports: StateKeys.ReportsKey.Value { get { - changes[StateKeys.ReportsKey()] as! StateKeys.ReportsKey.Value.ValueType + changes[StateKeys.ReportsKey()]!.value()! } set { - changes[StateKeys.ReportsKey()] = newValue + changes[StateKeys.ReportsKey()] = .init(newValue) } } // τ: The most recent block’s τimeslot. - public var timeslot: StateKeys.TimeslotKey.Value.ValueType { + public var timeslot: StateKeys.TimeslotKey.Value { get { - changes[StateKeys.TimeslotKey()] as! StateKeys.TimeslotKey.Value.ValueType + changes[StateKeys.TimeslotKey()]!.value()! } set { - changes[StateKeys.TimeslotKey()] = newValue + changes[StateKeys.TimeslotKey()] = .init(newValue) } } // χ: The privileged service indices. - public var privilegedServices: StateKeys.PrivilegedServicesKey.Value.ValueType { + public var privilegedServices: StateKeys.PrivilegedServicesKey.Value { get { - changes[StateKeys.PrivilegedServicesKey()] as! StateKeys.PrivilegedServicesKey.Value.ValueType + changes[StateKeys.PrivilegedServicesKey()]!.value()! } set { - changes[StateKeys.PrivilegedServicesKey()] = newValue + changes[StateKeys.PrivilegedServicesKey()] = .init(newValue) } } // π: The activity statistics for the validators. - public var activityStatistics: StateKeys.ActivityStatisticsKey.Value.ValueType { + public var activityStatistics: StateKeys.ActivityStatisticsKey.Value { get { - changes[StateKeys.ActivityStatisticsKey()] as! StateKeys.ActivityStatisticsKey.Value.ValueType + changes[StateKeys.ActivityStatisticsKey()]!.value()! } set { - changes[StateKeys.ActivityStatisticsKey()] = newValue + changes[StateKeys.ActivityStatisticsKey()] = .init(newValue) } } // δ: The (prior) state of the service accounts. - public subscript(serviceAccount index: ServiceIndex) -> StateKeys.ServiceAccountKey.Value.ValueType? { + public subscript(serviceAccount index: ServiceIndex) -> StateKeys.ServiceAccountKey.Value? { get { - changes[StateKeys.ServiceAccountKey(index: index)] as? StateKeys.ServiceAccountKey.Value.ValueType + changes[StateKeys.ServiceAccountKey(index: index)]!.value()! } set { - changes[StateKeys.ServiceAccountKey(index: index)] = newValue + changes[StateKeys.ServiceAccountKey(index: index)] = .init(newValue) } } // s - public subscript(serviceAccount index: ServiceIndex, storageKey key: Data32) -> StateKeys.ServiceAccountStorageKey.Value.ValueType? { + public subscript(serviceAccount index: ServiceIndex, storageKey key: Data32) -> StateKeys.ServiceAccountStorageKey.Value? { get { - changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)] as? StateKeys.ServiceAccountStorageKey.Value.ValueType + changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)]!.value()! } set { - changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)] = newValue + changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)] = .init(newValue) } } // p public subscript( serviceAccount index: ServiceIndex, preimageHash hash: Data32 - ) -> StateKeys.ServiceAccountPreimagesKey.Value.ValueType? { + ) -> StateKeys.ServiceAccountPreimagesKey.Value? { get { - changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)] as? StateKeys.ServiceAccountPreimagesKey.Value.ValueType + changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)]!.value()! } set { - changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)] = newValue + changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)] = .init(newValue) } } // l public subscript( serviceAccount index: ServiceIndex, preimageHash hash: Data32, length length: UInt32 - ) -> StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType? { + ) -> StateKeys.ServiceAccountPreimageInfoKey.Value? { get { changes[ StateKeys.ServiceAccountPreimageInfoKey(index: index, hash: hash, length: length) - ] as? StateKeys.ServiceAccountPreimageInfoKey.Value.ValueType + ]!.value()! } set { - changes[StateKeys.ServiceAccountPreimageInfoKey(index: index, hash: hash, length: length)] = newValue + changes[StateKeys.ServiceAccountPreimageInfoKey(index: index, hash: hash, length: length)] = .init(newValue) } } } extension StateLayer { - public func toKV() -> some Sequence<(key: any StateKey, value: Codable & Sendable)> { - changes.map { (key: $0.key.base as! any StateKey, value: $0.value) } + public func toKV() -> some Sequence<(key: any StateKey, value: (Codable & Sendable)?)> { + changes.map { (key: $0.key.base as! any StateKey, value: $0.value.value()) } } } extension StateLayer { - public func read(_ key: Key) -> Key.Value.ValueType? { - changes[key] as? Key.Value.ValueType + public func read(_ key: Key) -> Key.Value? { + changes[key] as? Key.Value } - public mutating func write(_ key: Key, value: Key.Value.ValueType) { - changes[key] = value + public mutating func write(_ key: Key, value: Key.Value?) { + changes[key] = .init(value) } public subscript(key: any StateKey) -> (Codable & Sendable)? { get { - changes[AnyHashable(key)] + changes[AnyHashable(key)]?.value() } set { - changes[AnyHashable(key)] = newValue + changes[AnyHashable(key)] = .init(newValue) } } } diff --git a/Blockchain/Sources/Blockchain/State/StateTrie.swift b/Blockchain/Sources/Blockchain/State/StateTrie.swift index 6165accf..23845955 100644 --- a/Blockchain/Sources/Blockchain/State/StateTrie.swift +++ b/Blockchain/Sources/Blockchain/State/StateTrie.swift @@ -69,12 +69,12 @@ public enum StateTrieError: Error { } public class StateTrie { - private let backend: StateBackendProtocol + private let backend: StateBackend public private(set) var rootHash: Data32 private var nodes: [Data32: TrieNode] = [:] private var deleted: Set = [] - public init(rootHash: Data32, backend: StateBackendProtocol) { + public init(rootHash: Data32, backend: StateBackend) { self.rootHash = rootHash self.backend = backend } @@ -86,7 +86,7 @@ public class StateTrie { if let node = nodes[hash] { return node } - guard let data = try await backend.read(key: hash.data) else { + guard let data = try await backend.readTrieNode(key: hash.data) else { return nil } @@ -148,7 +148,7 @@ public class StateTrie { } } - try await backend.batchUpdate(ops) + try await backend.batchUpdateTrieNodes(ops) } private func insert( diff --git a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift index 137ab636..651710dd 100644 --- a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift +++ b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift @@ -118,7 +118,7 @@ public final class BlockAuthor: ServiceBase2, @unchecked Sendable { extrinsics: extrinsic.tickets ) - let unsignedHeader = Header.Unsigned( + let unsignedHeader = try await Header.Unsigned( parentHash: parentHash, priorStateRoot: state.stateRoot, extrinsicsHash: extrinsic.hash(), From 5cdfc1b0e21ca8aaacfdd95542cc133337d72e74 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Thu, 7 Nov 2024 17:10:36 +1300 Subject: [PATCH 03/10] many fixes --- .../Blockchain/State/InMemoryBackend.swift | 59 +++++++++----- .../Sources/Blockchain/State/State.swift | 71 ++--------------- .../Blockchain/State/StateBackend.swift | 69 ++++++++--------- .../State/StateBackendProtocol.swift | 27 +++++-- .../Sources/Blockchain/State/StateTrie.swift | 76 ++++++++++++++++--- .../Blockchain/Validator/BlockAuthor.swift | 2 +- .../Utils/SortedContainer/SortedArray.swift | 7 ++ 7 files changed, 171 insertions(+), 140 deletions(-) diff --git a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift index bbd4b67c..f8bb2803 100644 --- a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift +++ b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift @@ -12,17 +12,14 @@ private struct KVPair: Comparable, Sendable { } public actor InMemoryBackend: StateBackendProtocol { - private var store: SortedArray - private var refCounts: [Data: Int] + // we really should be using Heap or some other Tree based structure here + // but let's keep it simple for now + private var store: SortedArray = .init([]) + private var rawValues: [Data32: Data] = [:] + private var refCounts: [Data: Int] = [:] + private var rawValueRefCounts: [Data32: Int] = [:] - public init(store: [Data32: Data] = [:]) { - self.store = .init(store.map { KVPair(key: $0.key.data, value: $0.value) }) - refCounts = [:] - - for key in store.keys { - refCounts[key.data] = 1 - } - } + public init() {} public func read(key: Data) async throws -> Data? { let idx = store.insertIndex(KVPair(key: key, value: Data())) @@ -51,23 +48,47 @@ public actor InMemoryBackend: StateBackendProtocol { return resp } - public func batchRead(keys: [Data]) async throws -> [(key: Data, value: Data?)] { - var resp = [(key: Data, value: Data?)]() - for key in keys { - let value = try await read(key: key) - resp.append((key, value)) + public func batchUpdate(_ updates: [StateBackendOperation]) async throws { + for update in updates { + switch update { + case let .write(key, value): + let idx = store.insertIndex(KVPair(key: key, value: value)) + let item = store.array[safe: idx] + if let item, item.key == key { // found + // value is not used for ordering so this is safe + store.unsafeArrayAccess[idx].value = value + } else { // not found + store.insert(KVPair(key: key, value: value)) + } + case let .writeRawValue(key, value): + rawValues[key] = value + rawValueRefCounts[key, default: 0] += 1 + case .refIncrement: + break + case .refDecrement: + break + } } - return resp } - public func batchUpdate(_: [StateBackendOperation]) async throws {} + public func readValue(hash: Data32) async throws -> Data? { + rawValues[hash] + } - public func gc() async throws { + public func gc(callback: @Sendable (Data) -> Data32?) async throws { // check ref counts and remove keys with 0 ref count for (key, count) in refCounts where count == 0 { let idx = store.insertIndex(KVPair(key: key, value: Data())) - if store.array[safe: idx]?.key == key { + let item = store.array[safe: idx] + if let item, item.key == key { store.remove(at: idx) + if let rawValueKey = callback(item.value) { + rawValueRefCounts[rawValueKey, default: 0] -= 1 + if rawValueRefCounts[rawValueKey] == 0 { + rawValues.removeValue(forKey: rawValueKey) + rawValueRefCounts.removeValue(forKey: rawValueKey) + } + } } } } diff --git a/Blockchain/Sources/Blockchain/State/State.swift b/Blockchain/Sources/Blockchain/State/State.swift index d5ab6a62..ad74c68a 100644 --- a/Blockchain/Sources/Blockchain/State/State.swift +++ b/Blockchain/Sources/Blockchain/State/State.swift @@ -197,12 +197,8 @@ public struct State: Sendable { } } - public func save() async throws -> State { - let changes = layer.toKV() - let trie = StateTrie(rootHash: backend.rootHash, backend: backend) - try await trie.update(updates: changes) - let newRoot = trie.rootHash - return State(backend: backend.newBackend(rootHash: newRoot), layer: layer) + public func save() async throws { + try await backend.write(layer.toKV()) } } @@ -211,63 +207,6 @@ extension State { recentHistory.items.last.map(\.headerHash)! } - private class KVSequence: Sequence { - typealias Element = (key: Data32, value: Data) - - let seq: any Sequence<(key: Data32, value: Data)> - let layer: [Data32: Data?] - - init(state: State) async throws { - seq = try await state.backend.readAll() - var layer = [Data32: Data?]() - for (key, value) in state.layer.toKV() { - layer[key.encode()] = try value.map { try JamEncoder.encode($0) } - } - self.layer = layer - } - - func makeIterator() -> KVSequence.Iterator { - KVSequence.Iterator(iter: seq.makeIterator(), layer: layer) - } - - struct Iterator: IteratorProtocol { - typealias Element = (key: Data32, value: Data) - - var iter: any IteratorProtocol - var layerIterator: [Data32: Data?].Iterator? - let layer: [Data32: Data?] - - init(iter: any IteratorProtocol, layer: [Data32: Data?]) { - self.iter = iter - self.layer = layer - } - - mutating func next() -> KVSequence.Iterator.Element? { - if layerIterator != nil { - if let (key, value) = layerIterator?.next() { - if let value { - return (key, value) - } - return next() // skip this one - } - return nil - } - if let (key, value) = iter.next() { - if layer.keys.contains(key) { - return next() // skip this one - } - return (key, value) - } - layerIterator = layer.makeIterator() - return next() - } - } - } - - public func toKV() async throws -> some Sequence<(key: Data32, value: Data)> { - try await KVSequence(state: self) - } - public func asRef() -> StateRef { StateRef(self) } @@ -335,7 +274,7 @@ extension State: Dummy { } let rootHash = try! stateMerklize(kv: store) - let backend = StateBackend(InMemoryBackend(store: store), config: config, rootHash: rootHash) + let backend = StateBackend(InMemoryBackend(), config: config, rootHash: rootHash) let layer = StateLayer(changes: kv) @@ -489,4 +428,8 @@ public class StateRef: Ref, @unchecked Sendable { public static func dummy(config: ProtocolConfigRef, block: BlockRef?) -> StateRef { StateRef(State.dummy(config: config, block: block)) } + + public var stateRoot: Data32 { + fatalError("not implemented") + } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackend.swift b/Blockchain/Sources/Blockchain/State/StateBackend.swift index 55c5a336..6b324686 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackend.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackend.swift @@ -10,66 +10,59 @@ public enum StateBackendError: Error { public final class StateBackend: Sendable { private let impl: StateBackendProtocol private let config: ProtocolConfigRef - public let rootHash: Data32 + private let trie: StateTrie public init(_ impl: StateBackendProtocol, config: ProtocolConfigRef, rootHash: Data32) { self.impl = impl self.config = config - self.rootHash = rootHash + trie = StateTrie(rootHash: rootHash, backend: impl) } - public func read(_ key: Key) async throws -> Key.Value { - let encodedKey = key.encode().data - if let ret = try await impl.read(key: encodedKey) { + public var rootHash: Data32 { + get async { + await trie.rootHash + } + } + + public func read(_ key: Key) async throws -> Key.Value? { + let encodedKey = key.encode() + if let ret = try await trie.read(key: encodedKey) { guard let ret = try JamDecoder.decode(key.decodeType(), from: ret, withConfig: config) as? Key.Value else { throw StateBackendError.invalidData } return ret } if Key.optional { - return Key.Value?.none as! Key.Value + return nil } throw StateBackendError.missingState } public func batchRead(_ keys: [any StateKey]) async throws -> [(key: any StateKey, value: (Codable & Sendable)?)] { - let encodedKeys = keys.map { $0.encode().data } - let result = try await impl.batchRead(keys: encodedKeys) - return try zip(result, keys).map { data, key in - guard let rawValue = data.value else { - return (key: key, value: nil) - } - let value = try JamDecoder.decode(key.decodeType(), from: rawValue, withConfig: config) - return (key: key, value: value) + var ret = [(key: any StateKey, value: (Codable & Sendable)?)]() + ret.reserveCapacity(keys.count) + for key in keys { + try await ret.append((key, read(key))) } + return ret } - public func readAll() async throws -> [Data32: Data] { - let all = try await impl.readAll(prefix: Data(), startKey: nil, limit: nil) - var result = [Data32: Data]() - for (key, value) in all { - try result[Data32(key).unwrap()] = value - } - return result - } - - public func newBackend(rootHash: Data32) -> StateBackend { - StateBackend(impl: impl, config: config, rootHash: rootHash) - } -} - -// MARK: - TrieNode - -extension StateBackend { - public func readTrieNode(_ key: Data) async throws -> Data? { - try await impl.read(key: key) - } - - public func batchUpdateTrieNodes(_ ops: [StateBackendOperation]) async throws { - try await impl.batchUpdate(ops) + public func write(_ values: any Sequence<(key: any StateKey, value: (Codable & Sendable)?)>) async throws { + try await trie.update(values.map { try (key: $0.key.encode(), value: $0.value.map { try JamEncoder.encode($0) }) }) + try await trie.save() } public func gc() async throws { - try await impl.gc() + try await impl.gc { data in + guard data.count == 64 else { + // unexpected data size + return nil + } + let isRegularLeaf = data[0] & 0b1100_0000 == 0b1100_0000 + if isRegularLeaf { + return Data32(data.suffix(from: 32))! + } + return nil + } } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift index d38e5612..0ffe1a3d 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift @@ -1,20 +1,31 @@ import Foundation import Utils -public enum StateTrieBackendOperation: Sendable { +public enum StateBackendOperation: Sendable { case write(key: Data, value: Data) + case writeRawValue(key: Data32, value: Data) case refIncrement(key: Data) case refDecrement(key: Data) } -// key: trie node hash (32 bytes) -// value: trie node data (64 bytes) -public protocol StateTrieBackend: Sendable { +/// key: trie node hash (32 bytes) +/// value: trie node data (64 bytes) +/// ref counting requirements: +/// - write do not increment ref count, only explicit ref increment do +/// - lazy prune is used. e.g. when ref count is reduced to zero, the value will only be removed +/// when gc is performed +/// - raw value have its own ref counting +/// - writeRawValue increment ref count, and write if necessary +/// - raw value ref count is only decremented when connected trie node is removed during gc +public protocol StateBackendProtocol: Sendable { func read(key: Data) async throws -> Data? func readAll(prefix: Data, startKey: Data?, limit: UInt32?) async throws -> [(key: Data, value: Data)] - func batchRead(keys: [Data]) async throws -> [(key: Data, value: Data?)] - func batchUpdate(_ ops: [StateTrieBackendOperation]) async throws + func batchUpdate(_ ops: [StateBackendOperation]) async throws - // remove entries with zero ref count - func gc() async throws + // hash is the blake2b256 hash of the value + func readValue(hash: Data32) async throws -> Data? + + /// remove entries with zero ref count + /// callback returns a dependent raw value key if the data is regular leaf node + func gc(callback: @Sendable (Data) -> Data32?) async throws } diff --git a/Blockchain/Sources/Blockchain/State/StateTrie.swift b/Blockchain/Sources/Blockchain/State/StateTrie.swift index 23845955..89429e01 100644 --- a/Blockchain/Sources/Blockchain/State/StateTrie.swift +++ b/Blockchain/Sources/Blockchain/State/StateTrie.swift @@ -13,12 +13,14 @@ private struct TrieNode { var right: Data32 var type: TrieNodeType var isNew: Bool + var rawValue: Data? init(hash: Data32, data: Data64, isNew: Bool = false) { self.hash = hash left = Data32(data.data.prefix(32))! right = Data32(data.data.suffix(32))! self.isNew = isNew + rawValue = nil switch data.data[0] & 0b1100_0000 { case 0b1000_0000: type = .embeddedLeaf @@ -29,12 +31,13 @@ private struct TrieNode { } } - private init(left: Data32, right: Data32, type: TrieNodeType, isNew: Bool = false) { + private init(left: Data32, right: Data32, type: TrieNodeType, isNew: Bool, rawValue: Data?) { hash = Blake2b256.hash(left.data, right.data) self.left = left self.right = right self.type = type self.isNew = isNew + self.rawValue = rawValue } var encodedData: Data64 { @@ -45,21 +48,40 @@ private struct TrieNode { type == .branch } + var isLeaf: Bool { + !isBranch + } + + func isLeaf(key: Data32) -> Bool { + isLeaf && left.data[relative: 1 ..< 32] == key.data.prefix(31) + } + + var value: Data? { + if let rawValue { + return rawValue + } + guard type == .embeddedLeaf else { + return nil + } + let len = left.data[0] & 0b0011_1111 + return right.data[relative: 0 ..< Int(len)] + } + static func leaf(key: Data32, value: Data) -> TrieNode { var newKey = Data(capacity: 32) if value.count <= 32 { newKey.append(0b1000_0000 | UInt8(value.count)) newKey += key.data.prefix(31) let newValue = value + Data(repeating: 0, count: 32 - value.count) - return .init(left: Data32(newKey)!, right: Data32(newValue)!, type: .embeddedLeaf, isNew: true) + return .init(left: Data32(newKey)!, right: Data32(newValue)!, type: .embeddedLeaf, isNew: true, rawValue: value) } newKey.append(0b1100_0000) newKey += key.data.prefix(31) - return .init(left: Data32(newKey)!, right: value.blake2b256hash(), type: .regularLeaf, isNew: true) + return .init(left: Data32(newKey)!, right: value.blake2b256hash(), type: .regularLeaf, isNew: true, rawValue: value) } static func branch(left: Data32, right: Data32) -> TrieNode { - .init(left: left, right: right, type: .branch, isNew: true) + .init(left: left, right: right, type: .branch, isNew: true, rawValue: nil) } } @@ -68,25 +90,56 @@ public enum StateTrieError: Error { case invalidParent } -public class StateTrie { - private let backend: StateBackend +public actor StateTrie: Sendable { + private let backend: StateBackendProtocol public private(set) var rootHash: Data32 private var nodes: [Data32: TrieNode] = [:] private var deleted: Set = [] - public init(rootHash: Data32, backend: StateBackend) { + public init(rootHash: Data32, backend: StateBackendProtocol) { self.rootHash = rootHash self.backend = backend } + public func read(key: Data32) async throws -> Data? { + let node = try await find(hash: rootHash, key: key, depth: 0) + guard let node else { + return nil + } + if let value = node.value { + return value + } + return try await backend.readValue(hash: node.right) + } + + private func find(hash: Data32, key: Data32, depth: UInt8) async throws -> TrieNode? { + guard let node = try await get(hash: hash) else { + return nil + } + if node.isBranch { + let bitValue = bitAt(key.data, position: depth) + if bitValue { + return try await find(hash: node.right, key: key, depth: depth + 1) + } else { + return try await find(hash: node.left, key: key, depth: depth + 1) + } + } else if node.isLeaf(key: key) { + return node + } + return nil + } + private func get(hash: Data32) async throws -> TrieNode? { + if hash == Data32() { + return nil + } if deleted.contains(hash) { return nil } if let node = nodes[hash] { return node } - guard let data = try await backend.readTrieNode(key: hash.data) else { + guard let data = try await backend.read(key: hash.data) else { return nil } @@ -129,6 +182,9 @@ public class StateTrie { for node in nodes.values where node.isNew { ops.append(.write(key: node.hash.data, value: node.encodedData.data)) + if node.type == .regularLeaf { + try ops.append(.writeRawValue(key: node.right, value: node.rawValue.unwrap())) + } if node.isBranch { refChanges[node.left.data] = (refChanges[node.left.data] ?? 0) + 1 refChanges[node.right.data] = (refChanges[node.right.data] ?? 0) + 1 @@ -148,7 +204,7 @@ public class StateTrie { } } - try await backend.batchUpdateTrieNodes(ops) + try await backend.batchUpdate(ops) } private func insert( @@ -176,7 +232,7 @@ public class StateTrie { } private func insertLeafNode(existing: TrieNode, newKey: Data32, newValue: Data, depth: UInt8) async throws -> Data32 { - if existing.left == newKey { + if existing.isLeaf(key: newKey) { // update existing leaf let newLeaf = TrieNode.leaf(key: newKey, value: newValue) nodes[newLeaf.hash] = newLeaf diff --git a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift index 651710dd..137ab636 100644 --- a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift +++ b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift @@ -118,7 +118,7 @@ public final class BlockAuthor: ServiceBase2, @unchecked Sendable { extrinsics: extrinsic.tickets ) - let unsignedHeader = try await Header.Unsigned( + let unsignedHeader = Header.Unsigned( parentHash: parentHash, priorStateRoot: state.stateRoot, extrinsicsHash: extrinsic.hash(), diff --git a/Utils/Sources/Utils/SortedContainer/SortedArray.swift b/Utils/Sources/Utils/SortedContainer/SortedArray.swift index 3b86e327..a33aa42f 100644 --- a/Utils/Sources/Utils/SortedContainer/SortedArray.swift +++ b/Utils/Sources/Utils/SortedContainer/SortedArray.swift @@ -44,6 +44,13 @@ public struct SortedArray: SortedContainer { public mutating func remove(where predicate: (T) throws -> Bool) rethrows { try array.removeAll(where: predicate) } + + // mutate access to underlying array directly + // this is unsafe and should be used with care + public var unsafeArrayAccess: [T] { + _read { yield array } + _modify { yield &array } + } } extension SortedArray: Encodable where T: Encodable { From f6c08eb176eb7932d0dda219fdf5309250f8c554 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Thu, 7 Nov 2024 18:31:21 +1300 Subject: [PATCH 04/10] more fix --- Blockchain/Sources/Blockchain/Blockchain.swift | 4 +++- .../Blockchain/RuntimeProtocols/Runtime.swift | 6 ++++-- Blockchain/Sources/Blockchain/State/State.swift | 16 +++++++++++----- .../Sources/Blockchain/State/StateBackend.swift | 5 +++++ .../Blockchain/Validator/BlockAuthor.swift | 4 +++- .../Tests/BlockchainTests/BlockAuthorTests.swift | 15 ++++++++++++--- .../BlockchainDataProviderTests.swift | 4 ++-- Node/Sources/Node/Genesis.swift | 6 ++++-- Node/Sources/Node/Node.swift | 4 ++-- 9 files changed, 46 insertions(+), 18 deletions(-) diff --git a/Blockchain/Sources/Blockchain/Blockchain.swift b/Blockchain/Sources/Blockchain/Blockchain.swift index 49575762..8b578e33 100644 --- a/Blockchain/Sources/Blockchain/Blockchain.swift +++ b/Blockchain/Sources/Blockchain/Blockchain.swift @@ -39,9 +39,11 @@ public final class Blockchain: ServiceBase, @unchecked Sendable { let runtime = Runtime(config: config) let parent = try await dataProvider.getState(hash: block.header.parentHash) + let stateRoot = await parent.value.stateRoot let timeslot = timeProvider.getTime().timeToTimeslot(config: config) // TODO: figure out what is the best way to deal with block received a bit too early - let state = try await runtime.apply(block: block, state: parent, context: .init(timeslot: timeslot + 1)) + let context = Runtime.ApplyContext(timeslot: timeslot + 1, stateRoot: stateRoot) + let state = try await runtime.apply(block: block, state: parent, context: context) try await dataProvider.blockImported(block: block, state: state) diff --git a/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift b/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift index 5f506e4f..102fd5cf 100644 --- a/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift +++ b/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift @@ -35,9 +35,11 @@ public final class Runtime { public struct ApplyContext { public let timeslot: TimeslotIndex + public let stateRoot: Data32 - public init(timeslot: TimeslotIndex) { + public init(timeslot: TimeslotIndex, stateRoot _: Data32) { self.timeslot = timeslot + stateRoot = Data32() } } @@ -54,7 +56,7 @@ public final class Runtime { throw Error.invalidParentHash } - guard block.header.priorStateRoot == state.stateRoot else { + guard block.header.priorStateRoot == context.stateRoot else { throw Error.invalidHeaderStateRoot } diff --git a/Blockchain/Sources/Blockchain/State/State.swift b/Blockchain/Sources/Blockchain/State/State.swift index ad74c68a..3e9f6299 100644 --- a/Blockchain/Sources/Blockchain/State/State.swift +++ b/Blockchain/Sources/Blockchain/State/State.swift @@ -197,8 +197,18 @@ public struct State: Sendable { } } - public func save() async throws { + // TODO: we don't really want to write to the underlying backend here + // instead, it should be writting to a in memory layer + // and when actually saving the state, save the in memory layer to the presistent store + public func save() async throws -> Data32 { try await backend.write(layer.toKV()) + return await backend.rootHash + } + + public var stateRoot: Data32 { + get async { + await backend.rootHash + } } } @@ -428,8 +438,4 @@ public class StateRef: Ref, @unchecked Sendable { public static func dummy(config: ProtocolConfigRef, block: BlockRef?) -> StateRef { StateRef(State.dummy(config: config, block: block)) } - - public var stateRoot: Data32 { - fatalError("not implemented") - } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackend.swift b/Blockchain/Sources/Blockchain/State/StateBackend.swift index 6b324686..36ecf36f 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackend.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackend.swift @@ -52,6 +52,11 @@ public final class StateBackend: Sendable { try await trie.save() } + public func writeRaw(_ values: [(key: Data32, value: Data?)]) async throws { + try await trie.update(values) + try await trie.save() + } + public func gc() async throws { try await impl.gc { data in guard data.count == 64 else { diff --git a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift index 137ab636..9cea5de0 100644 --- a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift +++ b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift @@ -118,9 +118,11 @@ public final class BlockAuthor: ServiceBase2, @unchecked Sendable { extrinsics: extrinsic.tickets ) + let stateRoot = try await state.value.save() + let unsignedHeader = Header.Unsigned( parentHash: parentHash, - priorStateRoot: state.stateRoot, + priorStateRoot: stateRoot, extrinsicsHash: extrinsic.hash(), timeslot: timeslot, epoch: safroleResult.epochMark, diff --git a/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift b/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift index 16b809c4..6890933d 100644 --- a/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift +++ b/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift @@ -42,7 +42,10 @@ struct BlockAuthorTests { let block = try await blockAuthor.createNewBlock(timeslot: timeslot, claim: .right(pubkey)) // Verify block - try _ = await runtime.apply(block: block, state: genesisState, context: .init(timeslot: timeslot + 1)) + try _ = await runtime.apply(block: block, state: genesisState, context: .init( + timeslot: timeslot + 1, + stateRoot: genesisState.value.stateRoot + )) } @Test @@ -95,7 +98,10 @@ struct BlockAuthorTests { let block = try await blockAuthor.createNewBlock(timeslot: timeslot, claim: .left((ticket, devKey.bandersnatch))) // Verify block - try _ = await runtime.apply(block: block, state: newStateRef, context: .init(timeslot: timeslot + 1)) + try _ = await runtime.apply(block: block, state: newStateRef, context: .init( + timeslot: timeslot + 1, + stateRoot: newStateRef.value.stateRoot + )) } @Test @@ -122,7 +128,10 @@ struct BlockAuthorTests { let timeslot = timeProvider.getTime().timeToTimeslot(config: config) // Verify block - try _ = await runtime.apply(block: block.block, state: genesisState, context: .init(timeslot: timeslot + 1)) + try _ = await runtime.apply(block: block.block, state: genesisState, context: .init( + timeslot: timeslot + 1, + stateRoot: genesisState.value.stateRoot + )) } // TODO: test including extrinsic tickets from extrinsic pool diff --git a/Blockchain/Tests/BlockchainTests/BlockchainDataProviderTests.swift b/Blockchain/Tests/BlockchainTests/BlockchainDataProviderTests.swift index 27286e99..8b4320e0 100644 --- a/Blockchain/Tests/BlockchainTests/BlockchainDataProviderTests.swift +++ b/Blockchain/Tests/BlockchainTests/BlockchainDataProviderTests.swift @@ -78,11 +78,11 @@ struct BlockchainDataProviderTests { // Verify state exists #expect(try await provider.hasState(hash: block.hash)) - #expect(try await provider.getState(hash: block.hash).stateRoot == state.stateRoot) + #expect(try await provider.getState(hash: block.hash).value.stateRoot == state.value.stateRoot) // Test getting best state let bestState = try await provider.getBestState() - #expect(bestState.stateRoot == state.stateRoot) + #expect(await bestState.value.stateRoot == state.value.stateRoot) } @Test func testStateOperationsErrors() async throws { diff --git a/Node/Sources/Node/Genesis.swift b/Node/Sources/Node/Genesis.swift index 41aa55e1..1e419ac3 100644 --- a/Node/Sources/Node/Genesis.swift +++ b/Node/Sources/Node/Genesis.swift @@ -45,8 +45,10 @@ extension Genesis { let config = preset.config let (state, block) = try State.devGenesis(config: config) var kv = [String: Data]() - for (key, value) in try await state.value.toKV() { - kv[key.toHexString()] = value + for (key, value) in state.value.layer.toKV() { + if let value { + kv[key.encode().toHexString()] = try JamEncoder.encode(value) + } } return try ChainSpec( name: preset.rawValue, diff --git a/Node/Sources/Node/Node.swift b/Node/Sources/Node/Node.swift index 09a93907..438759d4 100644 --- a/Node/Sources/Node/Node.swift +++ b/Node/Sources/Node/Node.swift @@ -44,8 +44,8 @@ public class Node { let chainspec = try await genesis.load() let genesisBlock = try chainspec.getBlock() let genesisStateData = try chainspec.getState() - let rootHash = try stateMerklize(kv: genesisStateData) - let backend = try StateBackend(InMemoryBackend(store: genesisStateData), config: chainspec.getConfig(), rootHash: rootHash) + let backend = try StateBackend(InMemoryBackend(), config: chainspec.getConfig(), rootHash: Data32()) + try await backend.writeRaw(Array(genesisStateData)) let genesisState = try await State(backend: backend) let genesisStateRef = StateRef(genesisState) let protocolConfig = try chainspec.getConfig() From cac786f446fc4f7dc81152f71cf17866f72f687e Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Thu, 7 Nov 2024 22:25:43 +1300 Subject: [PATCH 05/10] almost there --- .../Blockchain/State/InMemoryBackend.swift | 8 +- .../State/StateBackendProtocol.swift | 2 +- .../Sources/Blockchain/State/StateTrie.swift | 70 ++++++--- .../BlockchainTests/StateTrieTests.swift | 137 ++++++++++++++++++ 4 files changed, 189 insertions(+), 28 deletions(-) create mode 100644 Blockchain/Tests/BlockchainTests/StateTrieTests.swift diff --git a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift index f8bb2803..0452f185 100644 --- a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift +++ b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift @@ -63,10 +63,10 @@ public actor InMemoryBackend: StateBackendProtocol { case let .writeRawValue(key, value): rawValues[key] = value rawValueRefCounts[key, default: 0] += 1 - case .refIncrement: - break - case .refDecrement: - break + case let .refIncrement(key): + refCounts[key, default: 0] += 1 + case let .refDecrement(key): + refCounts[key, default: 0] -= 1 } } } diff --git a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift index 0ffe1a3d..5c377dca 100644 --- a/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift +++ b/Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift @@ -8,7 +8,7 @@ public enum StateBackendOperation: Sendable { case refDecrement(key: Data) } -/// key: trie node hash (32 bytes) +/// key: trie node hash (31 bytes) /// value: trie node data (64 bytes) /// ref counting requirements: /// - write do not increment ref count, only explicit ref increment do diff --git a/Blockchain/Sources/Blockchain/State/StateTrie.swift b/Blockchain/Sources/Blockchain/State/StateTrie.swift index 89429e01..af5696e0 100644 --- a/Blockchain/Sources/Blockchain/State/StateTrie.swift +++ b/Blockchain/Sources/Blockchain/State/StateTrie.swift @@ -7,13 +7,20 @@ private enum TrieNodeType { case regularLeaf } +private func toId(hash: Data32) -> Data { + var id = hash.data + id[0] = id[0] & 0b0111_1111 // clear the highest bit + return id +} + private struct TrieNode { - var hash: Data32 - var left: Data32 - var right: Data32 - var type: TrieNodeType - var isNew: Bool - var rawValue: Data? + let hash: Data32 + let left: Data32 + let right: Data32 + let type: TrieNodeType + let isNew: Bool + let rawValue: Data? + let id: Data init(hash: Data32, data: Data64, isNew: Bool = false) { self.hash = hash @@ -29,6 +36,7 @@ private struct TrieNode { default: type = .branch } + id = toId(hash: hash) } private init(left: Data32, right: Data32, type: TrieNodeType, isNew: Bool, rawValue: Data?) { @@ -38,6 +46,7 @@ private struct TrieNode { self.type = type self.isNew = isNew self.rawValue = rawValue + id = toId(hash: hash) } var encodedData: Data64 { @@ -81,7 +90,9 @@ private struct TrieNode { } static func branch(left: Data32, right: Data32) -> TrieNode { - .init(left: left, right: right, type: .branch, isNew: true, rawValue: nil) + var left = left.data + left[0] = left[0] & 0b0111_1111 // clear the highest bit + return .init(left: Data32(left)!, right: right, type: .branch, isNew: true, rawValue: nil) } } @@ -93,8 +104,8 @@ public enum StateTrieError: Error { public actor StateTrie: Sendable { private let backend: StateBackendProtocol public private(set) var rootHash: Data32 - private var nodes: [Data32: TrieNode] = [:] - private var deleted: Set = [] + private var nodes: [Data: TrieNode] = [:] + private var deleted: Set = [] public init(rootHash: Data32, backend: StateBackendProtocol) { self.rootHash = rootHash @@ -133,13 +144,14 @@ public actor StateTrie: Sendable { if hash == Data32() { return nil } - if deleted.contains(hash) { + let id = toId(hash: hash) + if deleted.contains(id) { return nil } - if let node = nodes[hash] { + if let node = nodes[id] { return node } - guard let data = try await backend.read(key: hash.data) else { + guard let data = try await backend.read(key: id) else { return nil } @@ -168,20 +180,23 @@ public actor StateTrie: Sendable { var refChanges = [Data: Int]() // process deleted nodes - for hash in deleted { - let node = try await get(hash: hash).unwrap() + let deletedCopy = deleted + deleted.removeAll() + for id in deletedCopy { + guard let node = nodes[id] else { + continue + } if node.isBranch { // assign -1 to not worry about duplicates refChanges[node.hash.data] = -1 refChanges[node.left.data] = -1 refChanges[node.right.data] = -1 } - nodes.removeValue(forKey: hash) + nodes.removeValue(forKey: id) } - deleted.removeAll() for node in nodes.values where node.isNew { - ops.append(.write(key: node.hash.data, value: node.encodedData.data)) + ops.append(.write(key: node.id, value: node.encodedData.data)) if node.type == .regularLeaf { try ops.append(.writeRawValue(key: node.right, value: node.rawValue.unwrap())) } @@ -196,7 +211,11 @@ public actor StateTrie: Sendable { nodes.removeAll() + let zeros = Data(repeating: 0, count: 32) for (key, value) in refChanges { + if key == zeros { + continue + } if value > 0 { ops.append(.refIncrement(key: key)) } else if value < 0 { @@ -210,7 +229,11 @@ public actor StateTrie: Sendable { private func insert( hash: Data32, key: Data32, value: Data, depth: UInt8 ) async throws -> Data32 { - let parent = try await get(hash: hash).unwrap(orError: StateTrieError.invalidParent) + guard let parent = try await get(hash: hash) else { + let node = TrieNode.leaf(key: key, value: value) + saveNode(node: node) + return node.hash + } removeNode(hash: hash) if parent.isBranch { @@ -235,7 +258,7 @@ public actor StateTrie: Sendable { if existing.isLeaf(key: newKey) { // update existing leaf let newLeaf = TrieNode.leaf(key: newKey, value: newValue) - nodes[newLeaf.hash] = newLeaf + saveNode(node: newLeaf) return newLeaf.hash } @@ -297,13 +320,14 @@ public actor StateTrie: Sendable { } private func removeNode(hash: Data32) { - deleted.insert(hash) - nodes.removeValue(forKey: hash) + let id = toId(hash: hash) + deleted.insert(id) + nodes.removeValue(forKey: id) } private func saveNode(node: TrieNode) { - nodes[node.hash] = node - deleted.remove(node.hash) // TODO: maybe this is not needed + nodes[node.id] = node + deleted.remove(node.id) // TODO: maybe this is not needed } } diff --git a/Blockchain/Tests/BlockchainTests/StateTrieTests.swift b/Blockchain/Tests/BlockchainTests/StateTrieTests.swift new file mode 100644 index 00000000..43f520af --- /dev/null +++ b/Blockchain/Tests/BlockchainTests/StateTrieTests.swift @@ -0,0 +1,137 @@ +@testable import Blockchain +import Foundation +import Testing +import Utils + +struct StateTrieTests { + let backend = InMemoryBackend() + + // MARK: - Basic Operations Tests + + @Test + func testEmptyTrie() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let key = Data32.random() + let value = try await trie.read(key: key) + #expect(value == nil) + } + + @Test + func testInsertAndRetrieveSingleValue() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let key = Data([1]).blake2b256hash() + let value = Data("test value".utf8) + + try await trie.update([(key: key, value: value)]) + try await trie.save() + + let retrieved = try await trie.read(key: key) + #expect(retrieved == value) + } + + @Test + func testInsertAndRetrieveMultipleValues() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let pairs = (0 ..< 5).map { i in + let data = Data(String(i).utf8) + return (key: data.blake2b256hash(), value: data) + } + + try await trie.update(pairs) + try await trie.save() + + for (i, (key, value)) in pairs.enumerated() { + let retrieved = try await trie.read(key: key) + #expect(retrieved == value, "Failed at index \(i)") + } + } + + // MARK: - Update Tests + + @Test + func testUpdateExistingValue() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let key = Data32.random() + let value1 = Data("value1".utf8) + let value2 = Data("value2".utf8) + + try await trie.update([(key: key, value: value1)]) + try await trie.save() + + try await trie.update([(key: key, value: value2)]) + try await trie.save() + + let retrieved = try await trie.read(key: key) + #expect(retrieved == value2) + } + + @Test + func testDeleteValue() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let key = Data32.random() + let value = Data("test".utf8) + + try await trie.update([(key: key, value: value)]) + try await trie.save() + + try await trie.update([(key: key, value: nil)]) + try await trie.save() + + let retrieved = try await trie.read(key: key) + #expect(retrieved == nil) + } + + // MARK: - Large Value Tests + + @Test + func testLargeValue() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let key = Data32.random() + let value = Data(repeating: 0xFF, count: 1000) // Value larger than 32 bytes + + try await trie.update([(key: key, value: value)]) + try await trie.save() + + let retrieved = try await trie.read(key: key) + #expect(retrieved == value) + } + + // MARK: - Root Hash Tests + + @Test + func testRootHashChanges() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let initialRoot = await trie.rootHash + + let key = Data32.random() + let value = Data("test".utf8) + + try await trie.update([(key: key, value: value)]) + try await trie.save() + + let newRoot = await trie.rootHash + #expect(initialRoot != newRoot) + } + + @Test + func testRootHashConsistency() async throws { + let trie1 = StateTrie(rootHash: Data32(), backend: backend) + let trie2 = StateTrie(rootHash: Data32(), backend: backend) + + let pairs = (0 ..< 5).map { i in + let data = Data(String(i).utf8) + return (key: data.blake2b256hash(), value: data) + } + + // Apply same updates to both tries + try await trie1.update(pairs) + try await trie1.save() + + try await trie2.update(pairs) + try await trie2.save() + + #expect(await trie1.rootHash == trie2.rootHash) + } + + // TODO: test for gc, ref counting & pruning, raw value ref counting & cleaning +} From 4bafc337368e112c0e439a9c7b8f78ab11be4539 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Fri, 8 Nov 2024 14:03:00 +1300 Subject: [PATCH 06/10] fix optional --- Blockchain/Sources/Blockchain/State/StateLayer.swift | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Blockchain/Sources/Blockchain/State/StateLayer.swift b/Blockchain/Sources/Blockchain/State/StateLayer.swift index 6742dd45..6031e569 100644 --- a/Blockchain/Sources/Blockchain/State/StateLayer.swift +++ b/Blockchain/Sources/Blockchain/State/StateLayer.swift @@ -172,7 +172,7 @@ public struct StateLayer: @unchecked Sendable { // δ: The (prior) state of the service accounts. public subscript(serviceAccount index: ServiceIndex) -> StateKeys.ServiceAccountKey.Value? { get { - changes[StateKeys.ServiceAccountKey(index: index)]!.value()! + changes[StateKeys.ServiceAccountKey(index: index)]?.value() } set { changes[StateKeys.ServiceAccountKey(index: index)] = .init(newValue) @@ -182,7 +182,7 @@ public struct StateLayer: @unchecked Sendable { // s public subscript(serviceAccount index: ServiceIndex, storageKey key: Data32) -> StateKeys.ServiceAccountStorageKey.Value? { get { - changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)]!.value()! + changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)]?.value() } set { changes[StateKeys.ServiceAccountStorageKey(index: index, key: key)] = .init(newValue) @@ -194,7 +194,7 @@ public struct StateLayer: @unchecked Sendable { serviceAccount index: ServiceIndex, preimageHash hash: Data32 ) -> StateKeys.ServiceAccountPreimagesKey.Value? { get { - changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)]!.value()! + changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)]?.value() } set { changes[StateKeys.ServiceAccountPreimagesKey(index: index, hash: hash)] = .init(newValue) From e17f4e2a32d65cef198b3afebea71c1fb1d0ca83 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Fri, 8 Nov 2024 14:04:33 +1300 Subject: [PATCH 07/10] fix --- Blockchain/Sources/Blockchain/State/StateLayer.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Blockchain/Sources/Blockchain/State/StateLayer.swift b/Blockchain/Sources/Blockchain/State/StateLayer.swift index 6031e569..77838748 100644 --- a/Blockchain/Sources/Blockchain/State/StateLayer.swift +++ b/Blockchain/Sources/Blockchain/State/StateLayer.swift @@ -208,7 +208,7 @@ public struct StateLayer: @unchecked Sendable { get { changes[ StateKeys.ServiceAccountPreimageInfoKey(index: index, hash: hash, length: length) - ]!.value()! + ]?.value() } set { changes[StateKeys.ServiceAccountPreimageInfoKey(index: index, hash: hash, length: length)] = .init(newValue) From 4d75b2a9fb96308edf2701831cd9c073fa9ee2a3 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Fri, 8 Nov 2024 23:34:12 +1300 Subject: [PATCH 08/10] state trie fixed --- .../Blockchain/State/InMemoryBackend.swift | 30 ++++--- .../Sources/Blockchain/State/StateTrie.swift | 81 ++++++++++++------- .../BlockchainTests/StateTrieTests.swift | 64 ++++++++++++++- 3 files changed, 134 insertions(+), 41 deletions(-) diff --git a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift index 0452f185..af5dadce 100644 --- a/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift +++ b/Blockchain/Sources/Blockchain/State/InMemoryBackend.swift @@ -1,22 +1,25 @@ import Codec import Foundation +import TracingUtils import Utils -private struct KVPair: Comparable, Sendable { - var key: Data - var value: Data +private let logger = Logger(label: "InMemoryBackend") - public static func < (lhs: KVPair, rhs: KVPair) -> Bool { - lhs.key.lexicographicallyPrecedes(rhs.key) +public actor InMemoryBackend: StateBackendProtocol { + public struct KVPair: Comparable, Sendable { + var key: Data + var value: Data + + public static func < (lhs: KVPair, rhs: KVPair) -> Bool { + lhs.key.lexicographicallyPrecedes(rhs.key) + } } -} -public actor InMemoryBackend: StateBackendProtocol { // we really should be using Heap or some other Tree based structure here // but let's keep it simple for now - private var store: SortedArray = .init([]) + public private(set) var store: SortedArray = .init([]) private var rawValues: [Data32: Data] = [:] - private var refCounts: [Data: Int] = [:] + public private(set) var refCounts: [Data: Int] = [:] private var rawValueRefCounts: [Data32: Int] = [:] public init() {} @@ -92,4 +95,13 @@ public actor InMemoryBackend: StateBackendProtocol { } } } + + public func debugPrint() { + for item in store.array { + let refCount = refCounts[item.key, default: 0] + logger.info("key: \(item.key.toHexString())") + logger.info("value: \(item.value.toHexString())") + logger.info("ref count: \(refCount)") + } + } } diff --git a/Blockchain/Sources/Blockchain/State/StateTrie.swift b/Blockchain/Sources/Blockchain/State/StateTrie.swift index af5696e0..40e26a47 100644 --- a/Blockchain/Sources/Blockchain/State/StateTrie.swift +++ b/Blockchain/Sources/Blockchain/State/StateTrie.swift @@ -1,18 +1,15 @@ import Foundation +import TracingUtils import Utils +private let logger = Logger(label: "StateTrie") + private enum TrieNodeType { case branch case embeddedLeaf case regularLeaf } -private func toId(hash: Data32) -> Data { - var id = hash.data - id[0] = id[0] & 0b0111_1111 // clear the highest bit - return id -} - private struct TrieNode { let hash: Data32 let left: Data32 @@ -20,7 +17,6 @@ private struct TrieNode { let type: TrieNodeType let isNew: Bool let rawValue: Data? - let id: Data init(hash: Data32, data: Data64, isNew: Bool = false) { self.hash = hash @@ -36,7 +32,6 @@ private struct TrieNode { default: type = .branch } - id = toId(hash: hash) } private init(left: Data32, right: Data32, type: TrieNodeType, isNew: Bool, rawValue: Data?) { @@ -46,7 +41,6 @@ private struct TrieNode { self.type = type self.isNew = isNew self.rawValue = rawValue - id = toId(hash: hash) } var encodedData: Data64 { @@ -144,7 +138,7 @@ public actor StateTrie: Sendable { if hash == Data32() { return nil } - let id = toId(hash: hash) + let id = hash.data.suffix(31) if deleted.contains(id) { return nil } @@ -180,34 +174,33 @@ public actor StateTrie: Sendable { var refChanges = [Data: Int]() // process deleted nodes - let deletedCopy = deleted - deleted.removeAll() - for id in deletedCopy { + for id in deleted { guard let node = nodes[id] else { continue } if node.isBranch { // assign -1 to not worry about duplicates - refChanges[node.hash.data] = -1 - refChanges[node.left.data] = -1 - refChanges[node.right.data] = -1 + refChanges[node.hash.data.suffix(31)] = -1 + refChanges[node.left.data.suffix(31)] = -1 + refChanges[node.right.data.suffix(31)] = -1 } nodes.removeValue(forKey: id) } + deleted.removeAll() for node in nodes.values where node.isNew { - ops.append(.write(key: node.id, value: node.encodedData.data)) + ops.append(.write(key: node.hash.data.suffix(31), value: node.encodedData.data)) if node.type == .regularLeaf { try ops.append(.writeRawValue(key: node.right, value: node.rawValue.unwrap())) } if node.isBranch { - refChanges[node.left.data] = (refChanges[node.left.data] ?? 0) + 1 - refChanges[node.right.data] = (refChanges[node.right.data] ?? 0) + 1 + refChanges[node.left.data.suffix(31), default: 0] += 1 + refChanges[node.right.data.suffix(31), default: 0] += 1 } } // pin root node - refChanges[rootHash.data] = (refChanges[rootHash.data] ?? 0) + 1 + refChanges[rootHash.data.suffix(31), default: 0] += 1 nodes.removeAll() @@ -217,9 +210,9 @@ public actor StateTrie: Sendable { continue } if value > 0 { - ops.append(.refIncrement(key: key)) + ops.append(.refIncrement(key: key.suffix(31))) } else if value < 0 { - ops.append(.refDecrement(key: key)) + ops.append(.refDecrement(key: key.suffix(31))) } } @@ -234,9 +227,10 @@ public actor StateTrie: Sendable { saveNode(node: node) return node.hash } - removeNode(hash: hash) if parent.isBranch { + removeNode(hash: hash) + let bitValue = bitAt(key.data, position: depth) var left = parent.left var right = parent.right @@ -262,7 +256,7 @@ public actor StateTrie: Sendable { return newLeaf.hash } - let existingKeyBit = bitAt(existing.left.data, position: depth) + let existingKeyBit = bitAt(existing.left.data[1...], position: depth) let newKeyBit = bitAt(newKey.data, position: depth) if existingKeyBit == newKeyBit { @@ -292,9 +286,10 @@ public actor StateTrie: Sendable { private func delete(hash: Data32, key: Data32, depth: UInt8) async throws -> Data32 { let node = try await get(hash: hash).unwrap(orError: StateTrieError.invalidParent) - removeNode(hash: hash) if node.isBranch { + removeNode(hash: hash) + let bitValue = bitAt(key.data, position: depth) var left = node.left var right = node.right @@ -320,14 +315,44 @@ public actor StateTrie: Sendable { } private func removeNode(hash: Data32) { - let id = toId(hash: hash) + let id = hash.data.suffix(31) deleted.insert(id) nodes.removeValue(forKey: id) } private func saveNode(node: TrieNode) { - nodes[node.id] = node - deleted.remove(node.id) // TODO: maybe this is not needed + let id = node.hash.data.suffix(31) + nodes[id] = node + deleted.remove(id) // TODO: maybe this is not needed + } + + public func debugPrint() async throws { + func printNode(_ hash: Data32, depth: UInt8) async throws { + let prefix = String(repeating: " ", count: Int(depth)) + if hash == Data32() { + logger.info("\(prefix) nil") + return + } + let node = try await get(hash: hash) + guard let node else { + return logger.info("\(prefix) ????") + } + logger.info("\(prefix)\(node.hash.toHexString()) \(node.type)") + if node.isBranch { + logger.info("\(prefix) left:") + try await printNode(node.left, depth: depth + 1) + + logger.info("\(prefix) right:") + try await printNode(node.right, depth: depth + 1) + } else { + logger.info("\(prefix) key: \(node.left.toHexString())") + if let value = node.value { + logger.info("\(prefix) value: \(value.toHexString())") + } + } + } + + try await printNode(rootHash, depth: 0) } } diff --git a/Blockchain/Tests/BlockchainTests/StateTrieTests.swift b/Blockchain/Tests/BlockchainTests/StateTrieTests.swift index 43f520af..60526eaa 100644 --- a/Blockchain/Tests/BlockchainTests/StateTrieTests.swift +++ b/Blockchain/Tests/BlockchainTests/StateTrieTests.swift @@ -1,8 +1,20 @@ -@testable import Blockchain import Foundation import Testing +import TracingUtils import Utils +@testable import Blockchain + +private let logger = Logger(label: "StateTrieTests") + +private func merklize(_ data: some Sequence<(key: Data32, value: Data)>) -> Data32 { + var dict = [Data32: Data]() + for (key, value) in data { + dict[key] = value + } + return try! stateMerklize(kv: dict) +} + struct StateTrieTests { let backend = InMemoryBackend() @@ -29,15 +41,59 @@ struct StateTrieTests { #expect(retrieved == value) } + @Test + func testInsertAndRetrieveSimple() async throws { + let trie = StateTrie(rootHash: Data32(), backend: backend) + let remainKey = Data(repeating: 0, count: 31) + let pairs = [ + (key: Data32(Data([0b0000_0000]) + remainKey)!, value: Data([0])), + (key: Data32(Data([0b1000_0000]) + remainKey)!, value: Data([1])), + (key: Data32(Data([0b0100_0000]) + remainKey)!, value: Data([2])), + (key: Data32(Data([0b1100_0000]) + remainKey)!, value: Data([3])), + ] + + for (i, pair) in pairs.enumerated() { + try await trie.update([(key: pair.key, value: pair.value)]) + + let expectedRoot = merklize(pairs[0 ... i]) + let trieRoot = await trie.rootHash + #expect(expectedRoot == trieRoot) + } + + for (i, (key, value)) in pairs.enumerated() { + let retrieved = try await trie.read(key: key) + #expect(retrieved == value, "Failed at index \(i)") + } + + try await trie.save() + + for (i, (key, value)) in pairs.enumerated() { + let retrieved = try await trie.read(key: key) + #expect(retrieved == value, "Failed at index \(i)") + } + } + @Test func testInsertAndRetrieveMultipleValues() async throws { let trie = StateTrie(rootHash: Data32(), backend: backend) - let pairs = (0 ..< 5).map { i in - let data = Data(String(i).utf8) + let pairs = (0 ..< 50).map { i in + let data = Data([UInt8(i)]) return (key: data.blake2b256hash(), value: data) } - try await trie.update(pairs) + for (i, pair) in pairs.enumerated() { + try await trie.update([(key: pair.key, value: pair.value)]) + + let expectedRoot = merklize(pairs[0 ... i]) + let trieRoot = await trie.rootHash + #expect(expectedRoot == trieRoot) + } + + for (i, (key, value)) in pairs.enumerated() { + let retrieved = try await trie.read(key: key) + #expect(retrieved == value, "Failed at index \(i)") + } + try await trie.save() for (i, (key, value)) in pairs.enumerated() { From 52cc13b9838f352ba7665718fa54afbf74d09f70 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Sat, 9 Nov 2024 00:23:22 +1300 Subject: [PATCH 09/10] working --- Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift | 4 ++-- Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift | 3 +-- Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift | 3 ++- Blockchain/Tests/BlockchainTests/StateTrieTests.swift | 3 --- 4 files changed, 5 insertions(+), 8 deletions(-) diff --git a/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift b/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift index 102fd5cf..843dee75 100644 --- a/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift +++ b/Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift @@ -37,9 +37,9 @@ public final class Runtime { public let timeslot: TimeslotIndex public let stateRoot: Data32 - public init(timeslot: TimeslotIndex, stateRoot _: Data32) { + public init(timeslot: TimeslotIndex, stateRoot: Data32) { self.timeslot = timeslot - stateRoot = Data32() + self.stateRoot = stateRoot } } diff --git a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift index 9cea5de0..785a027f 100644 --- a/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift +++ b/Blockchain/Sources/Blockchain/Validator/BlockAuthor.swift @@ -54,6 +54,7 @@ public final class BlockAuthor: ServiceBase2, @unchecked Sendable { // TODO: verify we are indeed the block author let state = try await dataProvider.getState(hash: parentHash) + let stateRoot = await state.value.stateRoot let epoch = timeslot.timeslotToEpochIndex(config: config) let pendingTickets = await extrinsicPool.getPendingTickets(epoch: epoch) @@ -118,8 +119,6 @@ public final class BlockAuthor: ServiceBase2, @unchecked Sendable { extrinsics: extrinsic.tickets ) - let stateRoot = try await state.value.save() - let unsignedHeader = Header.Unsigned( parentHash: parentHash, priorStateRoot: stateRoot, diff --git a/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift b/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift index 6890933d..4a0b3ed5 100644 --- a/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift +++ b/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift @@ -22,6 +22,7 @@ struct BlockAuthorTests { let config = services.config let timeProvider = services.timeProvider let genesisState = services.genesisState + let stateRoot = await genesisState.value.stateRoot let timeslot = timeProvider.getTime().timeToTimeslot(config: config) @@ -44,7 +45,7 @@ struct BlockAuthorTests { // Verify block try _ = await runtime.apply(block: block, state: genesisState, context: .init( timeslot: timeslot + 1, - stateRoot: genesisState.value.stateRoot + stateRoot: stateRoot )) } diff --git a/Blockchain/Tests/BlockchainTests/StateTrieTests.swift b/Blockchain/Tests/BlockchainTests/StateTrieTests.swift index 60526eaa..7b090ce0 100644 --- a/Blockchain/Tests/BlockchainTests/StateTrieTests.swift +++ b/Blockchain/Tests/BlockchainTests/StateTrieTests.swift @@ -1,12 +1,9 @@ import Foundation import Testing -import TracingUtils import Utils @testable import Blockchain -private let logger = Logger(label: "StateTrieTests") - private func merklize(_ data: some Sequence<(key: Data32, value: Data)>) -> Data32 { var dict = [Data32: Data]() for (key, value) in data { From 69b6a928927f4d1d6c63bc6fc0199a3ef20b00ef Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Sat, 9 Nov 2024 00:34:56 +1300 Subject: [PATCH 10/10] fixes --- Node/Tests/NodeTests/ChainSpecTests.swift | 6 ++++-- RPC/Sources/RPC/Handlers/ChainHandler.swift | 15 +++++++++------ 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/Node/Tests/NodeTests/ChainSpecTests.swift b/Node/Tests/NodeTests/ChainSpecTests.swift index 10a04055..bf09de32 100644 --- a/Node/Tests/NodeTests/ChainSpecTests.swift +++ b/Node/Tests/NodeTests/ChainSpecTests.swift @@ -18,12 +18,14 @@ struct ChainSpecTests { for preset in GenesisPreset.allCases { let genesis = Genesis.preset(preset) let chainspec = try await genesis.load() - let backend = try StateBackend(InMemoryBackend(store: chainspec.getState()), config: chainspec.getConfig(), rootHash: Data32()) + let backend = try StateBackend(InMemoryBackend(), config: chainspec.getConfig(), rootHash: Data32()) + let state = try chainspec.getState() + try await backend.writeRaw(state.map { (key: $0.key, value: $0.value) }) let block = try chainspec.getBlock() let config = try chainspec.getConfig() let recentHistory = try await backend.read(StateKeys.RecentHistoryKey()) - #expect(recentHistory.items.last?.headerHash == block.hash) + #expect(recentHistory?.items.last?.headerHash == block.hash) // Verify config matches preset #expect(config == preset.config) diff --git a/RPC/Sources/RPC/Handlers/ChainHandler.swift b/RPC/Sources/RPC/Handlers/ChainHandler.swift index 880d261a..47686eb5 100644 --- a/RPC/Sources/RPC/Handlers/ChainHandler.swift +++ b/RPC/Sources/RPC/Handlers/ChainHandler.swift @@ -38,18 +38,21 @@ struct ChainHandler { throw JSONError(code: -32602, message: "Invalid block hash") } let state = try await source.getState(hash: data32) + guard let state else { + return JSON.null + } // return state root for now - return [ - "stateRoot": state?.stateRoot.description, - "blockHash": hash.description, + return await [ + "stateRoot": state.value.stateRoot.toHexString(), + "blockHash": hash, ] } else { // return best block state by default let block = try await source.getBestBlock() let state = try await source.getState(hash: block.hash) - return [ - "stateRoot": state?.stateRoot.description, - "blockHash": block.hash.description, + return await [ + "stateRoot": state?.value.stateRoot.toHexString(), + "blockHash": block.hash.toHexString(), ] } }