From fbd95e5604f79a7523f72e31b2297fc404e571fb Mon Sep 17 00:00:00 2001
From: Liron Achdut <liron.achdut@gmail.com>
Date: Mon, 5 May 2025 16:34:20 +0300
Subject: [PATCH] feat: calculateReward to return 0 if claimed

---
 src/Hyperfund.sol      |  6 +++++-
 src/Hyperstaker.sol    | 29 ++++++++++++++++++++---------
 test/Hyperfund.t.sol   |  2 ++
 test/Hyperstaker.t.sol |  8 ++++++++
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/Hyperfund.sol b/src/Hyperfund.sol
index 39b6c6f..c83d955 100644
--- a/src/Hyperfund.sol
+++ b/src/Hyperfund.sol
@@ -112,7 +112,11 @@ contract Hyperfund is AccessControlUpgradeable, PausableUpgradeable, UUPSUpgrade
     /// @notice Issue a Hypercert fraction for a non-financial contribution, only callable by a manager
     /// @param _contributor address of the contributor to receive the Hypercert fraction
     /// @param _units amount of units to register as a non-financial contribution
-    function nonfinancialContribution(address _contributor, uint256 _units) external whenNotPaused onlyRole(MANAGER_ROLE) {
+    function nonfinancialContribution(address _contributor, uint256 _units)
+        external
+        whenNotPaused
+        onlyRole(MANAGER_ROLE)
+    {
         uint256 availableSupply = hypercertMinter.unitsOf(hypercertId);
         require(availableSupply >= _units, AmountExceedsAvailableSupply(availableSupply));
         _nonfinancialContribution(_contributor, _units);
diff --git a/src/Hyperstaker.sol b/src/Hyperstaker.sol
index 4f72a19..017d3e2 100644
--- a/src/Hyperstaker.sol
+++ b/src/Hyperstaker.sol
@@ -162,7 +162,7 @@ contract Hyperstaker is AccessControlUpgradeable, PausableUpgradeable, UUPSUpgra
         address staker = stakes[_hypercertId].staker;
         require(staker == msg.sender, NotStakerOfHypercert(staker));
         require(!isRoundClaimed(_hypercertId, _roundId), AlreadyClaimed());
-        uint256 reward = calculateReward(_hypercertId, _roundId);
+        uint256 reward = _calculateReward(_hypercertId, _roundId);
         require(reward != 0, NoRewardAvailable());
 
         _setRoundClaimed(_hypercertId, _roundId);
@@ -184,14 +184,10 @@ contract Hyperstaker is AccessControlUpgradeable, PausableUpgradeable, UUPSUpgra
     /// @param _roundId id of the round to calculate the reward for
     /// @return amount of the reward eligable for the staked Hypercert for the given round
     function calculateReward(uint256 _hypercertId, uint256 _roundId) public view returns (uint256) {
-        Round memory round = rounds[_roundId];
-        require(round.endTime != 0, RoundNotSet());
-        uint256 stakeStartTime = stakes[_hypercertId].stakingStartTime;
-        require(stakeStartTime != 0, NotStaked());
-        stakeStartTime = stakeStartTime < round.startTime ? round.startTime : stakeStartTime;
-        uint256 stakeDuration = stakeStartTime > round.endTime ? 0 : round.endTime - stakeStartTime;
-        return
-            round.totalRewards * hypercertMinter.unitsOf(_hypercertId) * stakeDuration / (totalUnits * round.duration);
+        if (isRoundClaimed(_hypercertId, _roundId)) {
+            return 0;
+        }
+        return _calculateReward(_hypercertId, _roundId);
     }
 
     /// @notice Check if a staked Hypercert had already claimed a reward for a given round
@@ -218,6 +214,21 @@ contract Hyperstaker is AccessControlUpgradeable, PausableUpgradeable, UUPSUpgra
 
     // INTERNAL FUNCTIONS
 
+    /// @notice Calculate the reward for a staked Hypercert for a given round
+    /// @param _hypercertId id of the Hypercert to calculate the reward for
+    /// @param _roundId id of the round to calculate the reward for
+    /// @return amount of the reward eligable for the staked Hypercert for the given round
+    function _calculateReward(uint256 _hypercertId, uint256 _roundId) internal view returns (uint256) {
+        Round memory round = rounds[_roundId];
+        require(round.endTime != 0, RoundNotSet());
+        uint256 stakeStartTime = stakes[_hypercertId].stakingStartTime;
+        require(stakeStartTime != 0, NotStaked());
+        stakeStartTime = stakeStartTime < round.startTime ? round.startTime : stakeStartTime;
+        uint256 stakeDuration = stakeStartTime > round.endTime ? 0 : round.endTime - stakeStartTime;
+        return
+            round.totalRewards * hypercertMinter.unitsOf(_hypercertId) * stakeDuration / (totalUnits * round.duration);
+    }
+
     /// @notice Get the hypercert type id for a given hypercert id
     /// @param _hypercertId id of the Hypercert to get the type id for
     /// @return hypercert type id for the given hypercert id
diff --git a/test/Hyperfund.t.sol b/test/Hyperfund.t.sol
index 30d2373..217536b 100644
--- a/test/Hyperfund.t.sol
+++ b/test/Hyperfund.t.sol
@@ -364,6 +364,8 @@ contract HyperfundTest is Test {
     }
 
     function test_RevertWhen_RedeemNotFraction() public {
+        vm.prank(manager);
+        hyperfund.nonfinancialContribution(contributor, amount);
         vm.startPrank(contributor);
         vm.recordLogs();
         hypercertMinter.mintClaim(contributor, totalUnits, "uri", IHypercertToken.TransferRestrictions.AllowAll);
diff --git a/test/Hyperstaker.t.sol b/test/Hyperstaker.t.sol
index 8d321c0..7e1945f 100644
--- a/test/Hyperstaker.t.sol
+++ b/test/Hyperstaker.t.sol
@@ -225,6 +225,14 @@ contract HyperstakerTest is Test {
         assertEq(hyperstaker.calculateReward(stakerHypercertId, 0), 0);
     }
 
+    function test_CalculateReward_0WhenRoundClaimed() public {
+        _setupStake();
+        _setupRewardEth();
+        vm.prank(staker);
+        hyperstaker.claimReward(stakerHypercertId, 0);
+        assertEq(hyperstaker.calculateReward(stakerHypercertId, 0), 0);
+    }
+
     function test_RevertWhen_CalculateRewardNoReward() public {
         _setupStake();
         vm.expectRevert(RoundNotSet.selector);