merge child with key

AcalaNetwork · Sep 12, 2023 · 1ca2a87 · 1ca2a87
1 parent a6de55c
commit 1ca2a87
Show file tree

Hide file tree

Showing 8 changed files with 139 additions and 109 deletions.
diff --git a/executor/src/lib.rs b/executor/src/lib.rs
@@ -20,9 +20,9 @@ fn setup_console(level: Option<String>) {
 const _: &'static str = r#"
 import { HexString } from '@polkadot/util/types';
 export interface JsCallback {
-	getStorage: (key: HexString, child?: HexString) => Promise<string | undefined>
+	getStorage: (key: HexString) => Promise<string | undefined>
 	getStateRoot: () => Promise<string>
-	getNextKey: (prefix: HexString, key: HexString, child?: HexString) => Promise<string | undefined>
+	getNextKey: (prefix: HexString, key: HexString) => Promise<string | undefined>
 	offchainGetStorage: (key: HexString) => Promise<string | undefined>
 	offchainTimestamp: () => Promise<number>
 	offchainRandomSeed: () => Promise<HexString>
@@ -36,18 +36,13 @@ extern "C" {
     pub type JsCallback;
 
     #[wasm_bindgen(structural, method, js_name = "getStorage")]
-    pub async fn get_storage(this: &JsCallback, key: JsValue, child: JsValue) -> JsValue;
+    pub async fn get_storage(this: &JsCallback, key: JsValue) -> JsValue;
 
     #[wasm_bindgen(structural, method, js_name = "getStateRoot")]
     pub async fn get_state_root(this: &JsCallback) -> JsValue;
 
     #[wasm_bindgen(structural, method, js_name = "getNextKey")]
-    pub async fn get_next_key(
-        this: &JsCallback,
-        prefix: JsValue,
-        key: JsValue,
-        child: JsValue,
-    ) -> JsValue;
+    pub async fn get_next_key(this: &JsCallback, prefix: JsValue, key: JsValue) -> JsValue;
 
     #[wasm_bindgen(structural, method, js_name = "offchainGetStorage")]
     pub async fn offchain_get_storage(this: &JsCallback, key: JsValue) -> JsValue;

diff --git a/executor/src/task.rs b/executor/src/task.rs
@@ -88,6 +88,8 @@ fn is_magic_signature(signature: &[u8]) -> bool {
     signature.starts_with(&[0xde, 0xad, 0xbe, 0xef]) && signature[4..].iter().all(|&b| b == 0xcd)
 }
 
+const DEFAULT_CHILD_STORAGE_PREFIX: &[u8] = b":child_storage:default:";
+
 pub async fn run_task(task: TaskCall, js: crate::JsCallback) -> Result<TaskResponse, String> {
     let mut storage_main_trie_changes = TrieDiff::default();
     let mut child_storage_changes: BTreeMap<Vec<u8>, Option<Vec<u8>>> = Default::default();
@@ -123,20 +125,22 @@ pub async fn run_task(task: TaskCall, js: crate::JsCallback) -> Result<TaskRespo
                 }
 
                 RuntimeHostVm::StorageGet(req) => {
-                    let key = HexString(req.key().as_ref().to_vec());
-                    let key = to_value(&key).map_err(|e| e.to_string())?;
-                    let child = req
-                        .child_trie()
-                        .map(|x| {
-                            b":child_storage:default:"
+                    let key = if let Some(child) = req.child_trie() {
+                        HexString(
+                            DEFAULT_CHILD_STORAGE_PREFIX
                                 .iter()
-                                .chain(x.as_ref())
+                                .chain(child.as_ref())
+                                .chain(req.key().as_ref())
                                 .map(|x| x.to_owned())
-                                .collect::<Vec<u8>>()
-                        })
-                        .map(HexString);
-                    let child = to_value(&child).map_err(|e| e.to_string())?;
-                    let value = js.get_storage(key, child).await;
+                                .collect::<Vec<u8>>(),
+                        )
+                    } else {
+                        HexString(req.key().as_ref().to_vec())
+                    };
+
+                    let key = to_value(&key).map_err(|e| e.to_string())?;
+
+                    let value = js.get_storage(key).await;
                     let value = if value.is_string() {
                         let encoded = from_value::<HexString>(value)
                             .map(|x| x.0)
@@ -161,26 +165,35 @@ pub async fn run_task(task: TaskCall, js: crate::JsCallback) -> Result<TaskRespo
                         // root_calculation, skip
                         req.inject_key(None::<Vec<_>>.map(|x| x.into_iter()))
                     } else {
-                        let prefix = HexString(
-                            nibbles_to_bytes_suffix_extend(req.prefix()).collect::<Vec<_>>(),
-                        );
-                        let key = HexString(
-                            nibbles_to_bytes_suffix_extend(req.key()).collect::<Vec<_>>(),
-                        );
-                        let child = req
-                            .child_trie()
-                            .map(|x| {
-                                b":child_storage:default:"
+                        let prefix = if let Some(child) = req.child_trie() {
+                            HexString(
+                                DEFAULT_CHILD_STORAGE_PREFIX
                                     .iter()
-                                    .chain(x.as_ref())
+                                    .chain(child.as_ref())
                                     .map(|x| x.to_owned())
-                                    .collect::<Vec<u8>>()
-                            })
-                            .map(HexString);
+                                    .chain(nibbles_to_bytes_suffix_extend(req.prefix()))
+                                    .collect::<Vec<u8>>(),
+                            )
+                        } else {
+                            HexString(
+                                nibbles_to_bytes_suffix_extend(req.prefix()).collect::<Vec<_>>(),
+                            )
+                        };
+                        let key = if let Some(child) = req.child_trie() {
+                            HexString(
+                                DEFAULT_CHILD_STORAGE_PREFIX
+                                    .iter()
+                                    .chain(child.as_ref())
+                                    .map(|x| x.to_owned())
+                                    .chain(nibbles_to_bytes_suffix_extend(req.key()))
+                                    .collect::<Vec<u8>>(),
+                            )
+                        } else {
+                            HexString(nibbles_to_bytes_suffix_extend(req.key()).collect::<Vec<_>>())
+                        };
                         let prefix = to_value(&prefix).map_err(|e| e.to_string())?;
                         let key = to_value(&key).map_err(|e| e.to_string())?;
-                        let child = to_value(&child).map_err(|e| e.to_string())?;
-                        let value = js.get_next_key(prefix, key, child).await;
+                        let value = js.get_next_key(prefix, key).await;
                         let value = if value.is_string() {
                             from_value::<HexString>(value)
                                 .map(|x| Some(x.0))
@@ -288,7 +301,7 @@ pub async fn run_task(task: TaskCall, js: crate::JsCallback) -> Result<TaskRespo
                         let key = nibbles_to_bytes_suffix_extend(key.iter().map(|x| x.to_owned()))
                             .collect::<Vec<_>>();
 
-                        let final_key = b":child_storage:default:"
+                        let final_key = DEFAULT_CHILD_STORAGE_PREFIX
                             .iter()
                             .chain(child.unwrap().iter())
                             .chain(key.iter())

diff --git a/packages/chopsticks/src/rpc/substrate/state.ts b/packages/chopsticks/src/rpc/substrate/state.ts
@@ -1,4 +1,4 @@
-import { Block } from '@acala-network/chopsticks-core'
+import { Block, isChild, mergeKey, stripChild } from '@acala-network/chopsticks-core'
 import { Handlers, ResponseError } from '../shared'
 import { defaultLogger } from '../../logger'
 
@@ -84,6 +84,22 @@ const handlers: Handlers = {
   state_unsubscribeStorage: async (_context, [subid], { unsubscribe }) => {
     unsubscribe(subid)
   },
+  childstate_getStorage: async (context, [child, key, hash]) => {
+    if (!isChild(child)) {
+      throw new ResponseError(-32000, 'Client error: Invalid child storage key')
+    }
+    const block = await context.chain.getBlock(hash)
+    return block?.get(mergeKey(child, key))
+  },
+  childstate_getKeysPaged: async (context, [child, prefix, pageSize, startKey, hash]) => {
+    if (!isChild(child)) {
+      throw new ResponseError(-32000, 'Client error: Invalid child storage key')
+    }
+    const block = await context.chain.getBlock(hash)
+    return block
+      ?.getKeysPaged({ prefix: mergeKey(child, prefix), pageSize, startKey: mergeKey(child, startKey) })
+      .then((keys) => keys.map(stripChild))
+  },
 }
 
 export default handlers
diff --git a/packages/core/src/api.ts b/packages/core/src/api.ts
@@ -1,7 +1,7 @@
 import { ExtDef } from '@polkadot/types/extrinsic/signedExtensions/types'
 import { HexString } from '@polkadot/util/types'
 import { ProviderInterface } from '@polkadot/rpc-provider/types'
-import { mergeKey } from './utils'
+import { mergeKey, splitChild, stripChild } from './utils'
 
 type ChainProperties = {
   ss58Format?: number
@@ -103,28 +103,31 @@ export class Api {
   }
 
   async getStorage(key: string, hash?: string) {
-    const params = [key]
-    if (hash) params.push(hash)
-    return this.#provider.send<string | null>('state_getStorage', params)
+    const parts = splitChild(key)
+    if (parts) {
+      const params = parts
+      if (hash) params.push(hash)
+      return this.#provider.send<string | null>('childstate_getStorage', params)
+    } else {
+      const params = [key]
+      if (hash) params.push(hash)
+      return this.#provider.send<string | null>('state_getStorage', params)
+    }
   }
 
   async getKeysPaged(prefix: string, pageSize: number, startKey: string, hash?: string) {
-    const params = [prefix, pageSize, startKey]
-    if (hash) params.push(hash)
-    return this.#provider.send<string[]>('state_getKeysPaged', params)
-  }
-
-  async getChildStorage(child: string, key: string, hash?: string) {
-    const params = [child, key]
-    if (hash) params.push(hash)
-    return this.#provider.send<string | null>('childstate_getStorage', params)
-  }
-
-  async getChildKeysPaged(child: string, prefix: string, pageSize: number, startKey: string, hash?: string) {
-    const params = [child, prefix, pageSize, startKey]
-    if (hash) params.push(hash)
-    return this.#provider
-      .send<string[]>('childstate_getKeysPaged', params)
-      .then((keys) => keys.map((k) => mergeKey(child, k)))
+    const parts = splitChild(prefix)
+    if (parts) {
+      const child = parts[0]
+      const params = [...parts, pageSize, stripChild(startKey)]
+      if (hash) params.push(hash)
+      return this.#provider
+        .send<string[]>('childstate_getKeysPaged', params)
+        .then((keys) => keys.map((k) => mergeKey(child, k)))
+    } else {
+      const params = [prefix, pageSize, startKey]
+      if (hash) params.push(hash)
+      return this.#provider.send<string[]>('state_getKeysPaged', params)
+    }
   }
 }
diff --git a/packages/core/src/blockchain/block.ts b/packages/core/src/blockchain/block.ts
@@ -110,8 +110,8 @@ export class Block {
     return this.#storages[this.#storages.length - 1] ?? this.#baseStorage
   }
 
-  async get(key: string, child?: string): Promise<string | undefined> {
-    const val = await this.storage.get(key, true, child)
+  async get(key: string): Promise<string | undefined> {
+    const val = await this.storage.get(key, true)
     switch (val) {
       case StorageValueKind.Deleted:
         return undefined
@@ -120,20 +120,15 @@ export class Block {
     }
   }
 
-  async getKeysPaged(options: {
-    prefix?: string
-    startKey?: string
-    pageSize: number
-    child?: string
-  }): Promise<string[]> {
+  async getKeysPaged(options: { prefix?: string; startKey?: string; pageSize: number }): Promise<string[]> {
     const layer = new StorageLayer(this.storage)
     await layer.fold()
 
     const prefix = options.prefix ?? '0x'
-    const startKey = options.startKey ?? prefix
+    const startKey = options.startKey ?? '0x'
     const pageSize = options.pageSize
 
-    return layer.getKeysPaged(prefix, pageSize, startKey, options.child)
+    return layer.getKeysPaged(prefix, pageSize, startKey)
   }
 
   pushStorageLayer(): StorageLayer {

diff --git a/packages/core/src/blockchain/storage-layer.ts b/packages/core/src/blockchain/storage-layer.ts
@@ -4,7 +4,6 @@ import _ from 'lodash'
 import { Api } from '../api'
 import { KeyValuePair } from '../db/entities'
 import { defaultLogger } from '../logger'
-import { mergeKey } from '../utils'
 import KeyCache from '../utils/key-cache'
 
 const logger = defaultLogger.child({ name: 'layer' })
@@ -19,11 +18,11 @@ export const enum StorageValueKind {
 export type StorageValue = string | StorageValueKind | undefined
 
 export interface StorageLayerProvider {
-  get(key: string, cache: boolean, child?: string): Promise<StorageValue>
+  get(key: string, cache: boolean): Promise<StorageValue>
   foldInto(into: StorageLayer): Promise<StorageLayerProvider | undefined>
   fold(): Promise<void>
 
-  getKeysPaged(prefix: string, pageSize: number, startKey: string, child?: string): Promise<string[]>
+  getKeysPaged(prefix: string, pageSize: number, startKey: string): Promise<string[]>
 }
 
 export class RemoteStorageLayer implements StorageLayerProvider {
@@ -38,20 +37,17 @@ export class RemoteStorageLayer implements StorageLayerProvider {
     this.#db = db
   }
 
-  async get(key: string, _cache: boolean, child?: string): Promise<StorageValue> {
-    const storageKey = mergeKey(child, key)
+  async get(key: string, _cache: boolean): Promise<StorageValue> {
     const keyValuePair = this.#db?.getRepository(KeyValuePair)
     if (this.#db) {
-      const res = await keyValuePair?.findOne({ where: { key: storageKey, blockHash: this.#at } })
+      const res = await keyValuePair?.findOne({ where: { key, blockHash: this.#at } })
       if (res) {
         return res.value ?? undefined
       }
     }
-    logger.trace({ at: this.#at, key, child }, 'RemoteStorageLayer get')
-    const data = child
-      ? await this.#api.getChildStorage(child, key, this.#at)
-      : await this.#api.getStorage(key, this.#at)
-    keyValuePair?.upsert({ key: storageKey, blockHash: this.#at, value: data }, ['key', 'blockHash'])
+    logger.trace({ at: this.#at, key }, 'RemoteStorageLayer get')
+    const data = await this.#api.getStorage(key, this.#at)
+    keyValuePair?.upsert({ key, blockHash: this.#at, value: data }, ['key', 'blockHash'])
     return data ?? undefined
   }
 
@@ -60,16 +56,12 @@ export class RemoteStorageLayer implements StorageLayerProvider {
   }
   async fold(): Promise<void> {}
 
-  async getKeysPaged(prefix: string, pageSize: number, startKey: string, child?: string): Promise<string[]> {
+  async getKeysPaged(prefix: string, pageSize: number, startKey: string): Promise<string[]> {
     if (pageSize > BATCH_SIZE) throw new Error(`pageSize must be less or equal to ${BATCH_SIZE}`)
-    logger.trace({ at: this.#at, prefix, pageSize, startKey, child }, 'RemoteStorageLayer getKeysPaged')
+    logger.trace({ at: this.#at, prefix, pageSize, startKey }, 'RemoteStorageLayer getKeysPaged')
     // can't handle keyCache without prefix
-    // can't handle keyCache for child keys
-    // TODO: cache child keys
-    if (prefix.length < 66 || !!child) {
-      return child
-        ? this.#api.getChildKeysPaged(child, prefix, pageSize, startKey, this.#at)
-        : this.#api.getKeysPaged(prefix, pageSize, startKey, this.#at)
+    if (prefix.length < 66) {
+      return this.#api.getKeysPaged(prefix, pageSize, startKey, this.#at)
     }
 
     let batchComplete = false
@@ -131,20 +123,19 @@ export class StorageLayer implements StorageLayerProvider {
     }
   }
 
-  async get(key: string, cache: boolean, child?: string): Promise<StorageValue | undefined> {
-    const storageKey = mergeKey(child, key)
-    if (storageKey in this.#store) {
-      return this.#store[storageKey]
+  async get(key: string, cache: boolean): Promise<StorageValue | undefined> {
+    if (key in this.#store) {
+      return this.#store[key]
     }
 
-    if (this.#deletedPrefix.some((dp) => storageKey.startsWith(dp))) {
+    if (this.#deletedPrefix.some((dp) => key.startsWith(dp))) {
       return StorageValueKind.Deleted
     }
 
     if (this.#parent) {
-      const val = this.#parent.get(key, false, child)
+      const val = this.#parent.get(key, false)
       if (cache) {
-        this.#store[storageKey] = val
+        this.#store[key] = val
       }
       return val
     }
@@ -208,12 +199,9 @@ export class StorageLayer implements StorageLayerProvider {
     }
   }
 
-  async getKeysPaged(prefix: string, pageSize: number, startKey: string, child?: string): Promise<string[]> {
-    const storagePrefix = mergeKey(child, prefix)
-    const storageStartKey = mergeKey(child, startKey)
-
-    if (!this.#deletedPrefix.some((dp) => storageStartKey.startsWith(dp))) {
-      const remote = (await this.#parent?.getKeysPaged(prefix, pageSize, startKey, child)) ?? []
+  async getKeysPaged(prefix: string, pageSize: number, startKey: string): Promise<string[]> {
+    if (!this.#deletedPrefix.some((dp) => startKey.startsWith(dp))) {
+      const remote = (await this.#parent?.getKeysPaged(prefix, pageSize, startKey)) ?? []
       for (const key of remote) {
         if (this.#deletedPrefix.some((dp) => key.startsWith(dp))) {
           continue
@@ -222,14 +210,14 @@ export class StorageLayer implements StorageLayerProvider {
       }
     }
 
-    let idx = _.sortedIndex(this.#keys, storageStartKey)
-    if (this.#keys[idx] === storageStartKey) {
+    let idx = _.sortedIndex(this.#keys, startKey)
+    if (this.#keys[idx] === startKey) {
       ++idx
     }
     const res: string[] = []
     while (res.length < pageSize) {
       const key: string = this.#keys[idx]
-      if (!key || !key.startsWith(storagePrefix)) {
+      if (!key || !key.startsWith(prefix)) {
         break
       }
       res.push(key)