From 2f97ab55912f2dedf1d36b7163fd28b2a6fa2dcb Mon Sep 17 00:00:00 2001
From: j75689 <j75689@gmail.com>
Date: Mon, 18 Sep 2023 17:46:09 +0800
Subject: [PATCH] perf: improve Claimed index

---
 contracts/MerkleDistributor.sol            | 9 ++++++---
 contracts/MerkleDistributor.template       | 8 +++++---
 contracts/interface/IMerkleDistributor.sol | 2 +-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/contracts/MerkleDistributor.sol b/contracts/MerkleDistributor.sol
index 5a598bc1..22158f79 100644
--- a/contracts/MerkleDistributor.sol
+++ b/contracts/MerkleDistributor.sol
@@ -18,6 +18,7 @@ contract MerkleDistributor is IMerkleDistributor, System {
     mapping(uint256 => uint256) private claimedBitMap;
 
     struct ClaimInfo {
+        bytes32 tokenHash;
         bytes32 node;
         uint256 index;
         address contractAddr;
@@ -38,10 +39,12 @@ contract MerkleDistributor is IMerkleDistributor, System {
     }
 
 
-    function claim(bytes32 tokenSymbol, uint256 amount, bytes calldata prefixNode, bytes calldata suffixNode, bytes calldata ownerSignature, bytes calldata approvalSignature, bytes32[] calldata merkleProof) external override {
+    function claim(bytes32 tokenSymbol, uint256 tokenIndex, uint256 amount, bytes calldata prefixNode, bytes calldata suffixNode, bytes calldata ownerSignature, bytes calldata approvalSignature, bytes32[] calldata merkleProof) external override {
         ClaimInfo memory claimInfo;
-        claimInfo.node = keccak256(abi.encodePacked(prefixNode, tokenSymbol, amount, suffixNode));
-        claimInfo.index = uint256(keccak256(abi.encodePacked(tokenSymbol, claimInfo.node)));
+        claimInfo.tokenHash = keccak256(abi.encodePacked(tokenIndex, tokenSymbol, amount));
+        claimInfo.node = keccak256(abi.encodePacked(prefixNode, claimInfo.tokenHash, suffixNode));
+        // TODO: consider the better way to index the claim info.
+        claimInfo.index = keccak256(tokenSymbol, bytes32(tokenIndex), claimInfo.node);
     
         // Check if the token is claimed.
         require(isClaimed(claimInfo.index), "AlreadyClaimed");
diff --git a/contracts/MerkleDistributor.template b/contracts/MerkleDistributor.template
index 0be0cfad..f3770bb4 100644
--- a/contracts/MerkleDistributor.template
+++ b/contracts/MerkleDistributor.template
@@ -31,6 +31,7 @@ contract MerkleDistributor is IMerkleDistributor, System {
     mapping(uint256 => uint256) private claimedBitMap;
 
     struct ClaimInfo {
+        bytes32 tokenHash;
         bytes32 node;
         uint256 index;
         address contractAddr;
@@ -51,10 +52,11 @@ contract MerkleDistributor is IMerkleDistributor, System {
     }
 
 
-    function claim(bytes32 tokenSymbol, uint256 amount, bytes calldata prefixNode, bytes calldata suffixNode, bytes calldata ownerSignature, bytes calldata approvalSignature, bytes32[] calldata merkleProof) external override {
+    function claim(bytes32 tokenSymbol, uint256 tokenIndex, uint256 amount, bytes calldata prefixNode, bytes calldata suffixNode, bytes calldata ownerSignature, bytes calldata approvalSignature, bytes32[] calldata merkleProof) external override {
         ClaimInfo memory claimInfo;
-        claimInfo.node = keccak256(abi.encodePacked(prefixNode, tokenSymbol, amount, suffixNode));
-        claimInfo.index = uint256(keccak256(abi.encodePacked(tokenSymbol, claimInfo.node)));
+        claimInfo.tokenHash = keccak256(abi.encodePacked(tokenIndex, tokenSymbol, amount));
+        claimInfo.node = keccak256(abi.encodePacked(prefixNode, claimInfo.tokenHash, suffixNode));
+        claimInfo.index = uint256(claimInfo.node) + tokenIndex;
     
         // Check if the token is claimed.
         require(isClaimed(claimInfo.index), "AlreadyClaimed");
diff --git a/contracts/interface/IMerkleDistributor.sol b/contracts/interface/IMerkleDistributor.sol
index 24003f73..374cfa38 100644
--- a/contracts/interface/IMerkleDistributor.sol
+++ b/contracts/interface/IMerkleDistributor.sol
@@ -7,7 +7,7 @@ interface IMerkleDistributor {
     // Returns true if the index has been marked claimed.
     function isClaimed(uint256 index) external view returns (bool);
     // Claim the given amount of the token to the given address. Reverts if the inputs are invalid.
-    function claim(bytes32 tokenSymbol, uint256 amount, bytes calldata prefixNode, bytes calldata suffixNode, bytes calldata ownerSignature, bytes calldata approvalSignature, bytes32[] calldata merkleProof) external;
+    function claim(bytes32 tokenSymbol, uint256 tokenIndex, uint256 amount, bytes calldata prefixNode, bytes calldata suffixNode, bytes calldata ownerSignature, bytes calldata approvalSignature, bytes32[] calldata merkleProof) external;
     // registerToken register a token to the merkle distributor.
     function registerToken(bytes32 tokenSymbol, address contractAddr, uint256 decimals, uint256 amount, bytes calldata ownerSignature, bytes calldata approvalSignature) external;
     // This event is triggered whenever a call to #claim succeeds.