Add prelude; Move pb implementations.

Else the implementations would get overwritten.

crates/tycho-substreams/src/lib.rs

@@ -1,4 +1,9 @@
 pub mod balances;
 pub mod contract;
 mod mock_store;
-pub mod pb;
+pub mod models;
+mod pb;
+
+pub mod prelude {
+    pub use super::models::*;
+}

crates/tycho-substreams/src/models.rs

@@ -0,0 +1,123 @@
+use substreams_ethereum::pb::eth::v2::{self as sf};
+
+// re-export the protobuf types here.
+pub use crate::pb::tycho::evm::v1::*;
+
+impl TransactionContractChanges {
+    /// Creates a new empty `TransactionContractChanges` instance.
+    pub fn new(tx: &Transaction) -> Self {
+        Self {
+            tx: Some(tx.clone()),
+            contract_changes: vec![],
+            component_changes: vec![],
+            balance_changes: vec![],
+        }
+    }
+}
+
+impl From<&sf::TransactionTrace> for Transaction {
+    fn from(tx: &sf::TransactionTrace) -> Self {
+        Self {
+            hash: tx.hash.clone(),
+            from: tx.from.clone(),
+            to: tx.to.clone(),
+            index: tx.index.into(),
+        }
+    }
+}
+
+impl From<&sf::Block> for Block {
+    fn from(block: &sf::Block) -> Self {
+        Self {
+            number: block.number,
+            hash: block.hash.clone(),
+            parent_hash: block
+                .header
+                .as_ref()
+                .expect("Block header not present")
+                .parent_hash
+                .clone(),
+            ts: block.timestamp_seconds(),
+        }
+    }
+}
+
+impl ProtocolComponent {
+    /// Creates a new empty `ProtocolComponent` instance.
+    ///
+    /// You can use the `with_*` methods to set the fields in a convience way.
+    pub fn new(id: &str, tx: &Transaction) -> Self {
+        Self {
+            id: id.to_string(),
+            tokens: vec![],
+            contracts: vec![],
+            static_att: vec![],
+            change: ChangeType::Creation.into(),
+            protocol_type: None,
+            tx: Some(tx.clone()),
+        }
+    }
+
+    /// Shorthand to create a component with a 1-1 relationship to a contract.
+    ///
+    /// Will set the component id to a hex encoded address with a 0x prefix
+    /// and add the contract to contracts attributes.
+    pub fn at_contract(id: &[u8], tx: &Transaction) -> Self {
+        Self {
+            id: format!("0x{}", hex::encode(id)),
+            tokens: vec![],
+            contracts: vec![id.to_vec()],
+            static_att: vec![],
+            change: ChangeType::Creation.into(),
+            protocol_type: None,
+            tx: Some(tx.clone()),
+        }
+    }
+
+    /// Replaces the tokens on this component.
+    pub fn with_tokens<B: AsRef<[u8]>>(mut self, tokens: &[B]) -> Self {
+        self.tokens = tokens
+            .iter()
+            .map(|e| e.as_ref().to_vec())
+            .collect::<Vec<Vec<u8>>>();
+        self
+    }
+
+    /// Replaces the contracts associated with this component.
+    pub fn with_contracts<B: AsRef<[u8]>>(mut self, contracts: &[B]) -> Self {
+        self.contracts = contracts
+            .iter()
+            .map(|e| e.as_ref().to_vec())
+            .collect::<Vec<Vec<u8>>>();
+        self
+    }
+
+    /// Replaces the static attributes on this component.
+    ///
+    /// The change type will be set to Creation.
+    pub fn with_attributes<K: AsRef<str>, V: AsRef<[u8]>>(mut self, attributes: &[(K, V)]) -> Self {
+        self.static_att = attributes
+            .iter()
+            .map(|(k, v)| Attribute {
+                name: k.as_ref().to_string(),
+                value: v.as_ref().to_vec(),
+                change: ChangeType::Creation.into(),
+            })
+            .collect::<Vec<Attribute>>();
+        self
+    }
+
+    /// Sets the protocol_type on this component.
+    ///
+    /// Will set the `financial_type` to FinancialType::Swap and the
+    /// `attribute_schema` to an empty list.
+    pub fn as_swap_type(mut self, name: &str, implementation_type: ImplementationType) -> Self {
+        self.protocol_type = Some(ProtocolType {
+            name: name.to_string(),
+            financial_type: FinancialType::Swap.into(),
+            attribute_schema: vec![],
+            implementation_type: implementation_type.into(),
+        });
+        self
+    }
+}

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

@@ -3,135 +3,8 @@ pub mod tycho {
     pub mod evm {
         // @@protoc_insertion_point(attribute:tycho.evm.v1)
         pub mod v1 {
-            use substreams_ethereum::pb::eth::v2::{self as sf};
             include!("tycho.evm.v1.rs");
             // @@protoc_insertion_point(tycho.evm.v1)
-
-            impl TransactionContractChanges {
-                /// Creates a new empty `TransactionContractChanges` instance.
-                pub fn new(tx: &Transaction) -> Self {
-                    Self {
-                        tx: Some(tx.clone()),
-                        contract_changes: vec![],
-                        component_changes: vec![],
-                        balance_changes: vec![],
-                    }
-                }
-            }
-
-            impl From<&sf::TransactionTrace> for Transaction {
-                fn from(tx: &sf::TransactionTrace) -> Self {
-                    Self {
-                        hash: tx.hash.clone(),
-                        from: tx.from.clone(),
-                        to: tx.to.clone(),
-                        index: tx.index.into(),
-                    }
-                }
-            }
-
-            impl From<&sf::Block> for Block {
-                fn from(block: &sf::Block) -> Self {
-                    Self {
-                        number: block.number,
-                        hash: block.hash.clone(),
-                        parent_hash: block
-                            .header
-                            .as_ref()
-                            .expect("Block header not present")
-                            .parent_hash
-                            .clone(),
-                        ts: block.timestamp_seconds(),
-                    }
-                }
-            }
-
-            impl ProtocolComponent {
-                /// Creates a new empty `ProtocolComponent` instance.
-                ///
-                /// You can use the `with_*` methods to set the fields in a convience way.
-                pub fn new(id: &str, tx: &Transaction) -> Self {
-                    Self {
-                        id: id.to_string(),
-                        tokens: vec![],
-                        contracts: vec![],
-                        static_att: vec![],
-                        change: ChangeType::Creation.into(),
-                        protocol_type: None,
-                        tx: Some(tx.clone()),
-                    }
-                }
-
-                /// Shorthand to create a component with a 1-1 relationship to a contract.
-                ///
-                /// Will set the component id to a hex encoded address with a 0x prefix
-                /// and add the contract to contracts attributes.
-                pub fn at_contract(id: &[u8], tx: &Transaction) -> Self {
-                    Self {
-                        id: format!("0x{}", hex::encode(id)),
-                        tokens: vec![],
-                        contracts: vec![id.to_vec()],
-                        static_att: vec![],
-                        change: ChangeType::Creation.into(),
-                        protocol_type: None,
-                        tx: Some(tx.clone()),
-                    }
-                }
-
-                /// Replaces the tokens on this component.
-                pub fn with_tokens<B: AsRef<[u8]>>(mut self, tokens: &[B]) -> Self {
-                    self.tokens = tokens
-                        .iter()
-                        .map(|e| e.as_ref().to_vec())
-                        .collect::<Vec<Vec<u8>>>();
-                    self
-                }
-
-                /// Replaces the contracts associated with this component.
-                pub fn with_contracts<B: AsRef<[u8]>>(mut self, contracts: &[B]) -> Self {
-                    self.contracts = contracts
-                        .iter()
-                        .map(|e| e.as_ref().to_vec())
-                        .collect::<Vec<Vec<u8>>>();
-                    self
-                }
-
-                /// Replaces the static attributes on this component.
-                ///
-                /// The change type will be set to Creation.
-                pub fn with_attributes<K: AsRef<str>, V: AsRef<[u8]>>(
-                    mut self,
-                    attributes: &[(K, V)],
-                ) -> Self {
-                    self.static_att = attributes
-                        .iter()
-                        .map(|(k, v)| Attribute {
-                            name: k.as_ref().to_string(),
-                            value: v.as_ref().to_vec(),
-                            change: ChangeType::Creation.into(),
-                        })
-                        .collect::<Vec<Attribute>>();
-                    self
-                }
-
-                /// Sets the protocol_type on this component.
-                ///
-                /// Will set the `financial_type` to FinancialType::Swap and the
-                /// `attribute_schema` to an empty list.
-                pub fn as_swap_type(
-                    mut self,
-                    name: &str,
-                    implementation_type: ImplementationType,
-                ) -> Self {
-                    self.protocol_type = Some(ProtocolType {
-                        name: name.to_string(),
-                        financial_type: FinancialType::Swap.into(),
-                        attribute_schema: vec![],
-                        implementation_type: implementation_type.into(),
-                    });
-                    self
-                }
-            }
         }
     }
 }

substreams/ethereum-balancer/src/modules.rs

@@ -9,12 +9,9 @@ use substreams::store::{
 };
 use substreams_ethereum::pb::eth;
 use substreams_ethereum::Event;
-use tycho_substreams::balances;
+use tycho_substreams::balances::aggregate_balances_changes;
 use tycho_substreams::contract::extract_contract_changes;
-use tycho_substreams::pb::tycho::evm::v1::{
+use tycho_substreams::prelude::*;
-    self as tycho, BalanceDelta, BlockBalanceDeltas, BlockTransactionProtocolComponents,
-    TransactionProtocolComponents,
-};
 
 const VAULT_ADDRESS: &[u8] = &hex!("BA12222222228d8Ba445958a75a0704d566BF2C8");
 
@@ -145,7 +142,7 @@ pub fn map_balance_deltas(
 ///  store key to ensure that there's a unique balance being tallied for each.
 #[substreams::handlers::store]
 pub fn store_balance_changes(deltas: BlockBalanceDeltas, store: StoreAddBigInt) {
-    balances::store_balance_changes(deltas, store);
+    tycho_substreams::balances::store_balance_changes(deltas, store);
 }
 
 /// This is the main map that handles most of the indexing of this substream.
@@ -161,11 +158,10 @@ pub fn map_changes(
     deltas: BlockBalanceDeltas,
     components_store: StoreGetInt64,
     balance_store: StoreDeltas, // Note, this map module is using the `deltas` mode for the store.
-) -> Result<tycho::BlockContractChanges> {
+) -> Result<BlockContractChanges> {
     // We merge contract changes by transaction (identified by transaction index) making it easy to
     //  sort them at the very end.
-    let mut transaction_contract_changes: HashMap<_, tycho::TransactionContractChanges> =
+    let mut transaction_contract_changes: HashMap<_, TransactionContractChanges> = HashMap::new();
-        HashMap::new();
 
     // `ProtocolComponents` are gathered from `map_pools_created` which just need a bit of work to
     //   convert into `TransactionContractChanges`
@@ -176,7 +172,7 @@ pub fn map_changes(
             let tx = tx_component.tx.as_ref().unwrap();
             transaction_contract_changes
                 .entry(tx.index)
-                .or_insert_with(|| tycho::TransactionContractChanges::new(&tx))
+                .or_insert_with(|| TransactionContractChanges::new(&tx))
                 .component_changes
                 .extend_from_slice(&tx_component.components);
         });
@@ -185,12 +181,12 @@ pub fn map_changes(
     //  `BlockBalanceDeltas`. We essentially just process the changes that occurred to the `store` this
     //  block. Then, these balance changes are merged onto the existing map of tx contract changes,
     //  inserting a new one if it doesn't exist.
-    balances::aggregate_balances_changes(balance_store, deltas)
+    aggregate_balances_changes(balance_store, deltas)
         .into_iter()
         .for_each(|(_, (tx, balances))| {
             transaction_contract_changes
                 .entry(tx.index)
-                .or_insert_with(|| tycho::TransactionContractChanges::new(&tx))
+                .or_insert_with(|| TransactionContractChanges::new(&tx))
                 .balance_changes
                 .extend(balances.into_iter().map(|(_, change)| change));
         });
@@ -208,7 +204,7 @@ pub fn map_changes(
 
     // Process all `transaction_contract_changes` for final output in the `BlockContractChanges`,
     //  sorted by transaction index (the key).
-    Ok(tycho::BlockContractChanges {
+    Ok(BlockContractChanges {
         block: Some((&block).into()),
         changes: transaction_contract_changes
             .drain()

substreams/ethereum-balancer/src/pool_factories.rs

@@ -1,11 +1,9 @@
-use substreams_ethereum::pb::eth::v2::{Call, Log};
-use substreams_ethereum::{Event, Function};
-
 use crate::abi;
 use substreams::hex;
-use tycho_substreams::pb::tycho::evm::v1::{ImplementationType, ProtocolComponent, Transaction};
-
 use substreams::scalar::BigInt;
+use substreams_ethereum::pb::eth::v2::{Call, Log};
+use substreams_ethereum::{Event, Function};
+use tycho_substreams::prelude::*;
 
 /// This trait defines some helpers for serializing and deserializing `Vec<BigInt` which is needed
 ///  to be able to encode the `normalized_weights` and `weights` `Attribute`s. This should also be