fix state trie (#366)

* fix fuzz test

* add new traces

* fix state trie hash calulation

* swift 6.2 fix

* update docker for swift 6.2

* misc

Author: qiweiii

.swift-version

@@ -1 +1 @@
-6.1
+6.2

Blockchain/Sources/Blockchain/State/State.swift

@@ -238,19 +238,7 @@ public struct State: Sendable {
 
     public var stateRoot: Data32 {
         get async {
-            // TODO: should use backend.rootHash after StateTrie is fixed
-            do {
-                let allKeys = try await backend.getKeys(nil, nil, nil)
-                var kv: [Data31: Data] = [:]
-                for (key, value) in allKeys {
-                    if let key31 = Data31(key) {
-                        kv[key31] = value
-                    }
-                }
-                return try stateMerklize(kv: kv)
-            } catch {
-                return await backend.rootHash
-            }
+            await backend.rootHash
         }
     }
 

Blockchain/Sources/Blockchain/State/StateBackend.swift

@@ -55,18 +55,18 @@ public final class StateBackend: Sendable {
                 break
             }
 
-            guard trieNodeData.count == 64 else {
+            guard trieNodeData.count == 65 else {
                 continue
             }
 
             let firstByte = trieNodeData[relative: 0]
-            let isLeaf = (firstByte & 0b1100_0000) == 0b1000_0000 || (firstByte & 0b1100_0000) == 0b1100_0000
+            let isLeaf = firstByte == 1 || firstByte == 2
 
             guard isLeaf else {
                 continue
             }
 
-            let stateKey = Data(trieNodeData[relative: 1 ..< 32])
+            let stateKey = Data(trieNodeData[relative: 2 ..< 33])
 
             if !prefixData.isEmpty, !stateKey.starts(with: prefixData) {
                 continue
@@ -94,7 +94,9 @@ public final class StateBackend: Sendable {
     }
 
     public func write(_ values: any Sequence<(key: Data31, value: (Codable & Sendable)?)>) async throws {
-        try await trie.update(values.map { try (key: $0.key, value: $0.value.map { try JamEncoder.encode($0) }) })
+        let updates: [(key: Data31, value: Data?)] = try values.map { try (key: $0.key, value: $0.value.map { try JamEncoder.encode($0) }) }
+
+        try await trie.update(updates)
         try await trie.save()
     }
 
@@ -109,13 +111,13 @@ public final class StateBackend: Sendable {
 
     public func gc() async throws {
         try await impl.gc { data in
-            guard data.count == 64 else {
+            guard data.count == 65 else {
                 // unexpected data size
                 return nil
             }
-            let isRegularLeaf = data[0] & 0b1100_0000 == 0b1100_0000
+            let isRegularLeaf = data[0] == 2 // type byte for regularLeaf
             if isRegularLeaf {
-                return Data32(data.suffix(from: 32))!
+                return Data32(data.suffix(from: 33))! // right child starts at byte 33
             }
             return nil
         }

Blockchain/Sources/Blockchain/State/StateBackendProtocol.swift

@@ -13,7 +13,7 @@ public protocol StateBackendIterator: Sendable {
 }
 
 /// key: trie node hash (31 bytes)
-/// value: trie node data (64 bytes)
+/// value: trie node data (65 bytes - includes node type + original child data)
 /// 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

Blockchain/Sources/Blockchain/State/StateTrie.swift

@@ -12,39 +12,77 @@ private enum TrieNodeType {
 
 private struct TrieNode {
     let hash: Data32
-    let left: Data32
-    let right: Data32
+    let left: Data32 // Original child hash/data
+    let right: Data32 // Original child hash/data
     let type: TrieNodeType
     let isNew: Bool
     let rawValue: Data?
 
-    init(hash: Data32, data: Data64, isNew: Bool = false) {
+    // Constructor for loading from storage (65-byte format: [type][left-32][right-32])
+    init(hash: Data32, data: Data, 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.first! & 0b1100_0000 {
-        case 0b1000_0000:
+
+        let typeByte = data[relative: 0]
+        switch typeByte {
+        case 0:
+            type = .branch
+        case 1:
             type = .embeddedLeaf
-        case 0b1100_0000:
+        case 2:
             type = .regularLeaf
         default:
             type = .branch
         }
+
+        left = Data32(data[relative: 1 ..< 33])! // bytes 1-32
+        right = Data32(data[relative: 33 ..< 65])! // bytes 33-64
     }
 
+    // Constructor for pure trie operations
     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
+        hash = Self.calculateHash(left: left, right: right, type: type)
+        self.left = left // Store original data
+        self.right = right // Store original data
         self.type = type
         self.isNew = isNew
         self.rawValue = rawValue
     }
 
-    var encodedData: Data64 {
-        Data64(left.data + right.data)!
+    // JAM spec compliant hash calculation
+    private static func calculateHash(left: Data32, right: Data32, type: TrieNodeType) -> Data32 {
+        switch type {
+        case .branch:
+            var leftForHashing = left.data
+            leftForHashing[leftForHashing.startIndex] = leftForHashing[leftForHashing.startIndex] & 0b0111_1111
+            return Blake2b256.hash(leftForHashing, right.data)
+        case .embeddedLeaf:
+            var leftForHashing = left.data
+            let valueLength = leftForHashing[leftForHashing.startIndex]
+            leftForHashing[leftForHashing.startIndex] = 0b1000_0000 | valueLength
+            return Blake2b256.hash(leftForHashing, right.data)
+        case .regularLeaf:
+            var leftForHashing = left.data
+            leftForHashing[leftForHashing.startIndex] = 0b1100_0000
+            return Blake2b256.hash(leftForHashing, right.data)
+        }
+    }
+
+    // New 65-byte storage format: [type:1][left:32][right:32]
+    var storageData: Data {
+        var data = Data(capacity: 65)
+
+        switch type {
+        case .branch: data.append(0)
+        case .embeddedLeaf: data.append(1)
+        case .regularLeaf: data.append(2)
+        }
+
+        data.append(left.data)
+        data.append(right.data)
+
+        return data
     }
 
     var isBranch: Bool {
@@ -66,27 +104,30 @@ private struct TrieNode {
         guard type == .embeddedLeaf else {
             return nil
         }
-        let len = left.data.first! & 0b0011_1111
+        // For embedded leaves: length is stored in first byte
+        let len = left.data[relative: 0]
         return right.data[relative: 0 ..< Int(len)]
     }
 
     static func leaf(key: Data31, value: Data) -> TrieNode {
-        var newKey = Data(capacity: 32)
         if value.count <= 32 {
-            newKey.append(0b1000_0000 | UInt8(value.count))
-            newKey += key.data
-            let newValue = value + Data(repeating: 0, count: 32 - value.count)
-            return .init(left: Data32(newKey)!, right: Data32(newValue)!, type: .embeddedLeaf, isNew: true, rawValue: value)
+            // Embedded leaf: store length + key, padded value
+            var keyData = Data(capacity: 32)
+            keyData.append(UInt8(value.count)) // Store length in first byte
+            keyData += key.data
+            let paddedValue = value + Data(repeating: 0, count: 32 - value.count)
+            return .init(left: Data32(keyData)!, right: Data32(paddedValue)!, type: .embeddedLeaf, isNew: true, rawValue: value)
+        } else {
+            // Regular leaf: store key, value hash
+            var keyData = Data(capacity: 32)
+            keyData.append(0x00) // Placeholder for first byte
+            keyData += key.data
+            return .init(left: Data32(keyData)!, right: value.blake2b256hash(), type: .regularLeaf, isNew: true, rawValue: value)
         }
-        newKey.append(0b1100_0000)
-        newKey += key.data
-        return .init(left: Data32(newKey)!, right: value.blake2b256hash(), type: .regularLeaf, isNew: true, rawValue: value)
     }
 
     static func branch(left: Data32, right: Data32) -> TrieNode {
-        var left = left.data
-        left[left.startIndex] = left[left.startIndex] & 0b0111_1111 // clear the highest bit
-        return .init(left: Data32(left)!, right: right, type: .branch, isNew: true, rawValue: nil)
+        .init(left: left, right: right, type: .branch, isNew: true, rawValue: nil)
     }
 }
 
@@ -148,12 +189,10 @@ public actor StateTrie {
         guard let data = try await backend.read(key: id) else {
             return nil
         }
-
-        guard let data64 = Data64(data) else {
+        guard data.count == 65 else {
             throw StateTrieError.invalidData
         }
-
-        let node = TrieNode(hash: hash, data: data64)
+        let node = TrieNode(hash: hash, data: data)
         saveNode(node: node)
         return node
     }
@@ -189,7 +228,7 @@ public actor StateTrie {
         deleted.removeAll()
 
         for node in nodes.values where node.isNew {
-            ops.append(.write(key: node.hash.data.suffix(31), value: node.encodedData.data))
+            ops.append(.write(key: node.hash.data.suffix(31), value: node.storageData))
             if node.type == .regularLeaf {
                 try ops.append(.writeRawValue(key: node.right, value: node.rawValue.unwrap()))
             }
@@ -256,7 +295,7 @@ public actor StateTrie {
             return newLeaf.hash
         }
 
-        let existingKeyBit = bitAt(existing.left.data[1...], position: depth)
+        let existingKeyBit = bitAt(existing.left.data[relative: 1...], position: depth)
         let newKeyBit = bitAt(newKey.data, position: depth)
 
         if existingKeyBit == newKeyBit {
@@ -306,6 +345,30 @@ public actor StateTrie {
             if left == Data32(), right == Data32() {
                 // this branch is empty
                 return Data32()
+            } else if left == Data32() {
+                // only right child remains - check if we can collapse
+                let rightNode = try await get(hash: right)
+                if let rightNode, rightNode.isLeaf {
+                    // Can collapse: right child is a leaf
+                    return right
+                } else {
+                    // Cannot collapse: right child is a branch that needs to maintain its depth
+                    let newBranch = TrieNode.branch(left: left, right: right)
+                    saveNode(node: newBranch)
+                    return newBranch.hash
+                }
+            } else if right == Data32() {
+                // only left child remains - check if we can collapse
+                let leftNode = try await get(hash: left)
+                if let leftNode, leftNode.isLeaf {
+                    // Can collapse: left child is a leaf
+                    return left
+                } else {
+                    // Cannot collapse: left child is a branch that needs to maintain its depth
+                    let newBranch = TrieNode.branch(left: left, right: right)
+                    saveNode(node: newBranch)
+                    return newBranch.hash
+                }
             }
 
             let newBranch = TrieNode.branch(left: left, right: right)
@@ -331,7 +394,7 @@ public actor StateTrie {
     private func saveNode(node: TrieNode) {
         let id = node.hash.data.suffix(31)
         nodes[id] = node
-        deleted.remove(id) // TODO: maybe this is not needed
+        deleted.remove(id)
     }
 
     public func debugPrint() async throws {

Blockchain/Sources/Blockchain/VMInvocations/HostCall/HostCalls.swift

@@ -925,6 +925,7 @@ public class Checkpoint: HostCall {
         y.nextAccountIndex = x.nextAccountIndex
         y.transfers = x.transfers
         y.yield = x.yield
+        y.provide = x.provide
     }
 }
 

Blockchain/Tests/BlockchainTests/StateTrieTests.swift

@@ -301,5 +301,98 @@ struct StateTrieTests {
         }
     }
 
-    // TODO: test for gc, ref counting & pruning, raw value ref counting & cleaning
+    @Test
+    func testNodeDeletion() async throws {
+        let trie = StateTrie(rootHash: Data32(), backend: backend)
+
+        // Create a scenario that will cause node deletion and tree restructuring
+        let keys = [
+            Data31(Data([0x00] + Data(repeating: 0, count: 30)))!, // 00000000...
+            Data31(Data([0x40] + Data(repeating: 0, count: 30)))!, // 01000000...
+            Data31(Data([0x80] + Data(repeating: 0, count: 30)))!, // 10000000...
+            Data31(Data([0xC0] + Data(repeating: 0, count: 30)))!, // 11000000...
+        ]
+        let values = [Data([0x01]), Data([0x02]), Data([0x03]), Data([0x04])]
+
+        // Insert all keys
+        try await trie.update(Array(zip(keys, values.map { $0 as Data? })))
+        try await trie.save()
+
+        // Verify all are readable
+        for (i, (key, expectedValue)) in zip(keys, values).enumerated() {
+            let readValue = try await trie.read(key: key)
+            #expect(readValue == expectedValue, "Key \(i) should be readable after batch insert")
+        }
+
+        // Delete some keys to trigger tree restructuring
+        try await trie.update([
+            (key: keys[1], value: nil), // Delete 01000000...
+            (key: keys[3], value: nil), // Delete 11000000...
+        ])
+        try await trie.save()
+
+        // Verify remaining keys are still readable and deleted keys return nil
+        let remainingKeys = [keys[0], keys[2]]
+        let remainingValues = [values[0], values[2]]
+
+        for (i, (key, expectedValue)) in zip(remainingKeys, remainingValues).enumerated() {
+            let readValue = try await trie.read(key: key)
+            #expect(readValue == expectedValue, "Remaining key \(i) should still be readable after deletions")
+        }
+
+        let deletedValue1 = try await trie.read(key: keys[1])
+        let deletedValue3 = try await trie.read(key: keys[3])
+        #expect(deletedValue1 == nil, "Deleted key should return nil")
+        #expect(deletedValue3 == nil, "Deleted key should return nil")
+
+        // Test re-adding a deleted key
+        try await trie.update([(key: keys[1], value: Data([0xFF]))])
+        try await trie.save()
+
+        let readdedValue = try await trie.read(key: keys[1])
+        #expect(readdedValue == Data([0xFF]), "Re-added key should be readable")
+
+        // Ensure other keys are still intact
+        for (i, (key, expectedValue)) in zip(remainingKeys, remainingValues).enumerated() {
+            let readValue = try await trie.read(key: key)
+            #expect(readValue == expectedValue, "Key \(i) should remain readable after re-adding deleted key")
+        }
+    }
+
+    @Test
+    func testConsecutiveSaveAndLoad() async throws {
+        let trie = StateTrie(rootHash: Data32(), backend: backend)
+
+        let key = Data31.random()
+        let value1 = Data("value1".utf8)
+        let value2 = Data("value2".utf8)
+
+        // Insert, save, read
+        try await trie.update([(key: key, value: value1)])
+        try await trie.save()
+
+        let read1 = try await trie.read(key: key)
+        #expect(read1 == value1, "Value should be readable after first save")
+
+        // Update same key, save, read
+        try await trie.update([(key: key, value: value2)])
+        try await trie.save()
+
+        let read2 = try await trie.read(key: key)
+        #expect(read2 == value2, "Updated value should be readable after second save")
+
+        // Multiple saves shouldn't change anything
+        try await trie.save()
+        try await trie.save()
+
+        let read3 = try await trie.read(key: key)
+        #expect(read3 == value2, "Value should remain after multiple saves")
+
+        let rootHash1 = await trie.rootHash
+
+        // Another save shouldn't change root hash
+        try await trie.save()
+        let rootHash2 = await trie.rootHash
+        #expect(rootHash1 == rootHash2, "Root hash should be stable across saves")
+    }
 }