State trie backend

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)
 

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
+            self.stateRoot = stateRoot
         }
     }
 
@@ -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
         }
 

Blockchain/Sources/Blockchain/State/InMemoryBackend.swift

@@ -1,42 +1,107 @@
 import Codec
 import Foundation
+import TracingUtils
 import Utils
 
-public actor InMemoryBackend: StateBackend {
-    private let config: ProtocolConfigRef
-    private var store: [Data32: Data]
+private let logger = Logger(label: "InMemoryBackend")
 
-    public init(config: ProtocolConfigRef, store: [Data32: Data] = [:]) {
-        self.config = config
-        self.store = store
+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 func readImpl(_ key: any StateKey) async throws -> (Codable & Sendable)? {
-        guard let value = store[key.encode()] else {
-            return nil
+    // we really should be using Heap or some other Tree based structure here
+    // but let's keep it simple for now
+    public private(set) var store: SortedArray<KVPair> = .init([])
+    private var rawValues: [Data32: Data] = [:]
+    public private(set) var refCounts: [Data: Int] = [:]
+    private var rawValueRefCounts: [Data32: Int] = [:]
+
+    public init() {}
+
+    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 try JamDecoder.decode(key.decodeType(), from: value, withConfig: config)
+        return nil
     }
 
-    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 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 batchWrite(_ changes: [(key: any StateKey, value: Codable & Sendable)]) async throws {
-        for (key, value) in changes {
-            store[key.encode()] = try JamEncoder.encode(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 let .refIncrement(key):
+                refCounts[key, default: 0] += 1
+            case let .refDecrement(key):
+                refCounts[key, default: 0] -= 1
+            }
         }
     }
 
-    public func readAll() async throws -> [Data32: Data] {
-        store
+    public func readValue(hash: Data32) async throws -> Data? {
+        rawValues[hash]
     }
 
-    public func stateRoot() async throws -> Data32 {
-        // TODO: store intermediate state so we can calculate the root efficiently
-        try stateMerklize(kv: store)
+    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()))
+            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)
+                    }
+                }
+            }
+        }
+    }
+
+    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)")
+        }
     }
 }

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?
     )
 }