fix: Fix bug in MerkleDistributor that doesn't decrement remainingAmount for merkle windows correctly (#193)

* fix: Fix bug in MerkleDistributor that doesn't decrement remainingAmount for merkle windows correctly

* Update MerkleDistributor.sol

* Update MerkleDistributor.sol

Author: nicholaspai

contracts/merkle-distributor/MerkleDistributor.sol

@@ -156,7 +156,6 @@ contract MerkleDistributor is Ownable {
                 merkleWindows[claims[nextI].windowIndex].rewardToken != currentRewardToken
                 // Next claim reward token is different than current one.
             ) {
-                merkleWindows[_claim.windowIndex].remainingAmount -= batchedAmount;
                 currentRewardToken.safeTransfer(_claim.account, batchedAmount);
                 batchedAmount = 0;
             }
@@ -172,7 +171,6 @@ contract MerkleDistributor is Ownable {
      */
     function claim(Claim memory _claim) external {
         _verifyAndMarkClaimed(_claim);
-        merkleWindows[_claim.windowIndex].remainingAmount -= _claim.amount;
         merkleWindows[_claim.windowIndex].rewardToken.safeTransfer(_claim.account, _claim.amount);
     }
 
@@ -193,6 +191,15 @@ contract MerkleDistributor is Ownable {
         return claimedWord & mask == mask;
     }
 
+    /**
+     * @notice Returns rewardToken set by admin for windowIndex.
+     * @param windowIndex merkle root to check.
+     * @return address Reward token address
+     */
+    function getRewardTokenForWindow(uint256 windowIndex) public view returns (address) {
+        return address(merkleWindows[windowIndex].rewardToken);
+    }
+
     /**
      * @notice Returns True if leaf described by {account, amount, accountIndex} is stored in Merkle root at given
      *         window index.
@@ -245,6 +252,7 @@ contract MerkleDistributor is Ownable {
 
         // Proof is correct and claim has not occurred yet, mark claimed complete.
         _setClaimed(_claim.windowIndex, _claim.accountIndex);
+        merkleWindows[_claim.windowIndex].remainingAmount -= _claim.amount;
         emit Claimed(
             msg.sender,
             _claim.windowIndex,

test/merkle-distributor/MerkleDistributor.ts

@@ -551,11 +551,63 @@ describe("MerkleDistributor", () => {
         const events = await merkleDistributor.queryFilter(eventFilter());
         expect(events.length).to.equal(allRecipients.length * 3);
       });
+      it("Can make multiple claims for one token across multiple windows with single leaf trees", async function () {
+        // This tests that claimMulti correctly decrements `remainingAmount` for each merkle window.
+        const window1RewardAmount = toBN(toWei("100"));
+        const window2RewardAmount = toBN(toWei("300"));
+
+        // Set two windows with trivial one leaf trees.
+        const reward1Recipients = [
+          {
+            account: accounts[3].address,
+            amount: window1RewardAmount.toString(),
+            accountIndex: 1,
+          },
+        ];
+        const reward2Recipients = [
+          {
+            account: accounts[3].address,
+            amount: window2RewardAmount.toString(),
+            accountIndex: 1,
+          },
+        ];
+        const merkleTree1 = new MerkleTree(reward1Recipients.map((item) => createLeaf(item)));
+        const nextWindowIndex = (await merkleDistributor.nextCreatedIndex()).toNumber();
+        await merkleDistributor
+          .connect(contractCreator)
+          .setWindow(window1RewardAmount, rewardToken.address, merkleTree1.getRoot(), "");
+        const merkleTree2 = new MerkleTree(reward2Recipients.map((item) => createLeaf(item)));
+        await merkleDistributor
+          .connect(contractCreator)
+          .setWindow(window2RewardAmount, rewardToken.address, merkleTree2.getRoot(), "");
+
+        batchedClaims = [
+          {
+            windowIndex: nextWindowIndex,
+            account: reward1Recipients[0].account,
+            accountIndex: reward1Recipients[0].accountIndex,
+            amount: reward1Recipients[0].amount,
+            merkleProof: merkleTree1.getProof(createLeaf(reward1Recipients[0])),
+          },
+          {
+            windowIndex: nextWindowIndex + 1,
+            account: reward2Recipients[0].account,
+            accountIndex: reward2Recipients[0].accountIndex,
+            amount: reward2Recipients[0].amount,
+            merkleProof: merkleTree2.getProof(createLeaf(reward2Recipients[0])),
+          },
+        ];
+
+        await merkleDistributor.claimMulti(batchedClaims);
+        // const eventFilter = merkleDistributor.filters.Claimed;
+        // const events = await merkleDistributor.queryFilter(eventFilter());
+        // expect(events.length).to.equal(batchedClaims.length);
+      });
       it("gas", async function () {
         const txn = await merkleDistributor.connect(contractCreator).claimMulti(batchedClaims);
         const receipt = await txn.wait();
         assertApproximate(
-          32860,
+          32185,
           Math.floor(receipt.gasUsed / (rewardLeafs.length * Object.keys(SamplePayouts.recipients).length)),
           0.02
         );