fix remove service (#363)

* fix service deletion

* fix AccumulateTests

* add tests

* commit tests

* add littleendian

Author: qiweiii

Blockchain/Sources/Blockchain/RuntimeProtocols/Accumulation.swift

@@ -142,7 +142,7 @@ public struct AccountChanges: Sendable {
             case let .newAccount(index, account):
                 try await accounts.addNew(serviceAccount: index, account: account)
             case let .removeAccount(index):
-                accounts.remove(serviceAccount: index)
+                try await accounts.remove(serviceAccount: index)
             case let .updateAccount(index, account):
                 accounts.set(serviceAccount: index, account: account)
             case let .updateStorage(index, key, value):

Blockchain/Sources/Blockchain/State/ServiceAccounts.swift

@@ -20,6 +20,8 @@ public protocol ServiceAccounts: Sendable {
         length: UInt32,
         value: StateKeys.ServiceAccountPreimageInfoKey.Value?
     ) async throws
+
+    mutating func remove(serviceAccount index: ServiceIndex) async throws
 }
 
 public class ServiceAccountsRef: Ref<ServiceAccounts>, @unchecked Sendable {}
@@ -83,8 +85,8 @@ public class ServiceAccountsMutRef: @unchecked Sendable {
         changes.addNewAccount(index: index, account: account)
     }
 
-    public func remove(serviceAccount index: ServiceIndex) {
-        ref.value.set(serviceAccount: index, account: nil)
+    public func remove(serviceAccount index: ServiceIndex) async throws {
+        try await ref.value.remove(serviceAccount: index)
         changes.addRemovedAccount(index: index)
     }
 

Blockchain/Sources/Blockchain/State/State.swift

@@ -461,6 +461,24 @@ extension State: ServiceAccounts {
         // update value
         layer[serviceAccount: index, preimageHash: hash, length: length] = value
     }
+
+    // TODO: consider if keeping track of service keys at backend level is better
+    public mutating func remove(serviceAccount index: ServiceIndex) async throws {
+        layer[serviceAccount: index] = nil
+
+        let serviceByte = UInt8(index & 0xFF)
+        var startKeyData = Data(repeating: 0, count: 31)
+        startKeyData[0] = serviceByte
+        let startKey = Data31(startKeyData)!
+
+        let candidateKeys = try await backend.getKeys(Data([serviceByte]), startKey, nil)
+
+        for (keyData, _) in candidateKeys {
+            if let key = Data31(keyData), StateKeys.isServiceKey(key, serviceIndex: index) {
+                layer[key] = nil
+            }
+        }
+    }
 }
 
 extension State: Safrole {

Blockchain/Sources/Blockchain/State/StateKeys.swift

@@ -24,7 +24,7 @@ private func constructKey(_ idx: UInt8) -> Data31 {
 private func constructKey(_ idx: UInt8, _ service: ServiceIndex) -> Data31 {
     var data = Data(repeating: 0, count: 31)
     data[0] = idx
-    withUnsafeBytes(of: service) { ptr in
+    withUnsafeBytes(of: service.littleEndian) { ptr in
         data[1] = ptr.load(as: UInt8.self)
         data[3] = ptr.load(fromByteOffset: 1, as: UInt8.self)
         data[5] = ptr.load(fromByteOffset: 2, as: UInt8.self)
@@ -40,7 +40,7 @@ private func constructKey(_ service: ServiceIndex, _ val: UInt32, _ data: Data)
     let h = valEncoded + data
     let a = h.blake2b256hash().data[relative: 0 ..< 27]
 
-    withUnsafeBytes(of: service) { servicePtr in
+    withUnsafeBytes(of: service.littleEndian) { servicePtr in
         a.withUnsafeBytes { aPtr in
             stateKey.append(servicePtr.load(as: UInt8.self))
             stateKey.append(aPtr.load(as: UInt8.self))
@@ -335,3 +335,24 @@ public enum StateKeys {
         }
     }
 }
+
+extension StateKeys {
+    public static func isServiceKey(_ key: Data31, serviceIndex: ServiceIndex) -> Bool {
+        let keyData = key.data
+        let serviceBytes = withUnsafeBytes(of: serviceIndex.littleEndian) { Data($0) }
+
+        // service details
+        if keyData[relative: 0] == 255 {
+            return keyData[relative: 1] == serviceBytes[relative: 0] &&
+                keyData[relative: 3] == serviceBytes[relative: 1] &&
+                keyData[relative: 5] == serviceBytes[relative: 2] &&
+                keyData[relative: 7] == serviceBytes[relative: 3]
+        }
+
+        // other service keys
+        return keyData[relative: 0] == serviceBytes[relative: 0] &&
+            keyData[relative: 2] == serviceBytes[relative: 1] &&
+            keyData[relative: 4] == serviceBytes[relative: 2] &&
+            keyData[relative: 6] == serviceBytes[relative: 3]
+    }
+}

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

@@ -1130,7 +1130,7 @@ public class Eject: HostCall {
             // accumulating service definitely exist
             var destAccount = try await x.state.accounts.value.get(serviceAccount: x.serviceIndex)!
             destAccount.balance += ejectAccount!.balance
-            x.state.accounts.remove(serviceAccount: ejectIndex)
+            try await x.state.accounts.remove(serviceAccount: ejectIndex)
             x.state.accounts.set(serviceAccount: x.serviceIndex, account: destAccount)
             state.writeRegister(Registers.Index(raw: 7), HostCallResultCode.OK.rawValue)
         } else {

JAMTests/Tests/JAMTests/jamtestnet/FuzzTests.swift

@@ -56,10 +56,13 @@ struct FuzzTests {
             // example: ("0.7.0/1754982630", "00000004")
         ],
         expectFailure: [
+            ("0.7.0/1756790723", "00000011"),
+            ("0.7.0/1756791458", "00000041"),
+            ("0.7.0/1756814312", "00000025"),
+            ("0.7.0/1756792661", "00000027"),
         ],
         ignore: [
-            ("0.7.0/1756548583", "00000009"), // TODO: one storage / preimage mismatch
-            ("0.7.0/1756548706", "00000094"), // TODO: backend
+            ("0.7.0/1756548583", "00000009"), // TODO: hard
         ]
     ))
     func v070(_ input: TestInput) async throws {

JAMTests/Tests/JAMTests/jamtestnet/TraceTest.swift

@@ -54,6 +54,13 @@ enum TraceTest {
                 )
             }
 
+            // make sure we don't have extra keys
+            let allKeys = try await stateRef.value.backend.getKeys(nil, nil, nil)
+            for (key, _) in allKeys {
+                let data31 = Data31(key)!
+                #expect(expectedStateDict[data31] != nil, "extra key in boka post state: \(data31.toHexString())")
+            }
+
             let stateRoot = await stateRef.value.stateRoot
             #expect(stateRoot == testcase.postState.root)
         case .failure:

JAMTests/Tests/JAMTests/w3f/AccumulateTests.swift

@@ -164,6 +164,13 @@ private struct FullAccumulateState: Accumulation {
         // update value
         preimageInfo[index, default: [:]][hash] = value
     }
+
+    mutating func remove(serviceAccount index: ServiceIndex) async throws {
+        accounts[index] = nil
+        storages[index] = nil
+        preimages[index] = nil
+        preimageInfo[index] = nil
+    }
 }
 
 struct AccumulateTests {