From 8d5379d1a49ecfb9bf086dfb4c31ce62d860c275 Mon Sep 17 00:00:00 2001
From: shashwatS22 <46936675+shashwatS22@users.noreply.github.com>
Date: Fri, 13 Oct 2023 10:48:56 +0530
Subject: [PATCH] added token balance from contract

---
 .../level-finance/src/modules/transaction.ts   | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/subgraphs/level-finance/src/modules/transaction.ts b/subgraphs/level-finance/src/modules/transaction.ts
index 44dba6c277..0ab5d66f23 100644
--- a/subgraphs/level-finance/src/modules/transaction.ts
+++ b/subgraphs/level-finance/src/modules/transaction.ts
@@ -33,21 +33,23 @@ export function transaction(
   const idx = pool.getInputTokens().indexOf(token.id);
   const inputTokenBalances = pool.pool.inputTokenBalances;
   const network = dataSource.network();
+  const poolContract = PoolContract.bind(Address.fromBytes(pool.getBytesID()));
   if (utils.equalsIgnoreCase(network, constants.Network.ARBITRUM_ONE)) {
-    const poolContract = PoolContract.bind(
-      Address.fromBytes(pool.getBytesID())
-    );
     const tokenBalance = utils.readValue(
       poolContract.try_poolBalances(Address.fromBytes(token.id)),
       constants.BIGINT_ZERO
     );
     inputTokenBalances[idx] = tokenBalance;
   } else {
-    if (transactionType === TransactionType.DEPOSIT) {
-      inputTokenBalances[idx] = inputTokenBalances[idx].plus(amount);
-    }
-    if (transactionType === TransactionType.WITHDRAW) {
-      inputTokenBalances[idx] = inputTokenBalances[idx].minus(amount);
+    const poolTokenResult = poolContract.try_poolTokens(
+      Address.fromBytes(token.id)
+    );
+    if (!poolTokenResult.reverted) {
+      let tokenBalance = poolTokenResult.value.getPoolBalance();
+      if (!tokenBalance) {
+        tokenBalance = constants.BIGINT_ZERO;
+      }
+      inputTokenBalances[idx] = tokenBalance;
     }
   }
   const amountsArray = new Array<BigInt>(poolInputTokens.length).fill(