From 178245253ae584ee84e8f5f4c8c1e20934a1ba41 Mon Sep 17 00:00:00 2001 From: 0xMemoryGrinder <35138272+0xMemoryGrinder@users.noreply.github.com> Date: Wed, 15 Nov 2023 21:59:39 -0500 Subject: [PATCH] feat(contracts): added MerkleDistributor --- .../src/abstracts/AMerkleDistributor.sol | 137 +++++++++++++ .../src/interfaces/IMerkleDistributor.sol | 37 ++++ contracts/src/utils/Errors.sol | 5 + contracts/src/utils/MerkleRootBuilder.sol | 76 +++++++ .../AMerkleDistributorTest.sol | 29 +++ .../AMerkleDistributor/AddRoundRewards.t.sol | 59 ++++++ contracts/test/AMerkleDistributor/Claim.t.sol | 95 +++++++++ .../test/AMerkleDistributor/Constructor.t.sol | 16 ++ .../GetProductRoundMerkleRoot.t.sol | 37 ++++ .../test/AMerkleDistributor/IsClaimed.t.sol | 68 +++++++ .../test/AMerkleDistributor/MultiClaim.t.sol | 186 ++++++++++++++++++ contracts/test/mock/AMerkleBitDistributor.sol | 9 + contracts/test/mock/AMerkleDistributor.sol | 9 + 13 files changed, 763 insertions(+) create mode 100644 contracts/src/abstracts/AMerkleDistributor.sol create mode 100644 contracts/src/interfaces/IMerkleDistributor.sol create mode 100644 contracts/src/utils/MerkleRootBuilder.sol create mode 100644 contracts/test/AMerkleDistributor/AMerkleDistributorTest.sol create mode 100644 contracts/test/AMerkleDistributor/AddRoundRewards.t.sol create mode 100644 contracts/test/AMerkleDistributor/Claim.t.sol create mode 100644 contracts/test/AMerkleDistributor/Constructor.t.sol create mode 100644 contracts/test/AMerkleDistributor/GetProductRoundMerkleRoot.t.sol create mode 100644 contracts/test/AMerkleDistributor/IsClaimed.t.sol create mode 100644 contracts/test/AMerkleDistributor/MultiClaim.t.sol create mode 100644 contracts/test/mock/AMerkleBitDistributor.sol create mode 100644 contracts/test/mock/AMerkleDistributor.sol diff --git a/contracts/src/abstracts/AMerkleDistributor.sol b/contracts/src/abstracts/AMerkleDistributor.sol new file mode 100644 index 0000000..a9fa34c --- /dev/null +++ b/contracts/src/abstracts/AMerkleDistributor.sol @@ -0,0 +1,137 @@ +//SPDX-License-Identifier: GPL-3.0-or-later +pragma solidity ^0.8.20; + +import { IMerkleDistributor } from "../interfaces/IMerkleDistributor.sol"; +import { AOperator } from "./AOperator.sol"; +import { MerkleProofLib } from "solady/utils/MerkleProofLib.sol"; +import { MerkleRootBuilder } from "../utils/MerkleRootBuilder.sol"; +import { SafeTransferLib } from "solady/utils/SafeTransferLib.sol"; +import { Errors } from "../utils/Errors.sol"; + +/** + * @title AMerkleDistributor + * @author 0xMemoryGrinder + * @dev Abstract contract which implements the merkle root update + */ +abstract contract AMerkleDistributor is IMerkleDistributor, AOperator { + // vault/pounder => currentRound + mapping(address => uint256) public currentRound; + // vault/pounder => round => merkel root + mapping(address => mapping(uint256 => bytes32)) public roundsMerkleRoots; + // vault/pounder => round => claimedIndex => claimedBitMap (packed boolean) + mapping(address => mapping(uint256 => mapping(uint256 => uint256))) public claimedStatus; + + // Events + event MerkleRootUpdated(address product, uint256 round, bytes32 merkleRoot); + event Claimed(address product, uint256 round, address token, address account, uint256 amount); + + constructor(address initialOperator) AOperator(initialOperator) {} + + function claim(address product, uint256 round, uint256 index, address token, address account, uint256 amount, bytes32[] calldata merkleProof) external override { + _claim(product, round, index, token, account, amount, merkleProof); + } + + function multiClaim(address product, address account, IMerkleDistributor.ClaimParams[] calldata params) external override { + _multiClaim(product, account, params); + } + + /** + * @dev Update the merkle root + * @param newMerkleRoot New merkle root after a harvest, set by the operator + */ + function addRoundRewards(address product, bytes32 newMerkleRoot) external onlyOperatorOrOwner { + if (product == address(0)) revert Errors.ZeroAddress(); + if (newMerkleRoot == bytes32(0)) revert Errors.ZeroValue(); + + uint256 round = currentRound[product]; + roundsMerkleRoots[product][round] = newMerkleRoot; + currentRound[product] = round + 1; + emit MerkleRootUpdated(product, round, newMerkleRoot); + } + + function getProductRoundMerkleRoot(address product, uint256 round) external view returns (bytes32) { + if (product == address(0)) revert Errors.ZeroAddress(); + if (roundsMerkleRoots[product][round] == bytes32(0)) revert Errors.MerkleDistributorNotInitialized(); + return roundsMerkleRoots[product][round]; + } + + /** + * @notice Checks if the rewards were claimed for an user on a given period + * @dev Checks if the rewards were claimed for an user (based on the index) on a given period + * @param product Address of the product + * @param round Round id + * @param index Index of the claim + * @return bool : true if already claimed + */ + function isClaimed(address product, uint256 round, uint256 index) public view returns (bool) { + uint256 claimedWordIndex = index >> 8; + uint256 claimedBitIndex = index & 0xff; + uint256 claimedWord = claimedStatus[product][round][claimedWordIndex]; + uint256 mask = (1 << claimedBitIndex); + return claimedWord & mask != 0; + } + + /** + * @dev Sets the rewards as claimed for the index on the given period + * @param product Address of the product + * @param round Round id + * @param index Index of the claim + */ + function _setClaimed(address product, uint256 round, uint256 index) private { + uint256 claimedWordIndex = index >> 8; + uint256 claimedBitIndex = index & 0xff; + claimedStatus[product][round][claimedWordIndex] |= (1 << claimedBitIndex); + } + + /** + * @dev Claim a token reward for a given account + * @param product Product address to claim + * @param round Round id + * @param token Token address to claim + * @param account Account to claim rewards for + * @param amount Rewards amount to claim + * @param merkleProof Merkle proof to validate the claim + */ + function _claim(address product, uint256 round, uint256 index, address token, address account, uint256 amount, bytes32[] calldata merkleProof) internal { + if (token == address(0) || account == address(0)) revert Errors.ZeroAddress(); + if (amount == 0) revert Errors.ZeroValue(); + if (roundsMerkleRoots[product][round] == bytes32(0)) revert Errors.MerkleDistributorNotInitialized(); + if (isClaimed(product, round, index)) revert Errors.AlreadyClaimed(); + + // Verify the claim info with the merkle proof + if (!_verifyProof(product, round, index, token, account, amount, merkleProof)) revert Errors.InvalidMerkleProof(); + // Transfer the reward to the account + SafeTransferLib.safeTransferFrom(token, address(this), account, amount); + _setClaimed(product, round, index); + emit Claimed(product, round, token, account, amount); + } + + /** + * @dev Verify the claim info with the merkle proof + */ + function _verifyProof(address product, uint256 round, uint256 index, address token, address account, uint256 amount, bytes32[] calldata proof) private view returns (bool) { + // Create merkle leaf + bytes32 leaf = keccak256(abi.encodePacked(index, token, account, amount)); + + // Check if the computed hash (root) is equal to the provided root + return MerkleProofLib.verifyCalldata(proof, roundsMerkleRoots[product][round], leaf); + } + + /** + * @dev Claim multiple token rewards for given accounts + * @param product Product address to claim rewards for + * @param account Account to claim rewards for + * @param params ClaimParams array containing the indexesn tokens, amounts and merkle proofs + */ + function _multiClaim(address product, address account, IMerkleDistributor.ClaimParams[] calldata params) internal { + uint256 length = params.length; + + if(length == 0) revert Errors.EmptyArray(); + + for(uint256 i; i < length;){ + _claim(product, params[i].round, params[i].index, params[i].token, account, params[i].amount, params[i].merkleProof); + + unchecked{ ++i; } + } + } +} \ No newline at end of file diff --git a/contracts/src/interfaces/IMerkleDistributor.sol b/contracts/src/interfaces/IMerkleDistributor.sol new file mode 100644 index 0000000..c2ee1e8 --- /dev/null +++ b/contracts/src/interfaces/IMerkleDistributor.sol @@ -0,0 +1,37 @@ +//SPDX-License-Identifier: GPL-3.0-or-later +pragma solidity ^0.8.20; + +/** + * @title IMerkleDistributor + * @author 0xMemoryGrinder + * @dev MerkleDistributor interface which updates the merkle root when the claims are made + */ +interface IMerkleDistributor { + struct ClaimParams { + address token; + uint256 round; + uint256 index; + uint256 amount; + bytes32[] merkleProof; + } + + /** + * @dev Claim a token reward for a given account + * @param product Product address to claim rewards for + * @param round Round to claim rewards for + * @param index Index in the merkle tree to claim rewards for + * @param token Token address to claim + * @param account Account to claim rewards for + * @param amount Rewards amount to claim + * @param merkleProof Merkle proof to validate the claim + */ + function claim(address product, uint256 round, uint256 index, address token, address account, uint256 amount, bytes32[] calldata merkleProof) external; + + /** + * @dev Claim multiple token rewards for given accounts + * @param product Product address to claim rewards for + * @param account Account to claim rewards for + * @param params ClaimParams array containing the tokens, amounts and merkle proofs + */ + function multiClaim(address product, address account, ClaimParams[] calldata params) external; +} \ No newline at end of file diff --git a/contracts/src/utils/Errors.sol b/contracts/src/utils/Errors.sol index be379a7..d954761 100644 --- a/contracts/src/utils/Errors.sol +++ b/contracts/src/utils/Errors.sol @@ -11,4 +11,9 @@ library Errors { // Swapper errors error SwapError(); error NotVault(); + + // MerkleDistributor errors + error MerkleDistributorNotInitialized(); + error InvalidMerkleProof(); + error AlreadyClaimed(); } diff --git a/contracts/src/utils/MerkleRootBuilder.sol b/contracts/src/utils/MerkleRootBuilder.sol new file mode 100644 index 0000000..6dfbba1 --- /dev/null +++ b/contracts/src/utils/MerkleRootBuilder.sol @@ -0,0 +1,76 @@ +//SPDX-License-Identifier GPL-3.0-or-later +pragma solidity ^0.8.20; + +/** + * @title MerkleRootBuilder + * @author 0xMemoryGrinder + * @dev Compute a merkle root from a new leaf and a merkle proof + * Adapted from solady library with returning the last hash instead of comparing it to the merkle root + */ +library MerkleRootBuilder { + /** + * @dev Compute a merkle root from a new leaf and a merkle proof + * @param leaf New leaf to be added to the merkle tree + * @param proof Merkle proof to validate the new leaf + * @return newRoot New merkle root + */ + function computeMerkleRoot(bytes32 leaf, bytes32[] memory proof) internal pure returns (bytes32 newRoot) { + /// @solidity memory-safe-assembly + assembly { + if mload(proof) { + // Initialize `offset` to the offset of `proof` elements in memory. + let offset := add(proof, 0x20) + // Left shift by 5 is equivalent to multiplying by 0x20. + let end := add(offset, shl(5, mload(proof))) + // Iterate over proof elements to compute root hash. + for {} 1 {} { + // Slot of `leaf` in scratch space. + // If the condition is true: 0x20, otherwise: 0x00. + let scratch := shl(5, gt(leaf, mload(offset))) + // Store elements to hash contiguously in scratch space. + // Scratch space is 64 bytes (0x00 - 0x3f) and both elements are 32 bytes. + mstore(scratch, leaf) + mstore(xor(scratch, 0x20), mload(offset)) + // Reuse `leaf` to store the hash to reduce stack operations. + leaf := keccak256(0x00, 0x40) + offset := add(offset, 0x20) + if iszero(lt(offset, end)) { break } + } + } + newRoot := leaf + } + } + + /** + * @dev Compute a merkle root from a new leaf and a merkle proof + * @param leaf New leaf to be added to the merkle tree + * @param proof Merkle proof to validate the new leaf + * @return newRoot New merkle root + */ + function computeMerkleRootCalldata(bytes32 leaf, bytes32[] calldata proof) internal pure returns (bytes32 newRoot) { + /// @solidity memory-safe-assembly + assembly { + if proof.length { + // Left shift by 5 is equivalent to multiplying by 0x20. + let end := add(proof.offset, shl(5, proof.length)) + // Initialize `offset` to the offset of `proof` in the calldata. + let offset := proof.offset + // Iterate over proof elements to compute root hash. + for {} 1 {} { + // Slot of `leaf` in scratch space. + // If the condition is true: 0x20, otherwise: 0x00. + let scratch := shl(5, gt(leaf, calldataload(offset))) + // Store elements to hash contiguously in scratch space. + // Scratch space is 64 bytes (0x00 - 0x3f) and both elements are 32 bytes. + mstore(scratch, leaf) + mstore(xor(scratch, 0x20), calldataload(offset)) + // Reuse `leaf` to store the hash to reduce stack operations. + leaf := keccak256(0x00, 0x40) + offset := add(offset, 0x20) + if iszero(lt(offset, end)) { break } + } + } + newRoot := leaf + } + } +} \ No newline at end of file diff --git a/contracts/test/AMerkleDistributor/AMerkleDistributorTest.sol b/contracts/test/AMerkleDistributor/AMerkleDistributorTest.sol new file mode 100644 index 0000000..db0a451 --- /dev/null +++ b/contracts/test/AMerkleDistributor/AMerkleDistributorTest.sol @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "../BaseTest.sol"; +import { AMerkleDistributorMock } from "../mock/AMerkleDistributor.sol"; +import { ERC20Mock } from "../mock/ERC20.sol"; + +contract AMerkleDistributorTest is BaseTest { + AMerkleDistributorMock distributor; + + address product1 = makeAddr("product1"); + address product2 = makeAddr("product2"); + address token1 = makeAddr("token1"); + address token2 = makeAddr("token2"); + address token3 = makeAddr("token3"); + address token4 = makeAddr("token4"); + + // Events + event MerkleRootUpdated(bytes32 merkleRoot); + event Claimed(address token, address account, uint256 amount); + + function setUp() public virtual { + vm.startPrank(owner); + + distributor = new AMerkleDistributorMock(admin, owner); + + vm.stopPrank(); + } +} diff --git a/contracts/test/AMerkleDistributor/AddRoundRewards.t.sol b/contracts/test/AMerkleDistributor/AddRoundRewards.t.sol new file mode 100644 index 0000000..8f2b5dc --- /dev/null +++ b/contracts/test/AMerkleDistributor/AddRoundRewards.t.sol @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "./AMerkleDistributorTest.sol"; + +contract AddRoundRewards is AMerkleDistributorTest { + + function generateSimpleMerkleRoot() public view returns (bytes32) { + bytes32 leaf1 = keccak256(abi.encodePacked(uint256(0), token1, bob, uint256(10000))); + bytes32 leaf2 = keccak256(abi.encodePacked(uint256(1), token1, alice, uint256(10000))); + bytes32 root = keccak256(leaf1 < leaf2 ? abi.encodePacked(leaf1, leaf2) : abi.encodePacked(leaf2, leaf1)); + + return root; + } + + function test_addRoundRewards_NormalSingle() public { + vm.prank(admin); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + assertEq(distributor.currentRound(product1), 1, "Should increment by 1 the currentRound"); + assertFalse(distributor.getProductRoundMerkleRoot(product1, 0) == 0, "Should have one round merkle root"); + } + + function test_addRoundRewards_NormalMultiple() public { + vm.startPrank(admin); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + vm.stopPrank(); + assertEq(distributor.currentRound(product1), 5, "Should increment by 5 the currentRound"); + assertFalse(distributor.getProductRoundMerkleRoot(product1, 4) == 0, "Should have 5 rounds merkle roots"); + } + + function test_addRoundRewards_ZeroValueRoot() public { + vm.prank(admin); + vm.expectRevert(Errors.ZeroValue.selector); + distributor.addRoundRewards(product1, bytes32(0)); + } + + function test_addRoundRewards_ZeroAddressProduct() public { + vm.prank(admin); + vm.expectRevert(Errors.ZeroAddress.selector); + distributor.addRoundRewards(address(0), generateSimpleMerkleRoot()); + } + + function test_addRoundRewards_NotOperatorOrOwner() public { + vm.prank(bob); + vm.expectRevert(Errors.NotOperatorOrOwner.selector); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + } + + function test_addRoundRewards_Owner() public { + vm.prank(owner); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + assertEq(distributor.currentRound(product1), 1, "Should increment by 1 the currentRound"); + assertFalse(distributor.getProductRoundMerkleRoot(product1, 0) == 0, "Should have one round merkle root"); + } +} diff --git a/contracts/test/AMerkleDistributor/Claim.t.sol b/contracts/test/AMerkleDistributor/Claim.t.sol new file mode 100644 index 0000000..3faefa0 --- /dev/null +++ b/contracts/test/AMerkleDistributor/Claim.t.sol @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "./AMerkleDistributorTest.sol"; + +contract Claim is AMerkleDistributorTest { + bytes32 leaf1 = keccak256(abi.encodePacked(uint256(0), token1, bob, uint256(10000))); + bytes32 leaf2 = keccak256(abi.encodePacked(uint256(1), token2, alice, uint256(10000))); + + function generateSimpleMerkleRoot() public view returns (bytes32) { + + bytes32 root = keccak256(leaf1 < leaf2 ? abi.encodePacked(leaf1, leaf2) : abi.encodePacked(leaf2, leaf1)); + + return root; + } + + function setUp() public virtual override { + super.setUp(); + + vm.startPrank(admin); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + vm.stopPrank(); + } + + function test_claim_Simple() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + distributor.claim(product1, 0, 0, token1, bob, 10000, proof); + } + + function test_claim_Simple2Round() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + vm.prank(admin); + distributor.addRoundRewards(product1, generateSimpleMerkleRoot()); + distributor.claim(product1, 1, 0, token1, bob, 10000, proof); + } + + function testFuzz_claim_InvalidRound(uint256 round) public { + vm.assume(round != 0); + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + vm.expectRevert(Errors.MerkleDistributorNotInitialized.selector); + distributor.claim(product1, round, 0, token1, bob, 10000, proof); + } + + function test_claim_InvalidMerkleProof() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf1; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.claim(product1, 0, 0, token1, bob, 10000, proof); + } + + function test_claim_InvalidIndex() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.claim(product1, 0, 1, token1, bob, 10000, proof); + } + + function test_claim_InvalidToken() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.claim(product1, 0, 0, token2, bob, 10000, proof); + } + + function test_claim_InvalidAmount() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.claim(product1, 0, 0, token1, bob, 10001, proof); + } + + function test_claim_InvalidAddress() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + vm.expectRevert(Errors.ZeroAddress.selector); + distributor.claim(product1, 0, 0, token1, address(0), 10000, proof); + } + + function test_claim_InvalidMerkleProofLength() public { + bytes32[] memory proof = new bytes32[](0); + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.claim(product1, 0, 0, token1, bob, 10000, proof); + } + + function test_claim_alreadyClaimed() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leaf2; + distributor.claim(product1, 0, 0, token1, bob, 10000, proof); + vm.expectRevert(Errors.AlreadyClaimed.selector); + distributor.claim(product1, 0, 0, token1, bob, 10000, proof); + } +} diff --git a/contracts/test/AMerkleDistributor/Constructor.t.sol b/contracts/test/AMerkleDistributor/Constructor.t.sol new file mode 100644 index 0000000..7d61b69 --- /dev/null +++ b/contracts/test/AMerkleDistributor/Constructor.t.sol @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "./AMerkleDistributorTest.sol"; + +contract Constructor is AMerkleDistributorTest { + function test_constructor_Normal() public { + assertEq(distributor.operator(), admin, "operator should be admin"); + assertEq(distributor.currentRound(product1), 0, "Current round should start at 0"); + } + + function test_constructor_ZeroAddressOperator() public { + vm.expectRevert(Errors.ZeroAddress.selector); + new AMerkleDistributorMock(zero, owner); + } +} diff --git a/contracts/test/AMerkleDistributor/GetProductRoundMerkleRoot.t.sol b/contracts/test/AMerkleDistributor/GetProductRoundMerkleRoot.t.sol new file mode 100644 index 0000000..4161c11 --- /dev/null +++ b/contracts/test/AMerkleDistributor/GetProductRoundMerkleRoot.t.sol @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "./AMerkleDistributorTest.sol"; + +contract GetProductRoundMerkleRoot is AMerkleDistributorTest { + bytes32 leafA1 = keccak256(abi.encodePacked(uint256(0), token1, bob, uint256(10000))); + bytes32 leafA2 = keccak256(abi.encodePacked(uint256(1), token2, alice, uint256(10000))); + bytes32 leafB1 = keccak256(abi.encodePacked(uint256(2), token3, bob, uint256(10000))); + bytes32 leafB2 = keccak256(abi.encodePacked(uint256(3), token4, alice, uint256(10000))); + bytes32 rootA = keccak256(leafA1 < leafA2 ? abi.encodePacked(leafA1, leafA2) : abi.encodePacked(leafA2, leafA1)); + bytes32 rootB = keccak256(leafB1 < leafB2 ? abi.encodePacked(leafB1, leafB2) : abi.encodePacked(leafB2, leafB1)); + + function setUp() public virtual override { + super.setUp(); + + vm.startPrank(admin); + distributor.addRoundRewards(product1, rootA); + distributor.addRoundRewards(product1, rootB); + vm.stopPrank(); + } + + function test_getProductRoundMerkleRoot_Simple() public { + assertEq(distributor.getProductRoundMerkleRoot(product1, 0), rootA, "Should return the correct merkle root"); + assertEq(distributor.getProductRoundMerkleRoot(product1, 1), rootB, "Should return the correct merkle root"); + } + + function test_getProductRoundMerkleRoot_InvalidRound() public { + vm.expectRevert(Errors.MerkleDistributorNotInitialized.selector); + distributor.getProductRoundMerkleRoot(product1, 2); + } + + function test_getProductRoundMerkleRoot_ZeroAddressProduct() public { + vm.expectRevert(Errors.ZeroAddress.selector); + distributor.getProductRoundMerkleRoot(address(0), 0); + } +} diff --git a/contracts/test/AMerkleDistributor/IsClaimed.t.sol b/contracts/test/AMerkleDistributor/IsClaimed.t.sol new file mode 100644 index 0000000..8428c23 --- /dev/null +++ b/contracts/test/AMerkleDistributor/IsClaimed.t.sol @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "./AMerkleDistributorTest.sol"; + +contract IsClaimed is AMerkleDistributorTest { + bytes32 leafA1 = keccak256(abi.encodePacked(uint256(0), token1, bob, uint256(10000))); + bytes32 leafA2 = keccak256(abi.encodePacked(uint256(1), token2, alice, uint256(10000))); + bytes32 leafB1 = keccak256(abi.encodePacked(uint256(0), token3, bob, uint256(10000))); + bytes32 leafB2 = keccak256(abi.encodePacked(uint256(1), token4, alice, uint256(10000))); + bytes32 rootA = keccak256(leafA1 < leafA2 ? abi.encodePacked(leafA1, leafA2) : abi.encodePacked(leafA2, leafA1)); + bytes32 rootB = keccak256(leafB1 < leafB2 ? abi.encodePacked(leafB1, leafB2) : abi.encodePacked(leafB2, leafB1)); + + function setUp() public virtual override { + super.setUp(); + + vm.startPrank(admin); + distributor.addRoundRewards(product1, rootA); + distributor.addRoundRewards(product1, rootB); + vm.stopPrank(); + } + + function test_isClaimed_simple() public { + assertFalse(distributor.isClaimed(product1, 0, 0)); + assertFalse(distributor.isClaimed(product1, 0, 1)); + assertFalse(distributor.isClaimed(product1, 1, 0)); + assertFalse(distributor.isClaimed(product1, 1, 1)); + } + + function test_isClaimed_claimed() public { + bytes32[] memory proof = new bytes32[](1); + proof[0] = leafA2; + distributor.claim(product1, 0, 0, token1, bob, 10000, proof); + assertTrue(distributor.isClaimed(product1, 0, 0)); + assertFalse(distributor.isClaimed(product1, 0, 1)); + assertFalse(distributor.isClaimed(product1, 1, 0)); + assertFalse(distributor.isClaimed(product1, 1, 1)); + + proof[0] = leafA1; + distributor.claim(product1, 0, 1, token2, alice, 10000, proof); + assertTrue(distributor.isClaimed(product1, 0, 0)); + assertTrue(distributor.isClaimed(product1, 0, 1)); + assertFalse(distributor.isClaimed(product1, 1, 0)); + assertFalse(distributor.isClaimed(product1, 1, 1)); + + proof[0] = leafB2; + distributor.claim(product1, 1, 0, token3, bob, 10000, proof); + assertTrue(distributor.isClaimed(product1, 0, 0)); + assertTrue(distributor.isClaimed(product1, 0, 1)); + assertTrue(distributor.isClaimed(product1, 1, 0)); + assertFalse(distributor.isClaimed(product1, 1, 1)); + + proof[0] = leafB1; + distributor.claim(product1, 1, 1, token4, alice, 10000, proof); + assertTrue(distributor.isClaimed(product1, 0, 0)); + assertTrue(distributor.isClaimed(product1, 0, 1)); + assertTrue(distributor.isClaimed(product1, 1, 0)); + assertTrue(distributor.isClaimed(product1, 1, 1)); + } + + function test_isClaimed_invalidRound() public { + assertFalse(distributor.isClaimed(product1, 2, 0)); + } + + function test_isClaimed_invalidIndex() public { + assertFalse(distributor.isClaimed(product1, 0, 2)); + } +} diff --git a/contracts/test/AMerkleDistributor/MultiClaim.t.sol b/contracts/test/AMerkleDistributor/MultiClaim.t.sol new file mode 100644 index 0000000..15c10f4 --- /dev/null +++ b/contracts/test/AMerkleDistributor/MultiClaim.t.sol @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import "./AMerkleDistributorTest.sol"; +import { IMerkleDistributor } from "../../src/interfaces/IMerkleDistributor.sol"; + +contract MultiClaim is AMerkleDistributorTest { + bytes32 leafA1 = keccak256(abi.encodePacked(uint256(0), token1, bob, uint256(10000))); + bytes32 leafA2 = keccak256(abi.encodePacked(uint256(1), token2, alice, uint256(10000))); + bytes32 leafB1 = keccak256(abi.encodePacked(uint256(0), token3, bob, uint256(10000))); + bytes32 leafB2 = keccak256(abi.encodePacked(uint256(1), token4, alice, uint256(10000))); + bytes32 rootA = keccak256(leafA1 < leafA2 ? abi.encodePacked(leafA1, leafA2) : abi.encodePacked(leafA2, leafA1)); + bytes32 rootB = keccak256(leafB1 < leafB2 ? abi.encodePacked(leafB1, leafB2) : abi.encodePacked(leafB2, leafB1)); + + function setUp() public virtual override { + super.setUp(); + + vm.startPrank(admin); + distributor.addRoundRewards(product1, rootA); + distributor.addRoundRewards(product1, rootB); + vm.stopPrank(); + } + + function test_multiClaim_simple() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA2; + params[1] = IMerkleDistributor.ClaimParams({ + round: 1, + index: 0, + token: token3, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB2; + distributor.multiClaim(product1, bob, params); + assertTrue(distributor.isClaimed(product1, 0, 0)); + assertFalse(distributor.isClaimed(product1, 0, 1)); + assertTrue(distributor.isClaimed(product1, 1, 0)); + assertFalse(distributor.isClaimed(product1, 1, 1)); + } + + function test_multiClaim_InvalidRound() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA2; + params[1] = IMerkleDistributor.ClaimParams({ + round: 2, + index: 0, + token: token3, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB2; + vm.expectRevert(Errors.MerkleDistributorNotInitialized.selector); + distributor.multiClaim(product1, bob, params); + } + + function test_multiClaim_InvalidIndex() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA2; + params[1] = IMerkleDistributor.ClaimParams({ + round: 1, + index: 2, + token: token3, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB2; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.multiClaim(product1, bob, params); + } + + function test_multiClaim_InvalidToken() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA2; + params[1] = IMerkleDistributor.ClaimParams({ + round: 1, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB2; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.multiClaim(product1, bob, params); + } + + function test_multiClaim_InvalidAmount() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA2; + params[1] = IMerkleDistributor.ClaimParams({ + round: 1, + index: 0, + token: token3, + amount: 10001, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB2; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.multiClaim(product1, bob, params); + } + + function test_multiClaim_InvalidProof() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA1; + params[1] = IMerkleDistributor.ClaimParams({ + round: 1, + index: 0, + token: token3, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB1; + vm.expectRevert(Errors.InvalidMerkleProof.selector); + distributor.multiClaim(product1, bob, params); + } + + function test_multiClaim_InvalidAddress() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](2); + params[0] = IMerkleDistributor.ClaimParams({ + round: 0, + index: 0, + token: token1, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[0].merkleProof[0] = leafA2; + params[1] = IMerkleDistributor.ClaimParams({ + round: 1, + index: 0, + token: token3, + amount: 10000, + merkleProof: new bytes32[](1) + }); + params[1].merkleProof[0] = leafB2; + vm.expectRevert(Errors.ZeroAddress.selector); + distributor.multiClaim(product1, zero, params); + } + + function test_multiClaim_noParams() public { + IMerkleDistributor.ClaimParams[] memory params = new IMerkleDistributor.ClaimParams[](0); + vm.expectRevert(Errors.EmptyArray.selector); + distributor.multiClaim(product1, bob, params); + } +} diff --git a/contracts/test/mock/AMerkleBitDistributor.sol b/contracts/test/mock/AMerkleBitDistributor.sol new file mode 100644 index 0000000..0445cf1 --- /dev/null +++ b/contracts/test/mock/AMerkleBitDistributor.sol @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import { AMerkleDistributor } from "../../src/abstracts/AMerkleDistributor.sol"; +import { Owned2Step } from "../../src/utils/Owned2Step.sol"; + +contract AMerkleDistributorMock is AMerkleDistributor { + constructor(address initialOperator, address initialOwner) Owned2Step(initialOwner) AMerkleDistributor(initialOperator) { } +} diff --git a/contracts/test/mock/AMerkleDistributor.sol b/contracts/test/mock/AMerkleDistributor.sol new file mode 100644 index 0000000..0445cf1 --- /dev/null +++ b/contracts/test/mock/AMerkleDistributor.sol @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Unlicensed +pragma solidity 0.8.20; + +import { AMerkleDistributor } from "../../src/abstracts/AMerkleDistributor.sol"; +import { Owned2Step } from "../../src/utils/Owned2Step.sol"; + +contract AMerkleDistributorMock is AMerkleDistributor { + constructor(address initialOperator, address initialOwner) Owned2Step(initialOwner) AMerkleDistributor(initialOperator) { } +}