feat: add claimable

script/Deploy.s.sol

@@ -75,7 +75,7 @@ contract Deploy is BaseDeploy {
 
     /// @notice DEMO only - Helper function to create an epoch
     function _createEpoch(uint256 id, uint256 startTimestamp, string memory description) private {
-        DripVault vault = new DripVault(MOCK_TOKEN);
+        DripVault vault = new DripVault(MOCK_TOKEN, address(challengeManager));
         Types.Epoch memory epoch = Types.Epoch({
             id: id,
             startTimestamp: startTimestamp,

src/Challenge.sol

@@ -74,9 +74,7 @@ contract Challenge is ERC721, IChallenge {
         emit ChallengeCreated(_tokenId, params.owner);
 
         // Add the challenge to the epoch
-        IChallengeManager(managerContract).addChallengeToEpoch(
-            params.epochId, _tokenId, params.owner, params.depositAmount
-        );
+        IChallengeManager(managerContract).addChallengeToEpoch(params.epochId, challenge);
 
         // Return challenge ID and increment the token counter
         challengeId = _tokenId;
@@ -112,6 +110,12 @@ contract Challenge is ERC721, IChallenge {
         require(block.timestamp >= dayStartTime && block.timestamp <= dayEndTime, ErrorsLib.NOT_IN_TIME_RANGE);
 
         challenge.dailyCompletionTimestamps[day] = block.timestamp;
+
+        // Update vault
+        (Types.Epoch memory epoch,) = IChallengeManager(managerContract).getEpochInfo(challenge.epochId);
+        DripVault vault = DripVault(epoch.vault);
+        vault.submitDailyCompletion(owner, challenge.dailyCompletionTimestamps);
+
         emit DailyCompletionSubmitted(tokenId, day);
     }
 }

src/ChallengeManager.sol

@@ -81,7 +81,7 @@ contract ChallengeManager is Ownable, IChallengeManager {
 
         uint256 endTimestamp = TimeLib.addDaysToTimestamp(start, durationInDays);
 
-        DripVault vault = new DripVault(asset);
+        DripVault vault = new DripVault(asset, address(this));
 
         Types.Epoch memory epoch = Types.Epoch({
             id: id,
@@ -100,17 +100,15 @@ contract ChallengeManager is Ownable, IChallengeManager {
         emit EpochStarted(id, start, endTimestamp);
     }
 
-    function addChallengeToEpoch(uint256 epochId, uint256 challengeId, address challengeOwner, uint256 depositAmount)
-        external
-        onlyChallenge
-    {
+    function addChallengeToEpoch(uint256 epochId, Types.Challenge calldata userChallenge) external onlyChallenge {
         Types.ChallengeManagerStorage storage $ = _getChallengeManagerStorage();
         IERC20($.epochs[epochId].asset).approve($.epochs[epochId].vault, type(uint256).max);
-        IERC4626($.epochs[epochId].vault).deposit(depositAmount, challengeOwner);
+        IERC4626($.epochs[epochId].vault).deposit(userChallenge.depositAmount, userChallenge.owner);
+        DripVault($.epochs[epochId].vault).setOwnerChallengeDays(userChallenge.owner, userChallenge.durationInDays);
 
-        $.epochChallenges[epochId].push(challengeId);
+        $.epochChallenges[epochId].push(userChallenge.id);
         EnumerableSet.AddressSet storage participants = $.epochParticipants[epochId];
-        participants.add(challengeOwner);
+        participants.add(userChallenge.owner);
         // Update epoch info
         $.epochs[epochId].totalDeposits = IERC4626($.epochs[epochId].vault).totalAssets();
         $.epochs[epochId].participantCount = participants.length();
@@ -137,6 +135,26 @@ contract ChallengeManager is Ownable, IChallengeManager {
         $.nextEpochId++;
     }
 
+    /**
+     * @notice Preview rewards in an epoch
+     */
+    function previewClaimRewards(address owner, uint256 epochId) public view returns (uint256) {
+        Types.Epoch memory epoch = _getChallengeManagerStorage().epochs[epochId];
+        DripVault vault = DripVault(epoch.vault);
+        return vault.previewClaim(owner);
+    }
+
+    /**
+     * @notice Claim rewards in an epoch
+     */
+    function claimRewards(address owner, uint256 epochId) external {
+        require(owner == msg.sender, ErrorsLib.NOT_AUTHORIZED);
+        Types.Epoch memory epoch = _getChallengeManagerStorage().epochs[epochId];
+        require(getEpochStatus(epochId) == Types.EpochStatus.Ended, ErrorsLib.INVALID_EPOCH_STATUS);
+        DripVault vault = DripVault(epoch.vault);
+        vault.claim(owner);
+    }
+
     /// @dev Returns the storage struct of ChallengeManager.
     function _getChallengeManagerStorage() private pure returns (Types.ChallengeManagerStorage storage $) {
         assembly {

src/DripVault.sol

@@ -2,7 +2,127 @@
 pragma solidity ^0.8.28;
 
 import {ERC4626, IERC20, ERC20} from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
+import {Types} from "./libraries/Types.sol";
+import {console} from "forge-std/console.sol";
 
 contract DripVault is ERC4626 {
-    constructor(address asset) ERC4626(IERC20(asset)) ERC20("Drip Vault Token", "DRIPVLT") {}
+    using EnumerableSet for EnumerableSet.AddressSet;
+    using Math for uint256;
+
+    EnumerableSet.AddressSet private owners;
+    // Owner => Original deposits
+    mapping(address => uint256) public ownerDeposits;
+    // Owner => Total challenge days
+    mapping(address => uint256) public ownerChallengeDays;
+    // Owner => Daily completions count
+    mapping(address => uint256) public ownerDailyCompletions;
+    // Owner => Claimable amount
+    mapping(address => uint256) public ownerClaimables;
+
+    address public challengeManager;
+    uint256 public totalDailyCompletionCount;
+
+    modifier onlyChallengeManager() {
+        require(msg.sender == challengeManager, "Only challenge managers can call this function");
+        _;
+    }
+
+    constructor(address asset, address _challengeManager)
+        ERC4626(IERC20(asset))
+        ERC20("Drip Vault Token", "DRIP-VLT")
+    {
+        challengeManager = _challengeManager;
+    }
+
+    function deposit(uint256 assets, address receiver) public override returns (uint256) {
+        ownerDeposits[receiver] += assets;
+        return super.deposit(assets, receiver);
+    }
+
+    function setOwnerChallengeDays(address owner, uint256 challengeDays) external {
+        ownerChallengeDays[owner] = challengeDays;
+    }
+
+    function previewClaim(address owner) external view returns (uint256) {
+        return ownerClaimables[owner];
+    }
+
+    function claim(address owner) external onlyChallengeManager {
+        uint256 claimable = ownerClaimables[owner];
+        require(claimable > 0, "No claimable amount");
+        uint256 userShares = IERC20(this).balanceOf(owner);
+        require(userShares > 0, "No shares to burn");
+        _burn(owner, userShares);
+        ownerClaimables[owner] = 0;
+        ERC20(asset()).transfer(owner, claimable);
+    }
+
+    function submitDailyCompletion(address owner, uint256[] calldata dailyCompletion) external {
+        uint256 completedDays = 0;
+        for (uint256 i = 0; i < dailyCompletion.length;) {
+            if (dailyCompletion[i] > 0) {
+                completedDays++;
+            }
+            unchecked {
+                i++;
+            }
+        }
+        owners.add(owner);
+        ownerDailyCompletions[owner] = completedDays;
+        totalDailyCompletionCount++;
+    }
+
+    function calculateOwnersClaimable() public {
+        uint256 ownerCount = owners.length(); // Cache length to avoid repeated calls
+        require(ownerCount > 0, "No owners to calculate");
+
+        uint256 totalClaimableDeposits = 0;
+        uint256 totalDeposits = totalAssets();
+        require(totalDeposits > 0, "Total deposits must be greater than zero");
+
+        // Cache remaining rewards calculation to reduce computations
+        uint256[] memory claimableDeposits = new uint256[](ownerCount);
+
+        // First pass: Calculate claimable deposits
+        for (uint256 i; i < ownerCount;) {
+            address owner = owners.at(i);
+
+            uint256 userDeposits = ownerDeposits[owner];
+            uint256 completedDays = ownerDailyCompletions[owner];
+            uint256 totalDays = ownerChallengeDays[owner];
+
+            uint256 claimable = 0;
+            if (totalDays > 0 && completedDays > 0) {
+                claimable = Math.mulDiv(userDeposits, completedDays, totalDays);
+            }
+
+            claimableDeposits[i] = claimable; // Cache claimable deposit for later
+            ownerClaimables[owner] = claimable; // Update state
+            totalClaimableDeposits += claimable;
+            unchecked {
+                i++;
+            }
+        }
+
+        // Calculate remaining rewards once
+        if (totalClaimableDeposits < totalDeposits) {
+            uint256 remainingRewards = totalDeposits - totalClaimableDeposits;
+
+            // Second pass: Add additional rewards
+            for (uint256 i; i < ownerCount;) {
+                address owner = owners.at(i);
+                uint256 userDeposits = ownerDeposits[owner];
+
+                if (userDeposits > 0) {
+                    uint256 additionalRewards = Math.mulDiv(userDeposits, remainingRewards, totalDeposits);
+                    ownerClaimables[owner] += additionalRewards; // Add to claimable
+                }
+                unchecked {
+                    i++;
+                }
+            }
+        }
+    }
 }

src/interfaces/IChallengeManager.sol

@@ -35,8 +35,7 @@ interface IChallengeManager {
     /**
      * @dev Adds a challenge to an epoch
      */
-    function addChallengeToEpoch(uint256 epochId, uint256 challengeId, address challengeOwner, uint256 depositAmount)
-        external;
+    function addChallengeToEpoch(uint256 epochId, Types.Challenge calldata challenge) external;
 
     /**
      * @dev Returns the challenges for an epoch
@@ -47,4 +46,14 @@ interface IChallengeManager {
      * @dev Returns if a token is whitelisted for depositing when creating a challenge
      */
     function isWhiteListedToken(address token) external view returns (bool isWhitelisted);
+
+    /**
+     * @dev Preview rewards in an epoch
+     */
+    function previewClaimRewards(address owner, uint256 epochId) external view returns (uint256);
+
+    /**
+     * @dev Claim rewards in an epoch
+     */
+    function claimRewards(address owner, uint256 epochId) external;
 }

src/libraries/Types.sol

@@ -84,11 +84,4 @@ library Types {
         // Token address => is whitelisted
         mapping(address => bool) isWhitelistedToken;
     }
-
-    struct VaultStorage {
-        // Epoch ID => total deposits
-        mapping(uint256 => uint256) epochTotalDeposits;
-        // Epoch ID => failed deposits
-        mapping(uint256 => uint256) epochFailedDeposits;
-    }
 }

test/ChallengeManager.t.sol

@@ -2,6 +2,9 @@
 pragma solidity ^0.8.28;
 
 import "./helpers/SetUp.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+import {DripVault} from "../src/DripVault.sol";
 import {ErrorsLib} from "../src/libraries/ErrorsLib.sol";
 import {Types} from "../src/libraries/Types.sol";
 
@@ -73,4 +76,97 @@ contract ChallengeManagerTest is SetUp {
         status = uint256(challengeManager.getEpochStatus(0));
         assertEq(status, uint256(Types.EpochStatus.Closed));
     }
+
+    function testPreviewRewards() public {
+        uint256 epochId = _createDummyEpoch();
+        (Types.Epoch memory epoch,) = challengeManager.getEpochInfo(epochId);
+        (uint256 profileId1, uint256 profileId2) = _createDummyProfiles();
+        DripVault vault = DripVault(epoch.vault);
+
+        // User1 creates a challenge by depositing 10 tokens
+        vm.startPrank(USER1);
+        deal(address(mockToken), USER1, 100);
+        _createDummyChallenge(profileId1, 10, 5);
+        assertEq(vault.ownerDeposits(USER1), 10, "User1 deposited 10 tokens");
+        assertEq(IERC20(address(vault)).balanceOf(USER1), 10, "User1 got 10 DRIP-VLT");
+        assertEq(IERC20(mockToken).balanceOf(USER1), 90, "User1 has 90 tokens left");
+        vm.stopPrank();
+
+        // User2 creates a challenge by depositing 15 tokens
+        vm.startPrank(USER2);
+        deal(address(mockToken), USER2, 100);
+        _createDummyChallenge(profileId2, 15, 3);
+        assertEq(vault.ownerDeposits(USER2), 15, "User2 deposited 15 tokens");
+        assertEq(IERC20(address(vault)).balanceOf(USER2), 15, "User2 got 15 DRIP-VLT");
+        assertEq(IERC20(mockToken).balanceOf(USER2), 85, "User2 has 85 tokens left");
+        vm.stopPrank();
+
+        // Total deposits
+        assertEq(vault.totalAssets(), 25, "Total deposits are 25 tokens");
+
+        _submitDummyDailyCompletion();
+        vm.warp(432002);
+
+        // Calculate claimable rewards
+        vault.calculateOwnersClaimable();
+
+        vm.startPrank(USER1);
+
+        assertEq(challengeManager.previewClaimRewards(USER1, epochId), 7, "User1 claimable rewards");
+        challengeManager.claimRewards(USER1, epochId);
+        assertEq(IERC20(address(vault)).balanceOf(USER1), 0, "User1 has 0 DRIP-VLT");
+        assertEq(IERC20(mockToken).balanceOf(USER1), 97, "User1 has 97 tokens");
+        vm.stopPrank();
+
+        vm.startPrank(USER2);
+        assertEq(challengeManager.previewClaimRewards(USER2, epochId), 17, "User2 claimable rewards");
+        challengeManager.claimRewards(USER2, epochId);
+        assertEq(IERC20(address(vault)).balanceOf(USER2), 0, "User2 has 0 DRIP-VLT");
+        assertEq(IERC20(mockToken).balanceOf(USER2), 102, "User2 has 102 tokens");
+        vm.stopPrank();
+    }
+
+    function testClaimRewards_EpochNotEnded() public {
+        uint256 epochId = _createDummyEpoch();
+
+        vm.prank(USER1);
+        vm.expectRevert(bytes(ErrorsLib.INVALID_EPOCH_STATUS));
+        challengeManager.claimRewards(USER1, epochId);
+    }
+
+    function _createDummyChallenge(uint256 profileId, uint256 amount, uint16 durationDays) private returns (uint256) {
+        mockToken.approve(address(challenge), amount);
+        return profile.createChallenge(profileId, "title", "desc", amount, durationDays);
+    }
+
+    function _submitDummyDailyCompletion() private {
+        uint256 user1ChallengeId = 0;
+        uint256 user2ChallengeId = 1;
+
+        // 1st day (1 - 86401)
+        vm.prank(USER1);
+        challenge.submitDailyCompletion(USER1, user1ChallengeId, 0);
+        vm.prank(USER2);
+        challenge.submitDailyCompletion(USER2, user2ChallengeId, 0);
+
+        // 2nd day (86401 - 172801)
+        vm.warp(86401);
+        vm.prank(USER1);
+        challenge.submitDailyCompletion(USER1, user1ChallengeId, 1);
+        vm.prank(USER2);
+        challenge.submitDailyCompletion(USER2, user2ChallengeId, 1);
+
+        // 3rd day (172801 - 259201)
+        vm.warp(172801);
+        vm.prank(USER2);
+        challenge.submitDailyCompletion(USER2, user2ChallengeId, 2);
+
+        // 4th day (259201 - 345601)
+        // No submission
+
+        // 5th day (345601 - 432001)
+        vm.warp(345601);
+        vm.prank(USER1);
+        challenge.submitDailyCompletion(USER1, user1ChallengeId, 4);
+    }
 }