State trie

Blockchain/Sources/Blockchain/Blockchain.swift

@@ -41,7 +41,7 @@ public final class Blockchain: ServiceBase, @unchecked Sendable {
             let parent = try await dataProvider.getState(hash: block.header.parentHash)
             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 runtime.apply(block: block, state: parent, context: .init(timeslot: timeslot + 1))
+            let state = try await runtime.apply(block: block, state: parent, context: .init(timeslot: timeslot + 1))
 
             try await dataProvider.blockImported(block: block, state: state)
 

Blockchain/Sources/Blockchain/RuntimeProtocols/AccumulateFunction.swift

@@ -74,6 +74,8 @@ public struct AccumlateResultContext {
 public protocol AccumulateFunction {
     func invoke(
         config: ProtocolConfigRef,
+        // prior accounts
+        accounts: ServiceAccounts,
         // u
         state: AccumulateState,
         // s
@@ -84,5 +86,5 @@ public protocol AccumulateFunction {
         arguments: [AccumulateArguments],
         initialIndex: ServiceIndex,
         timeslot: TimeslotIndex
-    ) throws -> (state: AccumulateState, transfers: [DeferredTransfers], result: Data32?, gas: Gas)
+    ) async throws -> (state: AccumulateState, transfers: [DeferredTransfers], result: Data32?, gas: Gas)
 }

Blockchain/Sources/Blockchain/RuntimeProtocols/Accumulation.swift

@@ -21,7 +21,7 @@ public struct AccumulationOutput {
     public var serviceAccounts: [ServiceIndex: ServiceAccount]
 }
 
-public protocol Accumulation {
+public protocol Accumulation: ServiceAccounts {
     var privilegedServices: PrivilegedServices { get }
     var validatorQueue: ConfigFixedSizeArray<
         ValidatorKey, ProtocolConfig.TotalNumberOfValidators
@@ -34,13 +34,12 @@ public protocol Accumulation {
         ProtocolConfig.TotalNumberOfCores
     > { get }
     var entropyPool: EntropyPool { get }
-    var serviceAccounts: [ServiceIndex: ServiceAccount] { get }
     var accumlateFunction: AccumulateFunction { get }
     var onTransferFunction: OnTransferFunction { get }
 }
 
 extension Accumulation {
-    public func update(config: ProtocolConfigRef, block: BlockRef, workReports: [WorkReport]) throws -> AccumulationOutput {
+    public func update(config: ProtocolConfigRef, block: BlockRef, workReports: [WorkReport]) async throws -> AccumulationOutput {
         var servicesGasRatio: [ServiceIndex: Gas] = [:]
         var servicesGas: [ServiceIndex: Gas] = [:]
 
@@ -50,15 +49,13 @@ extension Accumulation {
         }
 
         let totalGasRatio = workReports.flatMap(\.results).reduce(Gas(0)) { $0 + $1.gasRatio }
-        let totalMinimalGas = try workReports.flatMap(\.results)
-            .reduce(Gas(0)) {
-                try $0 + serviceAccounts[$1.serviceIndex].unwrap(orError: AccumulationError.invalidServiceIndex).minAccumlateGas
-            }
+        var totalMinimalGas = Gas(0)
         for report in workReports {
             for result in report.results {
                 servicesGasRatio[result.serviceIndex, default: Gas(0)] += result.gasRatio
-                servicesGas[result.serviceIndex, default: Gas(0)] += try serviceAccounts[result.serviceIndex]
-                    .unwrap(orError: AccumulationError.invalidServiceIndex).minAccumlateGas
+                let acc = try await get(serviceAccount: result.serviceIndex).unwrap(orError: AccumulationError.invalidServiceIndex)
+                totalMinimalGas += acc.minAccumlateGas
+                servicesGas[result.serviceIndex, default: Gas(0)] += acc.minAccumlateGas
             }
         }
         let remainingGas = config.value.coreAccumulationGas - totalMinimalGas
@@ -98,7 +95,7 @@ extension Accumulation {
             ProtocolConfig.TotalNumberOfCores
         >?
 
-        var newServiceAccounts = serviceAccounts
+        var newServiceAccounts = [ServiceIndex: ServiceAccount]()
 
         var transferReceivers = [ServiceIndex: [DeferredTransfers]]()
 
@@ -107,10 +104,11 @@ extension Accumulation {
                 assertionFailure("unreachable: service not found")
                 throw AccumulationError.invalidServiceIndex
             }
-            let (newState, transfers, commitment, _) = try accumlateFunction.invoke(
+            let (newState, transfers, commitment, _) = try await accumlateFunction.invoke(
                 config: config,
+                accounts: self,
                 state: AccumulateState(
-                    serviceAccounts: serviceAccounts,
+                    serviceAccounts: newServiceAccounts,
                     validatorQueue: validatorQueue,
                     authorizationQueue: authorizationQueue,
                     privilegedServices: privilegedServices
@@ -150,8 +148,9 @@ extension Accumulation {
         }
 
         for (service, transfers) in transferReceivers {
-            let acc = try serviceAccounts[service].unwrap(orError: AccumulationError.invalidServiceIndex)
-            guard let code = acc.preimages[acc.codeHash] else {
+            let acc = try await get(serviceAccount: service).unwrap(orError: AccumulationError.invalidServiceIndex)
+            let code = try await get(serviceAccount: service, preimageHash: acc.codeHash)
+            guard let code else {
                 continue
             }
             newServiceAccounts[service] = try onTransferFunction.invoke(

Blockchain/Sources/Blockchain/RuntimeProtocols/Guaranteeing.swift

@@ -36,9 +36,10 @@ public protocol Guaranteeing {
         >,
         ProtocolConfig.TotalNumberOfCores
     > { get }
-    var serviceAccounts: [ServiceIndex: ServiceAccount] { get }
     var recentHistory: RecentHistory { get }
     var offenders: Set<Ed25519PublicKey> { get }
+
+    func serviceAccount(index: ServiceIndex) -> ServiceAccountDetails?
 }
 
 extension Guaranteeing {
@@ -68,6 +69,12 @@ extension Guaranteeing {
         return toCoreAssignment(source, n: n, max: UInt32(config.value.totalNumberOfCores))
     }
 
+    public func requiredStorageKeys(extrinsic: ExtrinsicGuarantees) -> [any StateKey] {
+        extrinsic.guarantees
+            .flatMap(\.workReport.results)
+            .map { StateKeys.ServiceAccountKey(index: $0.serviceIndex) }
+    }
+
     public func update(
         config: ProtocolConfigRef,
         extrinsic: ExtrinsicGuarantees
@@ -131,7 +138,7 @@ extension Guaranteeing {
             }
 
             for result in report.results {
-                guard let acc = serviceAccounts[result.serviceIndex] else {
+                guard let acc = serviceAccount(index: result.serviceIndex) else {
                     throw .invalidServiceIndex
                 }
 

Blockchain/Sources/Blockchain/RuntimeProtocols/Runtime.swift

@@ -140,15 +140,15 @@ public final class Runtime {
         // TODO: validate other things
     }
 
-    public func apply(block: BlockRef, state prevState: StateRef, context: ApplyContext) throws(Error) -> StateRef {
-        let validatedBlock = try Result { try block.toValidated(config: config) }
+    public func apply(block: BlockRef, state prevState: StateRef, context: ApplyContext) async throws(Error) -> StateRef {
+        let validatedBlock = try Result(catching: { try block.toValidated(config: config) })
             .mapError(Error.validateError)
             .get()
 
-        return try apply(block: validatedBlock, state: prevState, context: context)
+        return try await apply(block: validatedBlock, state: prevState, context: context)
     }
 
-    public func apply(block: Validated<BlockRef>, state prevState: StateRef, context: ApplyContext) throws(Error) -> StateRef {
+    public func apply(block: Validated<BlockRef>, state prevState: StateRef, context: ApplyContext) async throws(Error) -> StateRef {
         try validate(block: block, state: prevState, context: context)
         let block = block.value
 
@@ -170,9 +170,22 @@ public final class Runtime {
 
             // depends on Safrole and Disputes
             let availableReports = try updateReports(block: block, state: &newState)
-            let res = try newState.update(config: config, block: block, workReports: availableReports)
+            let res = try await newState.update(config: config, block: block, workReports: availableReports)
             newState.privilegedServices = res.privilegedServices
-            newState.serviceAccounts = res.serviceAccounts
+
+            for (service, account) in res.serviceAccounts {
+                newState[serviceAccount: service] = account.toDetails()
+                for (hash, value) in account.storage {
+                    newState[serviceAccount: service, storageKey: hash] = value
+                }
+                for (hash, value) in account.preimages {
+                    newState[serviceAccount: service, preimageHash: hash] = value
+                }
+                for (hashLength, value) in account.preimageInfos {
+                    newState[serviceAccount: service, preimageHash: hashLength.hash, length: hashLength.length] = value
+                }
+            }
+
             newState.authorizationQueue = res.authorizationQueue
             newState.validatorQueue = res.validatorQueue
 
@@ -311,28 +324,27 @@ public final class Runtime {
         return availableReports
     }
 
-    public func updatePreimages(block: BlockRef, state newState: inout State) throws {
+    public func updatePreimages(block: BlockRef, state newState: inout State) async throws {
         let preimages = block.extrinsic.preimages.preimages
 
         guard preimages.isSortedAndUnique() else {
             throw Error.preimagesNotSorted
         }
 
         for preimage in preimages {
-            guard var acc = newState.serviceAccounts[preimage.serviceIndex] else {
-                throw Error.invalidPreimageServiceIndex
-            }
-
             let hash = preimage.data.blake2b256hash()
-            let hashAndLength = HashAndLength(hash: hash, length: UInt32(preimage.data.count))
-            guard acc.preimages[hash] == nil, acc.preimageInfos[hashAndLength] == nil else {
+            let preimageData: Data? = try await newState.get(serviceAccount: preimage.serviceIndex, preimageHash: hash)
+            let info = try await newState.get(
+                serviceAccount: preimage.serviceIndex, preimageHash: hash, length: UInt32(preimage.data.count)
+            )
+            guard preimageData == nil, info == nil else {
                 throw Error.duplicatedPreimage
             }
 
-            acc.preimages[hash] = preimage.data
-            acc.preimageInfos[hashAndLength] = .init([newState.timeslot])
-
-            newState.serviceAccounts[preimage.serviceIndex] = acc
+            newState[serviceAccount: preimage.serviceIndex, preimageHash: hash] = preimage.data
+            newState[
+                serviceAccount: preimage.serviceIndex, preimageHash: hash, length: UInt32(preimage.data.count)
+            ] = .init([newState.timeslot])
         }
     }
 

Blockchain/Sources/Blockchain/State/InMemoryBackend.swift

@@ -0,0 +1,42 @@
+import Codec
+import Foundation
+import Utils
+
+public actor InMemoryBackend: StateBackend {
+    private let config: ProtocolConfigRef
+    private var store: [Data32: Data]
+
+    public init(config: ProtocolConfigRef, store: [Data32: Data] = [:]) {
+        self.config = config
+        self.store = store
+    }
+
+    public func readImpl(_ key: any StateKey) async throws -> (Codable & Sendable)? {
+        guard let value = store[key.encode()] else {
+            return nil
+        }
+        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 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() async throws -> [Data32: Data] {
+        store
+    }
+
+    public func stateRoot() async throws -> Data32 {
+        // TODO: store intermediate state so we can calculate the root efficiently
+        try stateMerklize(kv: store)
+    }
+}

Blockchain/Sources/Blockchain/State/ServiceAccounts.swift

@@ -0,0 +1,21 @@
+import Foundation
+import Utils
+
+public protocol ServiceAccounts {
+    func get(serviceAccount index: ServiceIndex) async throws -> ServiceAccountDetails?
+    func get(serviceAccount index: ServiceIndex, storageKey key: Data32) async throws -> Data?
+    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
+
+    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
+    )
+}