From 3944650a0ce41c8ced2a54072262fc321f4c9143 Mon Sep 17 00:00:00 2001
From: Ryan Sauge <ryan.sauge@protonmail.com>
Date: Wed, 10 Apr 2024 10:36:58 +0200
Subject: [PATCH] SnapshotModule - update snapshotInfoBatch

---
 contracts/interfaces/ICMTATSnapshot.sol         |  8 ++++++++
 .../internal/ERC20SnapshotModuleInternal.sol    | 17 +++++++++++++++++
 .../ERC20SnapshotModuleUtils.js                 | 15 ++++++++++-----
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/contracts/interfaces/ICMTATSnapshot.sol b/contracts/interfaces/ICMTATSnapshot.sol
index 8dcb4881..32792e5b 100644
--- a/contracts/interfaces/ICMTATSnapshot.sol
+++ b/contracts/interfaces/ICMTATSnapshot.sol
@@ -27,4 +27,12 @@ interface ICMTATSnapshot {
     * @return ownerBalances array with the balance of each address, the total supply
     */
     function snapshotInfoBatch(uint256 time, address[] calldata addresses) external view returns (uint256[] memory ownerBalances, uint256 totalSupply);
+
+    /**
+    * @notice Return snapshotBalanceOf for each address in the array and the total supply
+    * @return ownerBalances array with the balance of each address, the total supply
+    */
+    function snapshotInfoBatch(uint256[] calldata times, address[] calldata addresses) external view returns (uint256[][] memory ownerBalances, uint256[] memory totalSupply);
+
+
 }
diff --git a/contracts/modules/internal/ERC20SnapshotModuleInternal.sol b/contracts/modules/internal/ERC20SnapshotModuleInternal.sol
index 47cc8771..6dd70ded 100644
--- a/contracts/modules/internal/ERC20SnapshotModuleInternal.sol
+++ b/contracts/modules/internal/ERC20SnapshotModuleInternal.sol
@@ -55,6 +55,23 @@ abstract contract ERC20SnapshotModuleInternal is SnapshotModuleBase, ERC20Upgrad
         totalSupply = snapshotTotalSupply(time);
     }
 
+    /**
+    * @notice Return snapshotBalanceOf for each address in the array and the total supply
+    * @return ownerBalances array with the balance of each address, the total supply
+    */
+    function snapshotInfoBatch(uint256[] calldata times, address[] calldata addresses) public view returns (uint256[][] memory ownerBalances, uint256[] memory totalSupply) {
+        ownerBalances = new uint256[][](times.length);
+        totalSupply = new uint256[](times.length);
+        for(uint256 iT = 0; iT < times.length; ++iT){
+            /*ownerBalances[iT] = new uint256[](addresses.length);
+            for(uint256 jA = 0; jA < addresses.length; ++jA){
+                ownerBalances[iT][jA]  = snapshotBalanceOf(times[iT], addresses[jA]);
+            }
+            totalSupply[iT] = snapshotTotalSupply(times[iT]);*/
+            (ownerBalances[iT], totalSupply[iT]) = snapshotInfoBatch(times[iT],addresses);
+        }
+    }
+
     /** 
     * @notice Return the number of tokens owned by the given owner at the time when the snapshot with the given time was created.
     * @return value stored in the snapshot, or the actual balance if no snapshot
diff --git a/test/common/ERC20SnapshotModuleCommon/ERC20SnapshotModuleUtils/ERC20SnapshotModuleUtils.js b/test/common/ERC20SnapshotModuleCommon/ERC20SnapshotModuleUtils/ERC20SnapshotModuleUtils.js
index cdee1ed0..e0a8f18a 100644
--- a/test/common/ERC20SnapshotModuleCommon/ERC20SnapshotModuleUtils/ERC20SnapshotModuleUtils.js
+++ b/test/common/ERC20SnapshotModuleCommon/ERC20SnapshotModuleUtils/ERC20SnapshotModuleUtils.js
@@ -1,3 +1,4 @@
+const { BN } = require('@openzeppelin/test-helpers')
 const getUnixTimestamp = () => {
   return Math.round(new Date().getTime() / 1000)
 }
@@ -10,20 +11,24 @@ async function checkSnapshot (time, totalSupply, addresses, balances) {
   // Values before the snapshot
   (await this.cmtat.snapshotTotalSupply(time)).should.be.bignumber.equal(
     totalSupply
-  )
-  const result = await this.cmtat.snapshotInfoBatch(time, addresses)
+  );
+  const result = await this.cmtat.methods['snapshotInfoBatch(uint256,address[])'](time, addresses);
+  const times = [time]
+  const result2 = await this.cmtat.methods['snapshotInfoBatch(uint256[],address[])'](times, addresses);
   for (let i = 0; i < balances.length; ++i) {
     (
       await this.cmtat.snapshotBalanceOf(time, addresses[i])
     ).should.be.bignumber.equal(balances[i])
-    await this.cmtat.snapshotInfo(time, addresses[i])
-    const { 0: ownerBalance, 1: totalSupplyGet } = await this.cmtat.snapshotInfo(time, addresses[i])
+    await this.cmtat.snapshotInfo(time, addresses[i]);
+    const { 0: ownerBalance, 1: totalSupplyGet } = await this.cmtat.snapshotInfo(time, addresses[i]);
     // const [ownerBalance, totalSupplyGet]
     ownerBalance.should.be.bignumber.equal(balances[i])
     result[0][i].should.be.bignumber.equal(balances[i])
+    result2[0][0][i].should.be.bignumber.equal(balances[i]);
     totalSupplyGet.should.be.bignumber.equal(totalSupply)
   }
-  result[1].should.be.bignumber.equal(totalSupply)
+  result[1].should.be.bignumber.equal(totalSupply);
+  result2[1][0].should.be.bignumber.equal(totalSupply);
 }
 
 async function checkArraySnapshot (snapshots, snapshotsValue) {