perf: improve Claimed index

bnb-chain · Sep 18, 2023 · 2f97ab5 · 2f97ab5
1 parent c383163
commit 2f97ab5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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.