Skip to content

Commit

Permalink
feat(contracts): added MerkleDistributor
Browse files Browse the repository at this point in the history
  • Loading branch information
0xmemorygrinder committed Dec 11, 2023
1 parent 7e95f52 commit 1782452
Show file tree
Hide file tree
Showing 13 changed files with 763 additions and 0 deletions.
137 changes: 137 additions & 0 deletions contracts/src/abstracts/AMerkleDistributor.sol
Original file line number Diff line number Diff line change
@@ -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; }
}
}
}
37 changes: 37 additions & 0 deletions contracts/src/interfaces/IMerkleDistributor.sol
Original file line number Diff line number Diff line change
@@ -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;
}
5 changes: 5 additions & 0 deletions contracts/src/utils/Errors.sol
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,9 @@ library Errors {
// Swapper errors
error SwapError();
error NotVault();

// MerkleDistributor errors
error MerkleDistributorNotInitialized();
error InvalidMerkleProof();
error AlreadyClaimed();
}
76 changes: 76 additions & 0 deletions contracts/src/utils/MerkleRootBuilder.sol
Original file line number Diff line number Diff line change
@@ -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
}
}
}
29 changes: 29 additions & 0 deletions contracts/test/AMerkleDistributor/AMerkleDistributorTest.sol
Original file line number Diff line number Diff line change
@@ -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();
}
}
59 changes: 59 additions & 0 deletions contracts/test/AMerkleDistributor/AddRoundRewards.t.sol
Original file line number Diff line number Diff line change
@@ -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");
}
}
Loading

0 comments on commit 1782452

Please sign in to comment.