fix(engine): get_proof_targets only add fetched accounts if they have new storage

crates/engine/tree/src/tree/root.rs

@@ -464,31 +464,37 @@ where
     }
 }
 
+/// Returns accounts only with those storages that were not already fetched, and
+/// if there are no such storages and the account itself was already fetched, the
+/// account shouldn't be included.
 fn get_proof_targets(
     state_update: &HashedPostState,
     fetched_proof_targets: &HashMap<B256, HashSet<B256>>,
 ) -> HashMap<B256, HashSet<B256>> {
-    state_update
-        .accounts
-        .keys()
-        .filter(|hashed_address| !fetched_proof_targets.contains_key(*hashed_address))
-        .map(|hashed_address| (*hashed_address, HashSet::default()))
-        .chain(state_update.storages.iter().map(|(hashed_address, storage)| {
-            let fetched_storage_proof_targets = fetched_proof_targets.get(hashed_address);
-            (
-                *hashed_address,
-                storage
-                    .storage
-                    .keys()
-                    .filter(|slot| {
-                        !fetched_storage_proof_targets
-                            .is_some_and(|targets| targets.contains(*slot))
-                    })
-                    .copied()
-                    .collect(),
-            )
-        }))
-        .collect()
+    let mut targets = HashMap::default();
+
+    // first collect all new accounts (not previously fetched)
+    for &hashed_address in state_update.accounts.keys() {
+        if !fetched_proof_targets.contains_key(&hashed_address) {
+            targets.insert(hashed_address, HashSet::default());
+        }
+    }
+
+    // then process storage slots for all accounts in the state update
+    for (hashed_address, storage) in &state_update.storages {
+        let fetched = fetched_proof_targets.get(hashed_address);
+        let mut changed_slots = storage
+            .storage
+            .keys()
+            .filter(|slot| !fetched.is_some_and(|f| f.contains(*slot)))
+            .peekable();
+
+        if changed_slots.peek().is_some() {
+            targets.entry(*hashed_address).or_default().extend(changed_slots);
+        }
+    }
+
+    targets
 }
 
 /// Updates the sparse trie with the given proofs and state, and returns the updated trie and the
@@ -793,4 +799,166 @@ mod tests {
         assert_eq!(ready.len(), 5);
         assert!(!sequencer.has_pending());
     }
+
+    fn create_get_proof_targets_state() -> HashedPostState {
+        let mut state = HashedPostState::default();
+
+        let addr1 = B256::random();
+        let addr2 = B256::random();
+        state.accounts.insert(addr1, Some(Default::default()));
+        state.accounts.insert(addr2, Some(Default::default()));
+
+        let mut storage = HashedStorage::default();
+        let slot1 = B256::random();
+        let slot2 = B256::random();
+        storage.storage.insert(slot1, U256::ZERO);
+        storage.storage.insert(slot2, U256::from(1));
+        state.storages.insert(addr1, storage);
+
+        state
+    }
+
+    #[test]
+    fn test_get_proof_targets_new_account_targets() {
+        let state = create_get_proof_targets_state();
+        let fetched = HashMap::default();
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        // should return all accounts as targets since nothing was fetched before
+        assert_eq!(targets.len(), state.accounts.len());
+        for addr in state.accounts.keys() {
+            assert!(targets.contains_key(addr));
+        }
+    }
+
+    #[test]
+    fn test_get_proof_targets_new_storage_targets() {
+        let state = create_get_proof_targets_state();
+        let fetched = HashMap::default();
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        // verify storage slots are included for accounts with storage
+        for (addr, storage) in &state.storages {
+            assert!(targets.contains_key(addr));
+            let target_slots = &targets[addr];
+            assert_eq!(target_slots.len(), storage.storage.len());
+            for slot in storage.storage.keys() {
+                assert!(target_slots.contains(slot));
+            }
+        }
+    }
+
+    #[test]
+    fn test_get_proof_targets_filter_already_fetched_accounts() {
+        let state = create_get_proof_targets_state();
+        let mut fetched = HashMap::default();
+
+        // select an account that has no storage updates
+        let fetched_addr = state
+            .accounts
+            .keys()
+            .find(|&&addr| !state.storages.contains_key(&addr))
+            .expect("Should have an account without storage");
+
+        // mark the account as already fetched
+        fetched.insert(*fetched_addr, HashSet::default());
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        // should not include the already fetched account since it has no storage updates
+        assert!(!targets.contains_key(fetched_addr));
+        // other accounts should still be included
+        assert_eq!(targets.len(), state.accounts.len() - 1);
+    }
+
+    #[test]
+    fn test_get_proof_targets_filter_already_fetched_storage() {
+        let state = create_get_proof_targets_state();
+        let mut fetched = HashMap::default();
+
+        // mark one storage slot as already fetched
+        let (addr, storage) = state.storages.iter().next().unwrap();
+        let mut fetched_slots = HashSet::default();
+        let fetched_slot = *storage.storage.keys().next().unwrap();
+        fetched_slots.insert(fetched_slot);
+        fetched.insert(*addr, fetched_slots);
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        // should not include the already fetched storage slot
+        let target_slots = &targets[addr];
+        assert!(!target_slots.contains(&fetched_slot));
+        assert_eq!(target_slots.len(), storage.storage.len() - 1);
+    }
+
+    #[test]
+    fn test_get_proof_targets_empty_state() {
+        let state = HashedPostState::default();
+        let fetched = HashMap::default();
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        assert!(targets.is_empty());
+    }
+
+    #[test]
+    fn test_get_proof_targets_mixed_fetched_state() {
+        let mut state = HashedPostState::default();
+        let mut fetched = HashMap::default();
+
+        let addr1 = B256::random();
+        let addr2 = B256::random();
+        let slot1 = B256::random();
+        let slot2 = B256::random();
+
+        state.accounts.insert(addr1, Some(Default::default()));
+        state.accounts.insert(addr2, Some(Default::default()));
+
+        let mut storage = HashedStorage::default();
+        storage.storage.insert(slot1, U256::ZERO);
+        storage.storage.insert(slot2, U256::from(1));
+        state.storages.insert(addr1, storage);
+
+        let mut fetched_slots = HashSet::default();
+        fetched_slots.insert(slot1);
+        fetched.insert(addr1, fetched_slots);
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        assert!(targets.contains_key(&addr2));
+        assert!(!targets[&addr1].contains(&slot1));
+        assert!(targets[&addr1].contains(&slot2));
+    }
+
+    #[test]
+    fn test_get_proof_targets_unmodified_account_with_storage() {
+        let mut state = HashedPostState::default();
+        let fetched = HashMap::default();
+
+        let addr = B256::random();
+        let slot1 = B256::random();
+        let slot2 = B256::random();
+
+        // don't add the account to state.accounts (simulating unmodified account)
+        // but add storage updates for this account
+        let mut storage = HashedStorage::default();
+        storage.storage.insert(slot1, U256::from(1));
+        storage.storage.insert(slot2, U256::from(2));
+        state.storages.insert(addr, storage);
+
+        assert!(!state.accounts.contains_key(&addr));
+        assert!(!fetched.contains_key(&addr));
+
+        let targets = get_proof_targets(&state, &fetched);
+
+        // verify that we still get the storage slots for the unmodified account
+        assert!(targets.contains_key(&addr));
+
+        let target_slots = &targets[&addr];
+        assert_eq!(target_slots.len(), 2);
+        assert!(target_slots.contains(&slot1));
+        assert!(target_slots.contains(&slot2));
+    }
 }