avoid JamCoder overhead in most of the stores

open-web3-stack · Dec 11, 2024 · 74b11e9 · 74b11e9
1 parent 7fb6e5e
commit 74b11e9
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 21 deletions.
diff --git a/Database/Sources/Database/BinaryCoder.swift b/Database/Sources/Database/BinaryCoder.swift
@@ -0,0 +1,97 @@
+import Blockchain
+import Foundation
+import RocksDBSwift
+import Utils
+
+protocol BinaryCodable {
+    func encode() throws -> Data
+    static func decode(from data: Data) throws -> Self
+}
+
+extension Data: BinaryCodable {
+    func encode() throws -> Data {
+        self
+    }
+
+    static func decode(from data: Data) throws -> Data {
+        data
+    }
+}
+
+extension Data32: BinaryCodable {
+    func encode() throws -> Data {
+        data
+    }
+
+    static func decode(from data: Data) throws -> Data32 {
+        try Data32(data).unwrap()
+    }
+}
+
+extension UInt32: BinaryCodable {
+    static func decode(from data: Data) throws -> Self {
+        guard data.count == 4 else {
+            throw DecodingError.dataCorrupted(
+                DecodingError.Context(
+                    codingPath: [],
+                    debugDescription: "Invalid data length"
+                )
+            )
+        }
+        return data.withUnsafeBytes { ptr in
+            ptr.loadUnaligned(as: Self.self)
+        }
+    }
+}
+
+extension Set<Data32>: BinaryCodable {
+    func encode() throws -> Data {
+        var data = Data(capacity: count * 32)
+        for element in self {
+            data.append(element.data)
+        }
+        return data
+    }
+
+    static func decode(from data: Data) throws -> Set<Element> {
+        guard data.count % 32 == 0 else {
+            throw DecodingError.dataCorrupted(
+                DecodingError.Context(
+                    codingPath: [],
+                    debugDescription: "Invalid data length"
+                )
+            )
+        }
+        var set = Set<Element>()
+        for i in stride(from: 0, to: data.count, by: 32) {
+            set.insert(Data32(data[relative: i ..< i + 32])!)
+        }
+        return set
+    }
+}
+
+struct BinaryCoder<Key: BinaryCodable, Value: BinaryCodable>: StoreCoder {
+    typealias Key = Key
+    typealias Value = Value
+
+    private let config: ProtocolConfigRef
+    private let prefix: Data?
+
+    init(config: ProtocolConfigRef, prefix: Data? = nil) {
+        self.config = config
+        self.prefix = prefix
+    }
+
+    func encode(key: Key) throws -> Data {
+        let encodedKey = try key.encode()
+        return prefix.map { $0 + encodedKey } ?? encodedKey
+    }
+
+    func encode(value: Value) throws -> Data {
+        try value.encode()
+    }
+
+    func decode(data: Data) throws -> Value {
+        try Value.decode(from: data)
+    }
+}
diff --git a/Database/Sources/Database/NoopCoder.swift → Database/Sources/Database/RawCoder.swift b/Database/Sources/Database/NoopCoder.swift → Database/Sources/Database/RawCoder.swift
@@ -1,7 +1,7 @@
 import Foundation
 import RocksDBSwift
 
-struct NoopCoder: StoreCoder {
+struct RawCoder: StoreCoder {
     typealias Key = Data
     typealias Value = Data
 

diff --git a/Database/Sources/Database/RocksDBBackend.swift b/Database/Sources/Database/RocksDBBackend.swift
@@ -14,32 +14,32 @@ public enum RocksDBBackendError: Error {
 public final class RocksDBBackend: Sendable {
     private let config: ProtocolConfigRef
     private let db: RocksDB<StoreId>
-    private let meta: Store<StoreId, NoopCoder>
+    private let meta: Store<StoreId, RawCoder>
     private let blocks: Store<StoreId, JamCoder<Data32, BlockRef>>
-    private let blockHashByTimeslot: Store<StoreId, JamCoder<TimeslotIndex, Set<Data32>>>
-    private let blockHashByNumber: Store<StoreId, JamCoder<UInt32, Set<Data32>>>
-    private let blockNumberByHash: Store<StoreId, JamCoder<Data32, UInt32>>
-    private let stateRootByHash: Store<StoreId, JamCoder<Data32, Data32>>
-    private let stateTrie: Store<StoreId, JamCoder<Data, Data>>
-    private let stateValue: Store<StoreId, JamCoder<Data32, Data>>
-    private let stateRefs: Store<StoreId, JamCoder<Data, UInt32>>
-    private let stateRefsRaw: Store<StoreId, JamCoder<Data32, UInt32>>
+    private let blockHashByTimeslot: Store<StoreId, BinaryCoder<TimeslotIndex, Set<Data32>>>
+    private let blockHashByNumber: Store<StoreId, BinaryCoder<UInt32, Set<Data32>>>
+    private let blockNumberByHash: Store<StoreId, BinaryCoder<Data32, UInt32>>
+    private let stateRootByHash: Store<StoreId, BinaryCoder<Data32, Data32>>
+    private let stateTrie: Store<StoreId, BinaryCoder<Data, Data>>
+    private let stateValue: Store<StoreId, BinaryCoder<Data32, Data>>
+    private let stateRefs: Store<StoreId, BinaryCoder<Data, UInt32>>
+    private let stateRefsRaw: Store<StoreId, BinaryCoder<Data32, UInt32>>
 
     public let genesisBlockHash: Data32
 
     public init(path: URL, config: ProtocolConfigRef, genesisBlock: BlockRef, genesisStateData: [Data32: Data]) async throws {
         self.config = config
         db = try RocksDB(path: path)
-        meta = Store(db: db, column: .meta, coder: NoopCoder())
+        meta = Store(db: db, column: .meta, coder: RawCoder())
         blocks = Store(db: db, column: .blocks, coder: JamCoder(config: config))
-        blockHashByTimeslot = Store(db: db, column: .blockIndexes, coder: JamCoder(config: config, prefix: Data([0])))
-        blockHashByNumber = Store(db: db, column: .blockIndexes, coder: JamCoder(config: config, prefix: Data([1])))
-        blockNumberByHash = Store(db: db, column: .blockIndexes, coder: JamCoder(config: config, prefix: Data([2])))
-        stateRootByHash = Store(db: db, column: .blockIndexes, coder: JamCoder(config: config, prefix: Data([3])))
-        stateTrie = Store(db: db, column: .state, coder: JamCoder(config: config, prefix: Data([0])))
-        stateValue = Store(db: db, column: .state, coder: JamCoder(config: config, prefix: Data([1])))
-        stateRefs = Store(db: db, column: .stateRefs, coder: JamCoder(config: config, prefix: Data([0])))
-        stateRefsRaw = Store(db: db, column: .stateRefs, coder: JamCoder(config: config, prefix: Data([1])))
+        blockHashByTimeslot = Store(db: db, column: .blockIndexes, coder: BinaryCoder(config: config, prefix: Data([0])))
+        blockHashByNumber = Store(db: db, column: .blockIndexes, coder: BinaryCoder(config: config, prefix: Data([1])))
+        blockNumberByHash = Store(db: db, column: .blockIndexes, coder: BinaryCoder(config: config, prefix: Data([2])))
+        stateRootByHash = Store(db: db, column: .blockIndexes, coder: BinaryCoder(config: config, prefix: Data([3])))
+        stateTrie = Store(db: db, column: .state, coder: BinaryCoder(config: config, prefix: Data([0])))
+        stateValue = Store(db: db, column: .state, coder: BinaryCoder(config: config, prefix: Data([1])))
+        stateRefs = Store(db: db, column: .stateRefs, coder: BinaryCoder(config: config, prefix: Data([0])))
+        stateRefsRaw = Store(db: db, column: .stateRefs, coder: BinaryCoder(config: config, prefix: Data([1])))
 
         genesisBlockHash = genesisBlock.hash
 

diff --git a/Database/Sources/RocksDBSwift/Store.swift b/Database/Sources/RocksDBSwift/Store.swift
@@ -1,8 +1,8 @@
 import Foundation
 
 public protocol StoreCoder<Key, Value>: Sendable {
-    associatedtype Key: Encodable
-    associatedtype Value: Decodable
+    associatedtype Key
+    associatedtype Value
 
     func encode(key: Key) throws -> Data
     func encode(value: Value) throws -> Data