Storage map iter (#148)

* Add example file

* Fmt

* Add KeyIter.

* Add iter method to store proc-macro.

* Fetch all values at once.

* Update docs.

* Run rustfmt.

Co-authored-by: Andrew Jones <ascjones@gmail.com>

src/lib.rs

@@ -55,6 +55,7 @@ use sc_rpc_api::state::ReadProof;
 use sp_core::{
     storage::{
         StorageChangeSet,
+        StorageData,
         StorageKey,
     },
     Bytes,
@@ -115,6 +116,7 @@ pub struct ClientBuilder<T: Runtime> {
     _marker: std::marker::PhantomData<T>,
     url: Option<String>,
     client: Option<jsonrpsee::Client>,
+    page_size: Option<u32>,
 }
 
 impl<T: Runtime> ClientBuilder<T> {
@@ -124,6 +126,7 @@ impl<T: Runtime> ClientBuilder<T> {
             _marker: std::marker::PhantomData,
             url: None,
             client: None,
+            page_size: None,
         }
     }
 
@@ -139,6 +142,12 @@ impl<T: Runtime> ClientBuilder<T> {
         self
     }
 
+    /// Set the page size.
+    pub fn set_page_size(mut self, size: u32) -> Self {
+        self.page_size = Some(size);
+        self
+    }
+
     /// Creates a new Client.
     pub async fn build(self) -> Result<Client<T>, Error> {
         let client = if let Some(client) = self.client {
@@ -164,6 +173,7 @@ impl<T: Runtime> ClientBuilder<T> {
             metadata: metadata?,
             runtime_version: runtime_version?,
             _marker: PhantomData,
+            page_size: self.page_size.unwrap_or(10),
         })
     }
 }
@@ -175,6 +185,7 @@ pub struct Client<T: Runtime> {
     metadata: Metadata,
     runtime_version: RuntimeVersion,
     _marker: PhantomData<(fn() -> T::Signature, T::Extra)>,
+    page_size: u32,
 }
 
 impl<T: Runtime> Clone for Client<T> {
@@ -185,6 +196,53 @@ impl<T: Runtime> Clone for Client<T> {
             metadata: self.metadata.clone(),
             runtime_version: self.runtime_version.clone(),
             _marker: PhantomData,
+            page_size: self.page_size,
+        }
+    }
+}
+
+/// Iterates over key value pairs in a map.
+pub struct KeyIter<T: Runtime, F: Store<T>> {
+    client: Client<T>,
+    _marker: PhantomData<F>,
+    count: u32,
+    hash: T::Hash,
+    start_key: Option<StorageKey>,
+    buffer: Vec<(StorageKey, StorageData)>,
+}
+
+impl<T: Runtime, F: Store<T>> KeyIter<T, F> {
+    /// Returns the next key value pair from a map.
+    pub async fn next(&mut self) -> Result<Option<(StorageKey, F::Returns)>, Error> {
+        loop {
+            if let Some((k, v)) = self.buffer.pop() {
+                return Ok(Some((k, Decode::decode(&mut &v.0[..])?)))
+            } else {
+                let keys = self
+                    .client
+                    .fetch_keys::<F>(self.count, self.start_key.take(), Some(self.hash))
+                    .await?;
+
+                if keys.is_empty() {
+                    return Ok(None)
+                }
+
+                self.start_key = keys.last().cloned();
+
+                let change_sets = self
+                    .client
+                    .rpc
+                    .query_storage_at(&keys, Some(self.hash))
+                    .await?;
+                for change_set in change_sets {
+                    for (k, v) in change_set.changes {
+                        if let Some(v) = v {
+                            self.buffer.push((k, v));
+                        }
+                    }
+                }
+                debug_assert_eq!(self.buffer.len(), self.count as usize);
+            }
         }
     }
 }
@@ -200,37 +258,58 @@ impl<T: Runtime> Client<T> {
         &self.metadata
     }
 
-    /// Fetch a StorageKey with default value.
-    pub async fn fetch_or_default<F: Store<T>>(
-        &self,
-        store: F,
-        hash: Option<T::Hash>,
-    ) -> Result<F::Returns, Error> {
-        let key = store.key(&self.metadata)?;
-        if let Some(data) = self.rpc.storage(key, hash).await? {
-            Ok(Decode::decode(&mut &data.0[..])?)
-        } else {
-            Ok(store.default(&self.metadata)?)
-        }
-    }
-
-    /// Fetch a StorageKey an optional storage key.
+    /// Fetch a StorageKey with an optional block hash.
     pub async fn fetch<F: Store<T>>(
         &self,
-        store: F,
+        store: &F,
         hash: Option<T::Hash>,
     ) -> Result<Option<F::Returns>, Error> {
         let key = store.key(&self.metadata)?;
-        if let Some(data) = self.rpc.storage(key, hash).await? {
+        if let Some(data) = self.rpc.storage(&key, hash).await? {
             Ok(Some(Decode::decode(&mut &data.0[..])?))
         } else {
             Ok(None)
         }
     }
 
-	/// Fetch up to `count` keys for a storage map in lexicographic order.
-	///
-	/// Supports pagination by passing a value to `start_key`.
+    /// Fetch a StorageKey that has a default value with an optional block hash.
+    pub async fn fetch_or_default<F: Store<T>>(
+        &self,
+        store: &F,
+        hash: Option<T::Hash>,
+    ) -> Result<F::Returns, Error> {
+        if let Some(data) = self.fetch(store, hash).await? {
+            Ok(data)
+        } else {
+            Ok(store.default(&self.metadata)?)
+        }
+    }
+
+    /// Returns an iterator of key value pairs.
+    pub async fn iter<F: Store<T>>(
+        &self,
+        hash: Option<T::Hash>,
+    ) -> Result<KeyIter<T, F>, Error> {
+        let hash = if let Some(hash) = hash {
+            hash
+        } else {
+            self.block_hash(None)
+                .await?
+                .expect("didn't pass a block number; qed")
+        };
+        Ok(KeyIter {
+            client: self.clone(),
+            hash,
+            count: self.page_size,
+            start_key: None,
+            buffer: Default::default(),
+            _marker: PhantomData,
+        })
+    }
+
+    /// Fetch up to `count` keys for a storage map in lexicographic order.
+    ///
+    /// Supports pagination by passing a value to `start_key`.
     pub async fn fetch_keys<F: Store<T>>(
         &self,
         count: u32,
@@ -238,7 +317,10 @@ impl<T: Runtime> Client<T> {
         hash: Option<T::Hash>,
     ) -> Result<Vec<StorageKey>, Error> {
         let prefix = <F as Store<T>>::prefix(&self.metadata)?;
-        let keys = self.rpc.storage_keys_paged(Some(prefix), count, start_key, hash).await?;
+        let keys = self
+            .rpc
+            .storage_keys_paged(Some(prefix), count, start_key, hash)
+            .await?;
         Ok(keys)
     }
 
@@ -513,6 +595,7 @@ mod tests {
                 SubxtClient::from_config(config, test_node::service::new_full)
                     .expect("Error creating subxt client"),
             )
+            .set_page_size(2)
             .build()
             .await
             .expect("Error creating client");
@@ -630,7 +713,21 @@ mod tests {
     #[async_std::test]
     async fn test_fetch_keys() {
         let (client, _) = test_client().await;
-        let keys = client.fetch_keys::<system::AccountStore<_>>(4, None, None).await.unwrap();
+        let keys = client
+            .fetch_keys::<system::AccountStore<_>>(4, None, None)
+            .await
+            .unwrap();
         assert_eq!(keys.len(), 4)
     }
+
+    #[async_std::test]
+    async fn test_iter() {
+        let (client, _) = test_client().await;
+        let mut iter = client.iter::<system::AccountStore<_>>(None).await.unwrap();
+        let mut i = 0;
+        while let Some(_) = iter.next().await.unwrap() {
+            i += 1;
+        }
+        assert_eq!(i, 4);
+    }
 }

src/rpc.rs

@@ -122,7 +122,7 @@ impl<T: Runtime> Rpc<T> {
     /// Fetch a storage key
     pub async fn storage(
         &self,
-        key: StorageKey,
+        key: &StorageKey,
         hash: Option<T::Hash>,
     ) -> Result<Option<StorageData>, Error> {
         let params = Params::Array(vec![to_json_value(key)?, to_json_value(hash)?]);
@@ -132,8 +132,8 @@ impl<T: Runtime> Rpc<T> {
     }
 
     /// Returns the keys with prefix with pagination support.
-	/// Up to `count` keys will be returned.
-	/// If `start_key` is passed, return next keys in storage in lexicographic order.
+    /// Up to `count` keys will be returned.
+    /// If `start_key` is passed, return next keys in storage in lexicographic order.
     pub async fn storage_keys_paged(
         &self,
         prefix: Option<StorageKey>,
@@ -170,6 +170,19 @@ impl<T: Runtime> Rpc<T> {
             .map_err(Into::into)
     }
 
+    /// Query historical storage entries
+    pub async fn query_storage_at(
+        &self,
+        keys: &[StorageKey],
+        at: Option<T::Hash>,
+    ) -> Result<Vec<StorageChangeSet<<T as System>::Hash>>, Error> {
+        let params = Params::Array(vec![to_json_value(keys)?, to_json_value(at)?]);
+        self.client
+            .request("state_queryStorage", params)
+            .await
+            .map_err(Into::into)
+    }
+
     /// Fetch the genesis hash
     pub async fn genesis_hash(&self) -> Result<T::Hash, Error> {
         let block_zero = Some(ListOrValue::Value(NumberOrHex::Number(0)));