Skip to content

Commit

Permalink
almost there
Browse files Browse the repository at this point in the history
  • Loading branch information
xlc committed Nov 7, 2024
1 parent f6c08eb commit cac786f
Show file tree
Hide file tree
Showing 4 changed files with 189 additions and 28 deletions.
8 changes: 4 additions & 4 deletions Blockchain/Sources/Blockchain/State/InMemoryBackend.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
70 changes: 47 additions & 23 deletions Blockchain/Sources/Blockchain/State/StateTrie.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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?) {
Expand All @@ -38,6 +46,7 @@ private struct TrieNode {
self.type = type
self.isNew = isNew
self.rawValue = rawValue
id = toId(hash: hash)
}

var encodedData: Data64 {
Expand Down Expand Up @@ -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)
}
}

Expand All @@ -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<Data32> = []
private var nodes: [Data: TrieNode] = [:]
private var deleted: Set<Data> = []

public init(rootHash: Data32, backend: StateBackendProtocol) {
self.rootHash = rootHash
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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()))
}
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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
}

Expand Down Expand Up @@ -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
}
}

Expand Down
137 changes: 137 additions & 0 deletions Blockchain/Tests/BlockchainTests/StateTrieTests.swift
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit cac786f

Please sign in to comment.