From 3c08359cf112e15c137dd5256b8dc8e9cd6c1626 Mon Sep 17 00:00:00 2001
From: Louise Poole <louisecarmenpoole@gmail.com>
Date: Mon, 27 Jan 2025 16:06:19 +0200
Subject: [PATCH] fix: allow setting of tokens balances on
 InterimContractChanges (#139)

---
 substreams/crates/tycho-substreams/src/models.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/substreams/crates/tycho-substreams/src/models.rs b/substreams/crates/tycho-substreams/src/models.rs
index 3dcee62..e41cfeb 100644
--- a/substreams/crates/tycho-substreams/src/models.rs
+++ b/substreams/crates/tycho-substreams/src/models.rs
@@ -515,6 +515,22 @@ impl InterimContractChange {
         self.code.clear();
         self.code.extend_from_slice(code);
     }
+
+    pub fn upsert_token_balance(&mut self, token: &[u8], balance: &[u8]) {
+        self.token_balances
+            .entry(token.to_vec())
+            .and_modify(|b| {
+                b.clear();
+                b.extend_from_slice(balance);
+            })
+            .or_insert_with(|| balance.to_vec());
+    }
+
+    pub fn upsert_token_balances(&mut self, balances: &HashMap<Vec<u8>, Vec<u8>>) {
+        for (token, balance) in balances.iter() {
+            self.upsert_token_balance(token, balance);
+        }
+    }
 }
 
 impl From<InterimContractChange> for Option<ContractChange> {