allow hashed prefix to be configured. (#33)

* allow hashed prefix to be configured. * Update core/src/state.rs * Update core/src/state.rs Co-authored-by: Liam Aharon <[email protected]> * Update core/src/state.rs Co-authored-by: Liam Aharon <[email protected]> * Update core/src/state.rs Co-authored-by: Liam Aharon <[email protected]> --------- Co-authored-by: Liam Aharon <[email protected]>
paritytech · Sep 25, 2023 · 22c7571 · 22c7571
1 parent 533a122
commit 22c7571
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/core/src/commands/follow_chain.rs b/core/src/commands/follow_chain.rs
@@ -140,6 +140,7 @@ where
                 at: Some(hex::encode(header.parent_hash().encode())),
                 pallet: vec![],
                 child_tree: true,
+                hashed_prefixes: vec![],
             });
             let ext = state
                 .to_ext::<Block, HostFns>(&shared, &executor, None, true)

diff --git a/core/src/state.rs b/core/src/state.rs
@@ -61,9 +61,15 @@ pub struct LiveState {
 
     /// A pallet to scrape. Can be provided multiple times. If empty, entire chain state will
     /// be scraped.
+    ///
+    /// This is equivalent to passing `xx_hash_64(pallet)` to `--hashed_prefixes`.
     #[arg(short, long, num_args = 1..)]
     pub pallet: Vec<String>,
 
+    /// Storage entry key prefixes to scrape and inject into the test externalities. Pass as 0x prefixed hex strings. By default, all keys are scraped and included.
+    #[arg(long = "prefix", value_parser = parse::hash, num_args = 1..)]
+    pub hashed_prefixes: Vec<String>,
+
     /// Fetch the child-keys as well.
     ///
     /// Default is `false`, if specific `--pallets` are specified, `true` otherwise. In other
@@ -132,11 +138,23 @@ impl State {
                 uri,
                 at,
                 child_tree,
+                hashed_prefixes,
             }) => {
                 let at = match at {
                     Some(at_str) => Some(hash_of::<Block>(at_str)?),
                     None => None,
                 };
+                let hashed_prefixes = hashed_prefixes
+                    .iter()
+                    .map(|p_str| {
+                        hex::decode(p_str).map_err(|e| {
+                            format!(
+                                "Error decoding `hashed_prefixes` hex string entry '{:?}' to bytes: {:?}",
+                                p_str, e
+                            )
+                        })
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
                 Builder::<Block>::new().mode(Mode::Online(OnlineConfig {
                     at,
                     transport: uri.to_owned().into(),
@@ -152,7 +170,7 @@ impl State {
                         [twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat(),
                         [twox_128(b"System"), twox_128(b"Number")].concat(),
                     ],
-                    hashed_prefixes: vec![],
+                    hashed_prefixes,
                 }))
             }
         };