refactor

packages/core/src/blockchain/index.ts

@@ -262,18 +262,15 @@ export class Blockchain {
       }
       if (this.lightClient) {
         try {
-          const blockData = await this.lightClient.queryBlock(number)
-          if (blockData && blockData.blocks.length > 0) {
-            const data = blockData.blocks[0]
-            const registry = await this.head.registry
-            const header = registry.createType<Header>('Header', hexToU8a(data.header))
-            const block = new Block(this, number, data.hash, undefined, {
-              header,
-              extrinsics: data.body,
-            })
-            this.#registerBlock(block)
-            return block
-          }
+          const data = await this.lightClient.queryBlock(number)
+          const registry = await this.head.registry
+          const header = registry.createType<Header>('Header', hexToU8a(data.header))
+          const block = new Block(this, number, data.hash, undefined, {
+            header,
+            extrinsics: data.body,
+          })
+          this.#registerBlock(block)
+          return block
         } catch (error) {
           logger.warn({ error }, `LightClient queryBlock ${number} failed`)
         }
@@ -295,16 +292,13 @@ export class Blockchain {
     const registry = await this.head.registry
     if (this.lightClient) {
       try {
-        const blockData = await this.lightClient.queryBlock(hash)
-        if (blockData && blockData.blocks.length > 0) {
-          const data = blockData.blocks[0]
-          const header = registry.createType<Header>('Header', hexToU8a(data.header))
-          const block = new Block(this, header.number.toNumber(), hash, undefined, {
-            header,
-            extrinsics: data.body,
-          })
-          return block
-        }
+        const data = await this.lightClient.queryBlock(hash)
+        const header = registry.createType<Header>('Header', hexToU8a(data.header))
+        const block = new Block(this, header.number.toNumber(), hash, undefined, {
+          header,
+          extrinsics: data.body,
+        })
+        return block
       } catch (error) {
         logger.warn({ error }, `LightClient queryBlock ${hash} failed`)
       }

packages/core/src/blockchain/storage-layer.ts

@@ -62,7 +62,7 @@ export class RemoteStorageLayer implements StorageLayerProvider {
 
     if (this.#lightClient) {
       try {
-        const entries = await this.#lightClient.queryStorage(this.#at as HexString, [key as HexString])
+        const entries = await this.#lightClient.queryStorage([key as HexString], this.#at as HexString)
         let maybeValue: HexString | undefined = undefined
         for (const [k, v] of entries) {
           this.#db?.saveStorage(this.#at as HexString, k, v)

packages/core/src/wasm-executor/browser-wasm-executor.js

@@ -54,12 +54,8 @@ const timerFinished = async (callback) => {
   return pkg.timer_finished(callback)
 }
 
-const storageRequest = async (chainId, req, callback) => {
-  return pkg.storage_request(chainId, req, callback)
-}
-
-const blocksRequest = async (chainId, req, callback) => {
-  return pkg.blocks_request(chainId, req, callback)
+const queryChain = async (chainId, requestId, request, retries, callback) => {
+  return pkg.query_chain(chainId, requestId, request, retries, callback)
 }
 
 const getPeers = async (chainId) => {
@@ -78,8 +74,7 @@ const wasmExecutor = {
   decodeProof,
   testing,
   startNetworkService,
-  storageRequest,
-  blocksRequest,
+  queryChain,
   getPeers,
   getLatestBlock,
   connectionStreamOpened,

packages/core/src/wasm-executor/executor.test.ts

@@ -179,9 +179,10 @@ describe('wasm', () => {
         '0x26aa394eea5630e07c48ae0c9558cef734abf5cb34d6244378cddbf18e849d96',
         '0x45323df7cc47150b3930e2666b0aa31362f8058e9dc65b738fce4a22e26fa4f2',
       ].map((key) =>
-        lightClient.queryStorage('0x15177d4bdc975077b85261c09503bf40932aae9d3a7a2e948870afe3432976be', [
-          key as HexString,
-        ]),
+        lightClient.queryStorage(
+          [key as HexString],
+          '0x15177d4bdc975077b85261c09503bf40932aae9d3a7a2e948870afe3432976be',
+        ),
       ),
     )
     expect(storage).toMatchSnapshot()

packages/core/src/wasm-executor/index.ts

@@ -10,12 +10,7 @@ import { defaultLogger, truncate } from '../logger.js'
 import { stripChildPrefix } from '../utils/index.js'
 
 import { LightClientConfig } from './light-client.js'
-import type {
-  BlockRequest,
-  JsLightClientCallback,
-  JsRuntimeCallback,
-  StorageRequest,
-} from '@acala-network/chopsticks-executor'
+import type { JsLightClientCallback, JsRuntimeCallback, Request } from '@acala-network/chopsticks-executor'
 export { JsRuntimeCallback }
 
 export type RuntimeVersion = {
@@ -82,8 +77,13 @@ export interface WasmExecutor {
   connectionReset: (connectionId: number, data: Uint8Array) => Promise<void>
   streamReset: (connectionId: number, streamId: number) => Promise<void>
   timerFinished: (callback: JsLightClientCallback) => Promise<void>
-  storageRequest: (chainId: number, req: StorageRequest, callback: JsLightClientCallback) => Promise<void>
-  blocksRequest: (chainId: number, req: BlockRequest, callback: JsLightClientCallback) => Promise<void>
+  queryChain: (
+    chainId: number,
+    requestId: number,
+    request: Request,
+    retries: number,
+    callback: JsLightClientCallback,
+  ) => Promise<void>
 }
 
 const logger = defaultLogger.child({ name: 'executor' })
@@ -242,14 +242,15 @@ export const startNetworkService = async (config: LightClientConfig, callback: J
   return worker.remote.startNetworkService(config, Comlink.proxy(callback))
 }
 
-export const storageRequest = async (chainId: number, req: StorageRequest, callback: JsLightClientCallback) => {
-  const worker = await getWorker()
-  return worker.remote.storageRequest(chainId, req, callback)
-}
-
-export const blocksRequest = async (chainId: number, req: BlockRequest, callback: JsLightClientCallback) => {
+export const queryChain = async (
+  chainId: number,
+  requestId: number,
+  request: Request,
+  retries: number,
+  callback: JsLightClientCallback,
+) => {
   const worker = await getWorker()
-  return worker.remote.blocksRequest(chainId, req, callback)
+  return worker.remote.queryChain(chainId, requestId, request, retries, callback)
 }
 
 export const getPeers = async (chainId: number) => {

packages/core/src/wasm-executor/light-client.ts

@@ -1,18 +1,17 @@
-import { BlockRequest, BlocksResponse, StorageRequest, StorageResponse } from '@acala-network/chopsticks-executor'
 import { HexString } from '@polkadot/util/types'
+import { Response } from '@acala-network/chopsticks-executor'
 import { WebSocket } from 'ws'
 import { stringToU8a } from '@polkadot/util'
 
 globalThis.WebSocket = typeof globalThis.WebSocket !== 'undefined' ? globalThis.WebSocket : (WebSocket as any)
 
 import { Deferred, defer } from '../utils/index.js'
 import {
-  blocksRequest,
   connectionReset,
   getLatestBlock,
   getPeers,
+  queryChain,
   startNetworkService,
-  storageRequest,
   streamMessage,
   streamWritableBytes,
   timerFinished,
@@ -85,8 +84,7 @@ export class LightClient {
   // blacklist of addresses that we have failed to connect to
   #blacklist: string[] = []
   #connections: Record<number, Connection> = {}
-  #storageResponse: Map<number, Deferred<StorageResponse>> = new Map()
-  #blockResponse: Map<number, Deferred<BlocksResponse>> = new Map()
+  #queryResponse: Map<number, Deferred<Response>> = new Map()
 
   #chainId = defer<number>()
 
@@ -159,14 +157,9 @@ export class LightClient {
     }
   }
 
-  async storageResponse(response: StorageResponse) {
-    this.#storageResponse.get(response.id)?.resolve(response)
-    this.#storageResponse.delete(response.id)
-  }
-
-  async blockResponse(response: BlocksResponse) {
-    this.#blockResponse.get(response.id)?.resolve(response)
-    this.#blockResponse.delete(response.id)
+  async queryResponse(requestId: number, response: Response) {
+    this.#queryResponse.get(requestId)?.resolve(response)
+    this.#queryResponse.delete(requestId)
   }
 
   streamSend(connectionId: number, data: Uint8Array) {
@@ -219,39 +212,60 @@ export class LightClient {
     }
   }
 
-  async queryStorage(blockHash: HexString, keys: HexString[]) {
+  async queryStorage(keys: HexString[], at: HexString) {
     const chainId = await this.#chainId.promise
-    const id = this.#requestId++
-    const deferred = defer<StorageResponse>()
-    this.#storageResponse.set(id, deferred)
-    await storageRequest(chainId, { id, blockHash, keys, retries: 10 } satisfies StorageRequest, this)
+    const requestId = this.#requestId++
+    const deferred = defer<Response>()
+    this.#queryResponse.set(requestId, deferred)
+    await queryChain(
+      chainId,
+      requestId,
+      {
+        storage: {
+          hash: at,
+          keys,
+        },
+      },
+      10,
+      this,
+    )
     const response = await deferred.promise
-    if (response.errorReason) {
-      throw new Error(response.errorReason)
+    if ('Error' in response) {
+      throw new Error(response.Error)
     }
-    return response.items
+    if ('Storage' in response) {
+      return response.Storage
+    }
+    throw new Error('Invalid response')
   }
 
   async queryBlock(block: HexString | number) {
     const chainId = await this.#chainId.promise
-    const id = this.#requestId++
-    const deferred = defer<BlocksResponse>()
-    this.#blockResponse.set(id, deferred)
-    await blocksRequest(
+    const requestId = this.#requestId++
+    const deferred = defer<Response>()
+    this.#queryResponse.set(requestId, deferred)
+    await queryChain(
       chainId,
+      requestId,
       {
-        id,
-        blockNumber: typeof block === 'number' ? block : null,
-        blockHash: typeof block === 'string' ? block : null,
-        retries: 10,
-      } satisfies BlockRequest,
+        block: {
+          number: typeof block === 'number' ? block : null,
+          hash: typeof block === 'string' ? block : null,
+          header: true,
+          body: true,
+        },
+      },
+      10,
       this,
     )
     const response = await deferred.promise
-    if (response.errorReason) {
-      throw new Error(response.errorReason)
+    if ('Error' in response) {
+      throw new Error(response.Error)
+    }
+    if ('Block' in response) {
+      return response.Block
     }
-    return response
+    throw new Error('Invalid response')
   }
 
   connectionStreamOpen(_connectionId: number) {}

packages/core/src/wasm-executor/node-wasm-executor.js

@@ -57,12 +57,8 @@ const timerFinished = async (callback) => {
   return pkg.timer_finished(callback)
 }
 
-const storageRequest = async (chainId, req, callback) => {
-  return pkg.storage_request(chainId, req, callback)
-}
-
-const blocksRequest = async (chainId, req, callback) => {
-  return pkg.blocks_request(chainId, req, callback)
+const queryChain = async (chainId, requestId, request, retries, callback) => {
+  return pkg.query_chain(chainId, requestId, request, retries, callback)
 }
 
 const getPeers = async (chainId) => {
@@ -81,8 +77,7 @@ const wasmExecutor = {
   decodeProof,
   testing,
   startNetworkService,
-  storageRequest,
-  blocksRequest,
+  queryChain,
   getPeers,
   getLatestBlock,
   connectionStreamOpened,