replace global functions with wasm js callback

executor/pkg/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@acala-network/chopsticks-executor",
   "description": "Chopsticks executor",
-  "version": "0.2.2",
+  "version": "0.2.3",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",

executor/src/lib.rs

@@ -4,7 +4,6 @@ use smoldot::json_rpc::methods::{HashHexString, HexString};
 use std::collections::BTreeMap;
 use wasm_bindgen::prelude::*;
 
-mod bindings;
 mod proof;
 mod task;
 
@@ -17,20 +16,32 @@ fn setup_console() {
 }
 
 #[wasm_bindgen]
-pub async fn get_runtime_version(code: JsValue) -> Result<JsValue, JsError> {
+extern "C" {
+    pub type JsCallback;
+
+    #[wasm_bindgen(structural, method, js_name = "getStorage")]
+    pub async fn get_storage(this: &JsCallback, key: JsValue) -> JsValue;
+
+    #[wasm_bindgen(structural, method, js_name = "getPrefixKeys")]
+    pub async fn get_prefix_keys(this: &JsCallback, key: JsValue) -> JsValue;
+
+    #[wasm_bindgen(structural, method, js_name = "getNextKey")]
+    pub async fn get_next_key(this: &JsCallback, key: JsValue) -> JsValue;
+}
+
+#[wasm_bindgen]
+pub async fn get_runtime_version(code: JsValue) -> Result<JsValue, JsValue> {
     setup_console();
 
     let code = serde_wasm_bindgen::from_value::<HexString>(code)?;
-    let runtime_version = task::runtime_version(code)
-        .await
-        .map_err(|e| JsError::new(e.as_str()))?;
+    let runtime_version = task::runtime_version(code).await?;
     let result = serde_wasm_bindgen::to_value(&runtime_version)?;
 
     Ok(result)
 }
 
 #[wasm_bindgen]
-pub async fn calculate_state_root(entries: JsValue) -> Result<JsValue, JsError> {
+pub async fn calculate_state_root(entries: JsValue) -> Result<JsValue, JsValue> {
     setup_console();
 
     let entries = serde_wasm_bindgen::from_value::<Vec<(HexString, HexString)>>(entries)?;
@@ -45,14 +56,13 @@ pub async fn decode_proof(
     root_trie_hash: JsValue,
     keys: JsValue,
     nodes: JsValue,
-) -> Result<JsValue, JsError> {
+) -> Result<JsValue, JsValue> {
     setup_console();
 
     let root_trie_hash = serde_wasm_bindgen::from_value::<HashHexString>(root_trie_hash)?;
     let keys = serde_wasm_bindgen::from_value::<Vec<HexString>>(keys)?;
     let nodes = serde_wasm_bindgen::from_value::<HexString>(nodes)?;
-    let entries =
-        proof::decode_proof(root_trie_hash, keys, nodes.0).map_err(|e| JsError::new(e.as_str()))?;
+    let entries = proof::decode_proof(root_trie_hash, keys, nodes.0)?;
     let result = serde_wasm_bindgen::to_value(&entries)?;
 
     Ok(result)
@@ -63,28 +73,25 @@ pub async fn create_proof(
     root_trie_hash: JsValue,
     nodes: JsValue,
     entries: JsValue,
-) -> Result<JsValue, JsError> {
+) -> Result<JsValue, JsValue> {
     setup_console();
 
     let root_trie_hash = serde_wasm_bindgen::from_value::<HashHexString>(root_trie_hash)?;
     let proof = serde_wasm_bindgen::from_value::<HexString>(nodes)?;
     let entries = serde_wasm_bindgen::from_value::<Vec<(HexString, HexString)>>(entries)?;
     let entries = BTreeMap::from_iter(entries.into_iter().map(|(key, value)| (key.0, value.0)));
-    let proof = proof::create_proof(root_trie_hash, proof.0, entries)
-        .map_err(|e| JsError::new(e.as_str()))?;
+    let proof = proof::create_proof(root_trie_hash, proof.0, entries)?;
     let result = serde_wasm_bindgen::to_value(&proof)?;
 
     Ok(result)
 }
 
 #[wasm_bindgen]
-pub async fn run_task(task: JsValue) -> Result<JsValue, JsError> {
+pub async fn run_task(task: JsValue, js: JsCallback) -> Result<JsValue, JsValue> {
     setup_console();
 
     let task = serde_wasm_bindgen::from_value::<task::TaskCall>(task)?;
-    let result = task::run_task(task)
-        .await
-        .map_err(|e| JsError::new(e.as_str()))?;
+    let result = task::run_task(task, js).await?;
     let result = serde_wasm_bindgen::to_value(&result)?;
 
     Ok(result)

executor/src/task.rs

@@ -15,8 +15,6 @@ use smoldot::{
 };
 use std::collections::BTreeMap;
 
-use crate::bindings::{get_next_key, get_prefix_keys, get_storage};
-
 #[derive(Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct RuntimeVersion {
@@ -62,7 +60,6 @@ impl RuntimeVersion {
 #[serde(rename_all = "camelCase")]
 pub struct TaskCall {
     wasm: HexString,
-    block_hash: HexString,
     calls: Option<Vec<(String, HexString)>>,
     storage: Vec<(HexString, Option<HexString>)>,
     mock_signature_host: bool,
@@ -87,7 +84,7 @@ fn is_magic_signature(signature: &[u8]) -> bool {
     signature.starts_with(&[0xde, 0xad, 0xbe, 0xef]) && signature[4..].iter().all(|&b| b == 0xcd)
 }
 
-pub async fn run_task(task: TaskCall) -> Result<TaskResponse, String> {
+pub async fn run_task(task: TaskCall, js: crate::JsCallback) -> Result<TaskResponse, String> {
     let mut storage_top_trie_changes = StorageDiff::from_iter(
         task.storage
             .into_iter()
@@ -117,8 +114,6 @@ pub async fn run_task(task: TaskCall) -> Result<TaskResponse, String> {
 
         log::trace!("Calling {}", call);
 
-        let hash = serde_wasm_bindgen::to_value(&task.block_hash).map_err(|e| e.to_string())?;
-
         let res = loop {
             vm = match vm {
                 RuntimeHostVm::Finished(res) => {
@@ -127,7 +122,7 @@ pub async fn run_task(task: TaskCall) -> Result<TaskResponse, String> {
                 RuntimeHostVm::StorageGet(req) => {
                     let key = HexString(req.key().as_ref().to_vec());
                     let key = serde_wasm_bindgen::to_value(&key).map_err(|e| e.to_string())?;
-                    let value = get_storage(hash.clone(), key).await;
+                    let value = js.get_storage(key).await;
                     let value = if value.is_string() {
                         let encoded = serde_wasm_bindgen::from_value::<HexString>(value)
                             .map(|x| x.0)
@@ -149,7 +144,7 @@ pub async fn run_task(task: TaskCall) -> Result<TaskResponse, String> {
                     } else {
                         let key = serde_wasm_bindgen::to_value(&HexString(prefix))
                             .map_err(|e| e.to_string())?;
-                        let keys = get_prefix_keys(hash.clone(), key).await;
+                        let keys = js.get_prefix_keys(key).await;
                         let keys = serde_wasm_bindgen::from_value::<Vec<HexString>>(keys)
                             .map(|x| x.into_iter().map(|x| x.0))
                             .map_err(|e| e.to_string())?;
@@ -159,11 +154,14 @@ pub async fn run_task(task: TaskCall) -> Result<TaskResponse, String> {
                 RuntimeHostVm::NextKey(req) => {
                     let key = HexString(req.key().as_ref().to_vec());
                     let key = serde_wasm_bindgen::to_value(&key).map_err(|e| e.to_string())?;
-                    let value = get_next_key(hash.clone(), key).await;
-                    let value = serde_wasm_bindgen::from_value::<HexString>(value)
-                        .map(|x| x.0)
-                        .map_err(|e| e.to_string())?;
-                    let value = if value.is_empty() { None } else { Some(value) };
+                    let value = js.get_next_key(key).await;
+                    let value = if value.is_string() {
+                        serde_wasm_bindgen::from_value::<HexString>(value)
+                            .map(|x| Some(x.0))
+                            .map_err(|e| e.to_string())?
+                    } else {
+                        None
+                    };
                     req.inject_key(value)
                 }
                 RuntimeHostVm::SignatureVerification(req) => {

src/blockchain/block.ts

@@ -10,7 +10,7 @@ import type { HexString } from '@polkadot/util/types'
 import { Blockchain } from '.'
 import { RemoteStorageLayer, StorageLayer, StorageLayerProvider, StorageValueKind } from './storage-layer'
 import { compactHex } from '../utils'
-import { getRuntimeVersion, runTask } from '../executor'
+import { getRuntimeVersion, runTask, taskHandler } from '../executor'
 import type { RuntimeVersion } from '../executor'
 
 export type TaskCallResponse = {
@@ -185,14 +185,16 @@ export class Block {
   get metadata(): Promise<HexString> {
     if (!this.#metadata) {
       this.#metadata = this.wasm.then(async (wasm) => {
-        const response = await runTask({
-          blockHash: this.hash as HexString,
-          wasm,
-          calls: [['Metadata_metadata', '0x']],
-          storage: [],
-          mockSignatureHost: this.#chain.mockSignatureHost,
-          allowUnresolvedImports: this.#chain.allowUnresolvedImports,
-        })
+        const response = await runTask(
+          {
+            wasm,
+            calls: [['Metadata_metadata', '0x']],
+            storage: [],
+            mockSignatureHost: this.#chain.mockSignatureHost,
+            allowUnresolvedImports: this.#chain.allowUnresolvedImports,
+          },
+          taskHandler(this)
+        )
         return compactHex(hexToU8a(response.Call.result))
       })
     }
@@ -215,14 +217,16 @@ export class Block {
     storage: [HexString, HexString | null][] = []
   ): Promise<TaskCallResponse> {
     const wasm = await this.wasm
-    const response = await runTask({
-      blockHash: this.hash as HexString,
-      wasm,
-      calls: [[method, args]],
-      storage,
-      mockSignatureHost: this.#chain.mockSignatureHost,
-      allowUnresolvedImports: this.#chain.allowUnresolvedImports,
-    })
+    const response = await runTask(
+      {
+        wasm,
+        calls: [[method, args]],
+        storage,
+        mockSignatureHost: this.#chain.mockSignatureHost,
+        allowUnresolvedImports: this.#chain.allowUnresolvedImports,
+      },
+      taskHandler(this)
+    )
     if (response.Call) return response.Call
     if (response.Error) throw Error(response.Error)
     throw Error('Unexpected response')

src/blockchain/index.ts

@@ -12,7 +12,6 @@ import { HeadState } from './head-state'
 import { InherentProvider } from './inherents'
 import { ResponseError } from '../rpc/shared'
 import { defaultLogger } from '../logger'
-import { setupBindings } from '../bindings'
 
 const logger = defaultLogger.child({ name: 'blockchain' })
 
@@ -57,8 +56,6 @@ export class Blockchain {
     this.#head = new Block(this, header.number, header.hash)
     this.#registerBlock(this.#head)
 
-    setupBindings(this)
-
     this.#txpool = new TxPool(this, inherentProvider, buildBlockMode)
 
     this.headState = new HeadState(this.#head)

src/executor.ts

@@ -11,6 +11,7 @@ import {
 } from '@polkadot/util'
 global.WebSocket = WebSocket
 
+import { Block } from './blockchain/block'
 import {
   calculate_state_root,
   create_proof,
@@ -63,16 +64,18 @@ export const createProof = async (trieRootHash: HexString, nodes: HexString[], e
   return { trieRootHash: result[0] as HexString, nodes: result[1] as HexString[] }
 }
 
-export const runTask = async (task: {
-  blockHash: HexString
-  wasm: HexString
-  calls: [string, HexString][]
-  storage: [HexString, HexString | null][]
-  mockSignatureHost: boolean
-  allowUnresolvedImports: boolean
-}) => {
+export const runTask = async (
+  task: {
+    wasm: HexString
+    calls: [string, HexString][]
+    storage: [HexString, HexString | null][]
+    mockSignatureHost: boolean
+    allowUnresolvedImports: boolean
+  },
+  callback: JsCallback
+) => {
   logger.trace({ task: { ...task, wasm: truncate(task.wasm) } }, 'taskRun')
-  const response = await run_task(task)
+  const response = await run_task(task, callback)
   if (response.Call) {
     logger.trace(
       { result: truncate(response.Call.result), storageDiff: truncateStorageDiff(response.Call.storageDiff) },
@@ -83,3 +86,24 @@ export const runTask = async (task: {
   }
   return response
 }
+
+interface JsCallback {
+  getStorage: (key: HexString) => Promise<string | undefined>
+  getPrefixKeys: (key: HexString) => Promise<string[]>
+  getNextKey: (key: HexString) => Promise<string | undefined>
+}
+
+export const taskHandler = (block: Block): JsCallback => {
+  return {
+    getStorage: async function (key: HexString) {
+      return block.get(key)
+    },
+    getPrefixKeys: async function (key: HexString) {
+      return block.getKeysPaged({ prefix: key, pageSize: 1000, startKey: key })
+    },
+    getNextKey: async function (key: HexString) {
+      const keys = await block.getKeysPaged({ prefix: key, pageSize: 1, startKey: key })
+      return keys[0]
+    },
+  }
+}