feat(contracts): added MerkleDistributor

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; }
+        }
+    }
+}

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;
+}

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();
 }

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
+        }
+     }
+}

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();
+    }
+}

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");
+    }
+}