From e016e11bff16bfe9fcd7dc4370613aa3ee09be87 Mon Sep 17 00:00:00 2001
From: Artem Vorotnikov <artem@vorotnikov.me>
Date: Fri, 24 Dec 2021 03:03:01 +0300
Subject: [PATCH] walk_dup

---
 src/kv/traits.rs          | 28 ++++++++++++++++++++++++++++
 src/stages/interhashes.rs | 12 +++++-------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/kv/traits.rs b/src/kv/traits.rs
index 225b1f0..f3f5510 100644
--- a/src/kv/traits.rs
+++ b/src/kv/traits.rs
@@ -188,6 +188,34 @@ where
     }
 }
 
+/// Walk over duplicates for some specific key.
+pub fn walk_dup<'tx: 'cur, 'cur, C, T>(
+    cursor: &'cur mut C,
+    start_key: T::Key,
+) -> impl Stream<Item = anyhow::Result<(T::Key, T::Value)>> + 'cur
+where
+    C: CursorDupSort<'tx, T>,
+    T: DupSort,
+    T::Key: TableDecode,
+    'tx: 'cur,
+{
+    try_stream! {
+        let start = cursor.seek_exact(start_key).await?;
+        if let Some(mut fv) = start {
+            loop {
+                yield fv;
+
+                match cursor.next_dup().await? {
+                    Some(fv1) => {
+                        fv = fv1;
+                    }
+                    None => break,
+                }
+            }
+        }
+    }
+}
+
 pub fn ttw<'a, T, E>(f: impl Fn(&T) -> bool + 'a) -> impl Fn(&Result<T, E>) -> bool + 'a {
     move |res| match res {
         Ok(v) => (f)(v),
diff --git a/src/stages/interhashes.rs b/src/stages/interhashes.rs
index ec657ea..3b8ba25 100644
--- a/src/stages/interhashes.rs
+++ b/src/stages/interhashes.rs
@@ -16,6 +16,7 @@ use async_trait::async_trait;
 use ethereum_types::*;
 use rlp::RlpStream;
 use std::{cmp, collections::BTreeMap};
+use tokio_stream::StreamExt;
 use tracing::*;
 
 fn hex_prefix(nibbles: &[u8], user_flag: bool) -> Vec<u8> {
@@ -476,13 +477,10 @@ where
         &mut self,
         address_hash: H256,
     ) -> anyhow::Result<Vec<(H256, U256)>> {
-        let mut storage = Vec::<(H256, U256)>::new();
-        let mut found = self.storage_cursor.seek_exact(address_hash).await?;
-        while let Some((_, storage_entry)) = found {
-            storage.push((storage_entry.0, storage_entry.1));
-            found = self.storage_cursor.next_dup().await?;
-        }
-        Ok(storage)
+        walk_dup(&mut self.storage_cursor, address_hash)
+            .map(|res| res.map(|(_, storage_entry)| storage_entry))
+            .collect()
+            .await
     }
 
     async fn visit_storage(&mut self, address_hash: H256) -> anyhow::Result<H256> {