feat: add native balance in StorageChanges (#292)

* feat: add native balance in StorageChanges

This will be used by the next update of DCI to properly track native balance for all the contract that it's indexing. Also added a performance optimization, we now ignore slot updates if the value after the transaction is the same than before.

* docs: add docs on how to generate test assets

---------

Co-authored-by: zizou <111426680+flopell@users.noreply.github.com>

proto/tycho/evm/v1/common.proto

@@ -192,6 +192,8 @@ message StorageChanges {
   bytes address = 1;
   // The contract's storage changes
   repeated ContractSlot slots = 2;
+  // [optional] The contract's balance change
+  optional bytes native_balance = 3;
 }
 
 // Aggregate entities

substreams/Cargo.lock

@@ -35,6 +35,12 @@ version = "0.13.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
 
+[[package]]
+name = "base64"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
+
 [[package]]
 name = "bigdecimal"
 version = "0.3.1"
@@ -250,7 +256,7 @@ dependencies = [
  "num-bigint",
  "substreams",
  "substreams-ethereum",
- "tycho-substreams 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "tycho-substreams 0.5.1",
 ]
 
 [[package]]
@@ -1415,7 +1421,7 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e2c7fca123abff659d15ed30da5b605fa954a29e912c94260c488d0d18f9107d"
 dependencies = [
- "base64",
+ "base64 0.13.1",
  "prost 0.11.9",
  "prost-types 0.11.9",
  "substreams",
@@ -1489,7 +1495,7 @@ name = "substreams-helper"
 version = "0.0.2"
 dependencies = [
  "anyhow",
- "base64",
+ "base64 0.13.1",
  "bigdecimal",
  "downcast-rs",
  "ethabi 18.0.0",
@@ -1512,7 +1518,7 @@ version = "0.0.2"
 source = "git+https://github.com/propeller-heads/tycho-protocol-sdk.git?tag=0.4.0#cfbf6812bdc9503ff51debcf5e171cd680b4d694"
 dependencies = [
  "anyhow",
- "base64",
+ "base64 0.13.1",
  "bigdecimal",
  "downcast-rs",
  "ethabi 18.0.0",
@@ -1535,7 +1541,7 @@ version = "0.0.2"
 source = "git+https://github.com/propeller-heads/tycho-protocol-sdk.git?rev=52d5021#52d502198e9aa964814ef5f139df0886c3eb7bb0"
 dependencies = [
  "anyhow",
- "base64",
+ "base64 0.13.1",
  "bigdecimal",
  "downcast-rs",
  "ethabi 18.0.0",
@@ -1558,7 +1564,7 @@ version = "0.0.2"
 source = "git+https://github.com/propeller-heads/tycho-protocol-sdk.git?rev=b8aeaa3#b8aeaa3dc6e7242a5dd23681921258ef2cb3c6dd"
 dependencies = [
  "anyhow",
- "base64",
+ "base64 0.13.1",
  "bigdecimal",
  "downcast-rs",
  "ethabi 18.0.0",
@@ -1758,22 +1764,6 @@ dependencies = [
  "substreams-ethereum",
 ]
 
-[[package]]
-name = "tycho-substreams"
-version = "0.5.1"
-dependencies = [
- "ethabi 18.0.0",
- "hex",
- "itertools 0.12.1",
- "num-bigint",
- "prost 0.11.9",
- "rstest",
- "serde",
- "serde_json",
- "substreams",
- "substreams-ethereum",
-]
-
 [[package]]
 name = "tycho-substreams"
 version = "0.5.1"
@@ -1791,6 +1781,23 @@ dependencies = [
  "substreams-ethereum",
 ]
 
+[[package]]
+name = "tycho-substreams"
+version = "0.6.0"
+dependencies = [
+ "base64 0.22.1",
+ "ethabi 18.0.0",
+ "hex",
+ "itertools 0.12.1",
+ "num-bigint",
+ "prost 0.11.9",
+ "rstest",
+ "serde",
+ "serde_json",
+ "substreams",
+ "substreams-ethereum",
+]
+
 [[package]]
 name = "typenum"
 version = "1.17.0"

substreams/crates/tycho-substreams/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tycho-substreams"
-version = "0.5.1"
+version = "0.6.0"
 edition = "2021"
 description = "Tycho substreams development kit, contains tycho-indexer block changes model and helper functions for common indexing tasks."
 repository = "https://github.com/propeller-heads/tycho-protocol-sdk/tree/main/substreams/crates/tycho-substreams"
@@ -24,3 +24,4 @@ serde_json = "1.0.120"
 
 [dev-dependencies]
 rstest = "0.24.0"
+base64 = "0.22.1"

substreams/crates/tycho-substreams/Readme.md

@@ -10,3 +10,7 @@ directory:
 ```bash
 buf generate --template substreams/crates/tycho-substreams/buf.gen.yaml --output substreams/crates/tycho-substreams/
 ```
+
+## Generate block test assets
+
+To be able to write complete unit tests, we rely on full block assets. These assets can be generated using the firecore tool from Substreams. More info in [Substreams documentation](https://docs.substreams.dev/reference-material/log-and-debug#generating-the-input-of-the-test)

substreams/crates/tycho-substreams/src/balances.rs

@@ -248,7 +248,7 @@ pub fn extract_balance_deltas_from_tx<F: Fn(&[u8], &[u8]) -> bool>(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{mock_store::MockStore, pb::tycho::evm::v1::BalanceDelta};
+    use crate::{pb::tycho::evm::v1::BalanceDelta, testing::mock_store::MockStore};
     use substreams::{
         pb::substreams::StoreDelta,
         prelude::{StoreGet, StoreNew},

substreams/crates/tycho-substreams/src/block_storage.rs

@@ -1,12 +1,12 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
-use substreams_ethereum::pb::{
+use substreams_ethereum::pb::eth::{
-    eth,
+    self,
-    eth::v2::{block::DetailLevel, StorageChange},
+    v2::{block::DetailLevel, BalanceChange, StorageChange},
 };
 
 use crate::{
-    models::{ContractSlot, StorageChanges, Transaction},
+    models::{ContractSlot, StorageChanges},
     pb::tycho::evm::v1::TransactionStorageChanges,
 };
 
@@ -30,8 +30,6 @@ pub fn get_block_storage_changes(block: &eth::v2::Block) -> Vec<TransactionStora
     let mut block_storage_changes = Vec::with_capacity(block.transaction_traces.len());
 
     for block_tx in block.transactions() {
-        let transaction: Transaction = block_tx.into();
-
         let mut changes_by_address: HashMap<Vec<u8>, Vec<StorageChange>> = HashMap::new();
         for storage_change in block_tx
             .calls
@@ -45,38 +43,270 @@ pub fn get_block_storage_changes(block: &eth::v2::Block) -> Vec<TransactionStora
                 .push(storage_change.clone());
         }
 
-        // For each address, sort by ordinal and collect latest changes per slot
+        let mut native_balance_changes_by_address: HashMap<Vec<u8>, Vec<BalanceChange>> =
-        let tx_storage_changes: Vec<StorageChanges> = changes_by_address
+            HashMap::new();
+        for balance_change in block_tx
+            .calls
+            .iter()
+            .filter(|call| !call.state_reverted)
+            .flat_map(|call| call.balance_changes.iter())
+        {
+            native_balance_changes_by_address
+                .entry(balance_change.address.clone())
+                .or_default()
+                .push(balance_change.clone());
+        }
+
+        // Collect all unique addresses from both storage changes and balance changes
+        let mut all_addresses = HashSet::new();
+        all_addresses.extend(changes_by_address.keys().cloned());
+        all_addresses.extend(
+            native_balance_changes_by_address
+                .keys()
+                .cloned(),
+        );
+
+        // For each address, collect both storage changes and balance changes
+        let tx_storage_changes: Vec<StorageChanges> = all_addresses
             .into_iter()
-            .map(|(address, mut changes)| {
+            .map(|address| {
-                changes.sort_unstable_by_key(|change| change.ordinal);
+                // Process storage changes for this address
+                let slots = if let Some(changes) = changes_by_address.get(&address) {
+                    let mut changes = changes.clone();
+                    changes.sort_unstable_by_key(|change| change.ordinal);
 
-                // Collect latest change per slot
+                    // Collect latest change per slot
-                let mut latest_changes: HashMap<Vec<u8>, ContractSlot> = HashMap::new();
+                    let mut latest_changes: HashMap<Vec<u8>, ContractSlot> = HashMap::new();
-                for change in changes {
+                    for change in changes {
-                    latest_changes
+                        latest_changes
-                        .entry(change.key.clone())
+                            .entry(change.key.clone())
-                        .and_modify(|slot| {
+                            .and_modify(|slot| {
-                            // Only update the latest value, previous value stays the first seen
+                                // Only update the latest value, previous value stays the first seen
-                            // one.
+                                // one.
-                            slot.value = change.new_value.clone();
+                                slot.value = change.new_value.clone();
-                        })
+                            })
-                        .or_insert(ContractSlot {
+                            .or_insert(ContractSlot {
-                            slot: change.key,
+                                slot: change.key,
-                            value: change.new_value,
+                                value: change.new_value,
-                            previous_value: change.old_value,
+                                previous_value: change.old_value,
+                            });
+                    }
+                    latest_changes.into_values().collect()
+                } else {
+                    vec![]
+                };
+
+                // Filter out slots that have the same value before and after the transaction
+                let slots = slots
+                    .into_iter()
+                    .filter(|slot| slot.previous_value != slot.value)
+                    .collect();
+
+                // Process native balance changes for this address
+                let native_balance = native_balance_changes_by_address
+                    .get(&address)
+                    .and_then(|balance_changes| {
+                        let (first, last) = balance_changes.iter().fold(
+                            (None, None),
+                            |(min, max): (Option<&BalanceChange>, Option<&BalanceChange>),
+                             change| {
+                                let new_min = match min {
+                                    None => Some(change),
+                                    Some(m) if change.ordinal < m.ordinal => Some(change),
+                                    _ => min,
+                                };
+                                let new_max = match max {
+                                    None => Some(change),
+                                    Some(m) if change.ordinal > m.ordinal => Some(change),
+                                    _ => max,
+                                };
+                                (new_min, new_max)
+                            },
+                        );
+
+                        let balance_before_tx = first.map(|f| {
+                            f.old_value
+                                .as_ref()
+                                .map(|b| b.bytes.clone())
+                                .unwrap_or_default()
+                        });
+                        let balance_after_tx = last.map(|l| {
+                            l.new_value
+                                .as_ref()
+                                .map(|b| b.bytes.clone())
+                                .unwrap_or_default()
                         });
-                }
 
-                StorageChanges { address, slots: latest_changes.into_values().collect() }
+                        (balance_before_tx != balance_after_tx).then_some(balance_after_tx.clone())
+                    })
+                    .flatten();
+
+                StorageChanges { address, slots, native_balance }
             })
             .collect();
 
         block_storage_changes.push(TransactionStorageChanges {
-            tx: Some(transaction),
+            tx: Some(block_tx.into()),
             storage_changes: tx_storage_changes,
         });
     }
 
     block_storage_changes
 }
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+    use crate::testing::assets::read_block;
+
+    #[test]
+    fn test_get_block_storage_changes_ethereum_block_23490768() {
+        let block = read_block("./assets/ethereum-block-23490768.binpb.base64");
+        let changes = get_block_storage_changes(&block);
+
+        let mut balance_map: HashMap<String, HashMap<String, String>> = HashMap::new();
+        let mut storage_map: HashMap<String, HashMap<String, HashMap<String, (String, String)>>> =
+            HashMap::new();
+        for change in changes {
+            let tx_hash = change.tx.unwrap().hash.clone();
+            let balance_tx_entry = balance_map
+                .entry(hex::encode(tx_hash.clone()))
+                .or_default();
+            let storage_tx_entry = storage_map
+                .entry(hex::encode(tx_hash.clone()))
+                .or_default();
+            for storage_change in change.storage_changes {
+                if let Some(native_balance) = storage_change.native_balance {
+                    balance_tx_entry.insert(
+                        hex::encode(storage_change.address.clone()),
+                        hex::encode(native_balance.clone()),
+                    );
+                }
+                for slot in storage_change.slots {
+                    let contract_tx_entry = storage_tx_entry
+                        .entry(hex::encode(storage_change.address.clone()))
+                        .or_default();
+                    contract_tx_entry.insert(
+                        hex::encode(slot.slot.clone()),
+                        (hex::encode(slot.previous_value.clone()), hex::encode(slot.value.clone())),
+                    );
+                }
+            }
+        }
+
+        // Assertions for https://etherscan.io/tx/0x44a34ba7400fa7004ec5037aeb1103a7c0cd8a83a95c4cd5cf9561c3c38db326#statechange
+        // Check balance changes
+        let balance_tx_entry = balance_map
+            .get("44a34ba7400fa7004ec5037aeb1103a7c0cd8a83a95c4cd5cf9561c3c38db326")
+            .unwrap();
+
+        assert_eq!(balance_tx_entry.len(), 4);
+        assert_eq!(
+            balance_tx_entry
+                .get("dadb0d80178819f2319190d340ce9a924f783711")
+                .unwrap(),
+            "052196f442fadb8314"
+        );
+
+        assert_eq!(
+            balance_tx_entry
+                .get("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2")
+                .unwrap(),
+            "0207150b274902c5e7871c"
+        );
+        assert_eq!(
+            balance_tx_entry
+                .get("ad01c20d5886137e056775af56915de824c8fce5")
+                .unwrap(),
+            "c83a1d6287cb5e"
+        );
+        assert_eq!(
+            balance_tx_entry
+                .get("638f1db9881a84af9835c6625d17b0af034234ad")
+                .unwrap(),
+            "0f69303da21468"
+        );
+
+        // Check storage changes
+        let storage_tx_entry = storage_map
+            .get("44a34ba7400fa7004ec5037aeb1103a7c0cd8a83a95c4cd5cf9561c3c38db326")
+            .unwrap();
+
+        assert_eq!(storage_tx_entry.len(), 3);
+
+        let storage_tx_entry_0f9e3401a5155a02c86353c3d9b24214876779dd = HashMap::from([
+            (
+                "0000000000000000000000000000000000000000000000000000000000000009".to_string(),
+                (
+                    "00000000000000000000000000000000009faeae5180599c05015fcfa242d3b0".to_string(),
+                    "00000000000000000000000000000000009faebb96f403f1913f425b3ea446e0".to_string(),
+                ),
+            ),
+            (
+                "000000000000000000000000000000000000000000000000000000000000000a".to_string(),
+                (
+                    "00000000000000000000000000f94f053f65617829584571d9de584cd219fb88".to_string(),
+                    "00000000000000000000000000f94f66e6e9d8f6688d6ca53ff9baae52e11cd8".to_string(),
+                ),
+            ),
+            (
+                "0000000000000000000000000000000000000000000000000000000000000008".to_string(),
+                (
+                    "68de8f37000000000001fb7a6a5bb2b548080000000560989aab8af59d9be89b".to_string(),
+                    "68de8f5b000000000001fb8b2909997ca55100000005606b52e81f19442026af".to_string(),
+                ),
+            ),
+        ]);
+        assert_eq!(
+            storage_tx_entry
+                .get("0f9e3401a5155a02c86353c3d9b24214876779dd")
+                .unwrap(),
+            &storage_tx_entry_0f9e3401a5155a02c86353c3d9b24214876779dd
+        );
+
+        let storage_tx_entry_11dfc652eb62c723ad8c2ae731fcede58ab07564 = HashMap::from([
+            (
+                "654f44e59f538551b5124259a61eaadb863c6c10cc9d43aa550237a76a7de0b0".to_string(),
+                (
+                    "000000000000000000000000000000000000000000000077c1c5e25db942af6a".to_string(),
+                    "0000000000000000000000000000000000000000000000a2c5f2bc08a7dea7a4".to_string(),
+                ),
+            ),
+            (
+                "6b12653da4ae5b17258ea9b02a62123c9305455af47b7dceea1b7137f7c69671".to_string(),
+                (
+                    "0000000000000000000000000000000000000000000001454f7d5d0ce8d4a21e".to_string(),
+                    "0000000000000000000000000000000000000000000001479313ef3e53b46bd0".to_string(),
+                ),
+            ),
+            (
+                "8f60e36f69a92730149f231ad2475b4aa8a8e50f4072f62a1f099ffc11d0f647".to_string(),
+                (
+                    "0000000000000000000000000000000000000000000560989aab8af59d9be89b".to_string(),
+                    "00000000000000000000000000000000000000000005606b52e81f19442026af".to_string(),
+                ),
+            ),
+        ]);
+        assert_eq!(
+            storage_tx_entry
+                .get("11dfc652eb62c723ad8c2ae731fcede58ab07564")
+                .unwrap(),
+            &storage_tx_entry_11dfc652eb62c723ad8c2ae731fcede58ab07564
+        );
+
+        let storage_tx_entry_c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2 = HashMap::from([(
+            "77f05379c72cc19907ba9648dcd0bda409fabc68ca111b532de62ffdb67e868f".to_string(),
+            (
+                "000000000000000000000000000000000000000000000001fb7a6a5bb2b54808".to_string(),
+                "000000000000000000000000000000000000000000000001fb8b2909997ca551".to_string(),
+            ),
+        )]);
+        assert_eq!(
+            storage_tx_entry
+                .get("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2")
+                .unwrap(),
+            &storage_tx_entry_c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2
+        );
+    }
+}

substreams/crates/tycho-substreams/src/lib.rs

@@ -4,11 +4,12 @@ pub mod balances;
 pub mod block_storage;
 pub mod contract;
 pub mod entrypoint;
-#[cfg(test)]
-mod mock_store;
 pub mod models;
 pub mod pb;
 
+#[cfg(test)]
+pub mod testing;
+
 pub mod prelude {
     pub use super::models::*;
 }

substreams/crates/tycho-substreams/src/pb/tycho.evm.v1.rs

@@ -245,6 +245,9 @@ pub struct StorageChanges {
     /// The contract's storage changes
     #[prost(message, repeated, tag="2")]
     pub slots: ::prost::alloc::vec::Vec<ContractSlot>,
+    /// \[optional\] The contract's balance change
+    #[prost(bytes="vec", optional, tag="3")]
+    pub native_balance: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
 }
 // Aggregate entities
 

substreams/crates/tycho-substreams/src/testing/assets.rs

@@ -0,0 +1,12 @@
+// Read a base64 encoded asset and return a decoded protobuf struct
+// Panics if the file does not exist or the base64 decoding fails
+pub fn read_block<B: prost::Message + Default>(filename: &str) -> B {
+    use base64::Engine;
+
+    let encoded = std::fs::read_to_string(filename).expect("Failed to read file");
+    let raw_bytes = base64::prelude::BASE64_STANDARD
+        .decode(&encoded)
+        .expect("Failed to decode base64");
+
+    B::decode(&*raw_bytes).expect("Not able to decode Block")
+}

substreams/crates/tycho-substreams/src/testing/mod.rs

@@ -0,0 +1,2 @@
+pub mod assets;
+pub mod mock_store;