From 5b0a5ccf4ff39f0e2f3999e122621890b02ba7f7 Mon Sep 17 00:00:00 2001
From: arigatodl <99845398+dulguun-staderlabs@users.noreply.github.com>
Date: Fri, 26 Apr 2024 13:23:18 +0800
Subject: [PATCH] Add batch balance read

---
 contracts/BatchReader.sol             | 34 ++++++++++++++++++---------
 contracts/interfaces/IBatchReader.sol |  2 ++
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/contracts/BatchReader.sol b/contracts/BatchReader.sol
index 30af166b..f0fc595d 100644
--- a/contracts/BatchReader.sol
+++ b/contracts/BatchReader.sol
@@ -1,27 +1,39 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
-pragma solidity 0.8.16;
+pragma solidity ^0.8.16;
 
 import "./interfaces/IPenalty.sol";
 import "./interfaces/IStaderConfig.sol";
 import "./interfaces/IBatchReader.sol";
 
+/// @title Batch Reader Contract
+/// @notice Provides batch operations for reading penalties and balances
 contract BatchReader is IBatchReader {
-    IPenalty public penaltyContract;
-    IStaderConfig immutable staderConfig;
+    IPenalty public immutable PENALTY_CONTRACT;
 
+    /// @param _staderConfig Address of the StaderConfig contract
     constructor(address _staderConfig) {
-        require(_staderConfig != address(0), "Invalid StaderConfig address");
+        require(_staderConfig != address(0), "BatchReader: invalid StaderConfig address");
 
-        staderConfig = IStaderConfig(_staderConfig);
-        penaltyContract = IPenalty(staderConfig.getPenaltyContract());
+        PENALTY_CONTRACT = IPenalty(IStaderConfig(_staderConfig).getPenaltyContract());
     }
 
-    /// @notice Reads the total penalty amounts for a batch of validators.
-    function getTotalPenaltyAmounts(bytes[] calldata _pubkeys) external view override returns (uint256[] memory) {
-        uint256[] memory penalties = new uint256[](_pubkeys.length);
+    /// @notice Reads the total penalty amounts for a batch of validators
+    /// @param _pubkeys Public keys of the validators
+    /// @return penalties Array of total penalty amounts for each validator
+    function getTotalPenaltyAmounts(bytes[] calldata _pubkeys) external view override returns (uint256[] memory penalties) {
+        penalties = new uint256[](_pubkeys.length);
         for (uint256 i = 0; i < _pubkeys.length; i++) {
-            penalties[i] = penaltyContract.totalPenaltyAmount(_pubkeys[i]);
+            penalties[i] = PENALTY_CONTRACT.totalPenaltyAmount(_pubkeys[i]);
+        }
+    }
+
+    /// @notice Gets balances for multiple addresses
+    /// @param _addresses Array of Ethereum addresses
+    /// @return balances Array of balances for the provided addresses
+    function getBalances(address[] calldata _addresses) external view override returns (uint256[] memory balances) {
+        balances = new uint256[](_addresses.length);
+        for (uint256 i = 0; i < _addresses.length; i++) {
+            balances[i] = _addresses[i].balance;
         }
-        return penalties;
     }
 }
diff --git a/contracts/interfaces/IBatchReader.sol b/contracts/interfaces/IBatchReader.sol
index d33de782..6a2df044 100644
--- a/contracts/interfaces/IBatchReader.sol
+++ b/contracts/interfaces/IBatchReader.sol
@@ -3,4 +3,6 @@ pragma solidity 0.8.16;
 
 interface IBatchReader {
     function getTotalPenaltyAmounts(bytes[] calldata _pubkeys) external view returns (uint256[] memory);
+
+    function getBalances(address[] calldata _addresses) external view returns (uint256[] memory);
 }