From eb3ca2aeec19955148153855194e7a6d965350a4 Mon Sep 17 00:00:00 2001
From: AntonAndell <anton@andell.eu>
Date: Thu, 27 Feb 2025 11:03:47 +0100
Subject: [PATCH] fix: optimize rewards query

---
 .../balanced/score/core/rewards/DataSourceImpl.java  | 12 +-----------
 .../balanced/score/core/rewards/RewardsImpl.java     | 11 ++++++++---
 2 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/DataSourceImpl.java b/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/DataSourceImpl.java
index b74d80ef4..45d6759cb 100644
--- a/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/DataSourceImpl.java
+++ b/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/DataSourceImpl.java
@@ -257,17 +257,7 @@ public void addRewardToken(Address token) {
 
     @SuppressWarnings("unchecked")
     public Map<String, BigInteger> loadCurrentSupply(String owner) {
-        // Bad handling that is only relevant during migration, otherwise it will always succeed on first scenario
-        try {
-            return (Map<String, BigInteger>) Context.call(getContractAddress(), "getBalanceAndSupply", getName(), owner);
-        } catch (Exception e) {
-            try {
-                return (Map<String, BigInteger>) Context.call(getContractAddress(), "getBalanceAndSupply", getName(), Address.fromString(owner));
-            } catch (Exception _e) {
-                return Map.of("_totalSupply", BigInteger.ZERO,
-                        "_balance", BigInteger.ZERO);
-            }
-        }
+        return (Map<String, BigInteger>) Context.call(getContractAddress(), "getBalanceAndSupply", getName(), owner);
     }
 
     public Map<Address, BigInteger> updateSingleUserData(BigInteger currentTime, BalanceData balances, String user, boolean readOnlyContext) {
diff --git a/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/RewardsImpl.java b/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/RewardsImpl.java
index bba22a8ab..475da6a85 100644
--- a/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/RewardsImpl.java
+++ b/core-contracts/Rewards/src/main/java/network/balanced/score/core/rewards/RewardsImpl.java
@@ -194,15 +194,20 @@ public BigInteger getBalnHolding(String _holder) {
     public Map<String, BigInteger> getRewards(String _holder) {
         Map<String, BigInteger> accruedRewards = getHoldings(_holder);
         int dataSourcesCount = DataSourceDB.size();
+        BigInteger currentTime = getTime();
+
         for (int i = 0; i < dataSourcesCount; i++) {
             String name = DataSourceDB.names.get(i);
             DataSourceImpl dataSource = DataSourceDB.get(name);
-            BigInteger currentTime = getTime();
             BalanceData balances = new BalanceData();
+            balances.prevWorkingBalance = dataSource.getWorkingBalance(_holder);
+            if (balances.prevWorkingBalance.compareTo(BigInteger.ZERO) <= 0) {
+                continue;
+            }
+
             balances.prevBalance = dataSource.getBalance(_holder);
             balances.prevSupply = dataSource.getTotalSupply();
-            balances.prevWorkingBalance = dataSource.getWorkingBalance(_holder, true);
-            balances.prevWorkingSupply = dataSource.getWorkingSupply(true);
+            balances.prevWorkingSupply = dataSource.getWorkingSupply();
             Map<Address, BigInteger> sourceRewards = dataSource.updateSingleUserData(currentTime, balances, _holder, true);
             for (Map.Entry<Address, BigInteger> entry : sourceRewards.entrySet()) {
                 accruedRewards.put(entry.getKey().toString(), accruedRewards.get(entry.getKey().toString()).add(entry.getValue()));