From 86eeaff7f38dcb9da74be5d00f595f1cd46e34d3 Mon Sep 17 00:00:00 2001
From: Joonatan Saarhelo <joon.saar@gmail.com>
Date: Fri, 24 May 2024 19:01:34 +0200
Subject: [PATCH] return iterators

---
 src/world_diff.rs | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/world_diff.rs b/src/world_diff.rs
index a782d0ce..0fadbb5d 100644
--- a/src/world_diff.rs
+++ b/src/world_diff.rs
@@ -139,20 +139,25 @@ impl WorldDiff {
         self.storage_changes.as_ref()
     }
 
-    pub fn get_storage_changes(&self) -> BTreeMap<(H160, U256), (u16, Option<U256>, U256)> {
-        let mut result = BTreeMap::new();
-        for (key, &(tx_number, value)) in self.storage_changes.as_ref() {
-            if self.storage_initial_values[key] != Some(value) {
-                result.insert(*key, (tx_number, self.storage_initial_values[key], value));
-            }
-        }
-        result
+    pub fn get_storage_changes(
+        &self,
+    ) -> impl Iterator<Item = ((H160, U256), (u16, Option<U256>, U256))> + '_ {
+        self.storage_changes
+            .as_ref()
+            .iter()
+            .filter_map(|(key, &(tx_number, value))| {
+                if self.storage_initial_values[key] == Some(value) {
+                    None
+                } else {
+                    Some((*key, (tx_number, self.storage_initial_values[key], value)))
+                }
+            })
     }
 
     pub fn get_storage_changes_after(
         &self,
         snapshot: &Snapshot,
-    ) -> BTreeMap<(H160, U256), StorageChange> {
+    ) -> impl Iterator<Item = ((H160, U256), StorageChange)> + '_ {
         self.storage_changes
             .changes_after(snapshot.storage_changes)
             .into_iter()
@@ -168,7 +173,6 @@ impl WorldDiff {
                     },
                 )
             })
-            .collect()
     }
 
     pub(crate) fn read_transient_storage(&mut self, contract: H160, key: U256) -> U256 {