fix(contract): add reward sweeping in SpaceDelegationFacet (#1435)

Introduce reward sweeping when updating space operators. This ensures
that any unclaimed rewards are properly transferred to the current
operator. Additionally, refactored validations and removed redundant
checks for enhanced readability and maintainability.

contracts/scripts/deployments/facets/DeploySpaceDelegation.s.sol

@@ -17,6 +17,7 @@ contract DeploySpaceDelegation is Deployer, FacetHelper {
     addSelector(SpaceDelegationFacet.getSpaceDelegation.selector);
     addSelector(SpaceDelegationFacet.getSpaceDelegationsByOperator.selector);
     addSelector(SpaceDelegationFacet.setRiverToken.selector);
+    addSelector(SpaceDelegationFacet.riverToken.selector);
     addSelector(SpaceDelegationFacet.getTotalDelegation.selector);
     addSelector(SpaceDelegationFacet.setMainnetDelegation.selector);
     addSelector(SpaceDelegationFacet.setSpaceFactory.selector);

contracts/src/base/registry/facets/delegation/SpaceDelegationFacet.sol

@@ -2,26 +2,34 @@
 pragma solidity ^0.8.23;
 
 // interfaces
-
 import {ISpaceDelegation} from "contracts/src/base/registry/facets/delegation/ISpaceDelegation.sol";
 import {IMainnetDelegation} from "contracts/src/tokens/river/base/delegation/IMainnetDelegation.sol";
 import {IERC173} from "contracts/src/diamond/facets/ownable/IERC173.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {IVotesEnumerable} from "contracts/src/diamond/facets/governance/votes/enumerable/IVotesEnumerable.sol";
 import {IArchitect} from "contracts/src/factory/facets/architect/IArchitect.sol";
+import {IRewardsDistributionBase} from "contracts/src/base/registry/facets/distribution/v2/IRewardsDistribution.sol";
 
 // libraries
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {SpaceDelegationStorage} from "contracts/src/base/registry/facets/delegation/SpaceDelegationStorage.sol";
 import {CustomRevert} from "contracts/src/utils/libraries/CustomRevert.sol";
 import {NodeOperatorStorage, NodeOperatorStatus} from "contracts/src/base/registry/facets/operator/NodeOperatorStorage.sol";
+import {StakingRewards} from "contracts/src/base/registry/facets/distribution/v2/StakingRewards.sol";
+import {RewardsDistributionStorage} from "contracts/src/base/registry/facets/distribution/v2/RewardsDistributionStorage.sol";
 
 // contracts
 import {OwnableBase} from "contracts/src/diamond/facets/ownable/OwnableBase.sol";
 import {Facet} from "contracts/src/diamond/facets/Facet.sol";
 
-contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
+contract SpaceDelegationFacet is
+  ISpaceDelegation,
+  IRewardsDistributionBase,
+  OwnableBase,
+  Facet
+{
   using EnumerableSet for EnumerableSet.AddressSet;
+  using StakingRewards for StakingRewards.Layout;
 
   function __SpaceDelegation_init(
     address riverToken_
@@ -42,16 +50,14 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
     address space,
     address operator
   ) external onlySpaceOwner(space) {
-    if (space == address(0))
-      CustomRevert.revertWith(SpaceDelegation__InvalidAddress.selector);
     if (operator == address(0))
       CustomRevert.revertWith(SpaceDelegation__InvalidAddress.selector);
 
     SpaceDelegationStorage.Layout storage ds = SpaceDelegationStorage.layout();
 
     address currentOperator = ds.operatorBySpace[space];
 
-    if (currentOperator != address(0) && currentOperator == operator)
+    if (currentOperator == operator)
       CustomRevert.revertWith(SpaceDelegation__AlreadyDelegated.selector);
 
     NodeOperatorStorage.Layout storage nodeOperatorDs = NodeOperatorStorage
@@ -65,12 +71,14 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
     if (nodeOperatorDs.statusByOperator[operator] == NodeOperatorStatus.Exiting)
       CustomRevert.revertWith(SpaceDelegation__InvalidOperator.selector);
 
-    //remove the space from the current operator
+    _sweepSpaceRewardsIfNecessary(space, currentOperator);
+
+    // remove the space from the current operator
     ds.spacesByOperator[currentOperator].remove(space);
 
-    //overwrite the operator for this space
+    // overwrite the operator for this space
     ds.operatorBySpace[space] = operator;
-    //add the space to this new operator array
+    // add the space to this new operator array
     ds.spacesByOperator[operator].add(space);
     ds.spaceDelegationTime[space] = block.timestamp;
 
@@ -90,6 +98,8 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
       CustomRevert.revertWith(SpaceDelegation__InvalidAddress.selector);
     }
 
+    _sweepSpaceRewardsIfNecessary(space, operator);
+
     ds.operatorBySpace[space] = address(0);
     ds.spacesByOperator[operator].remove(space);
     ds.spaceDelegationTime[space] = 0;
@@ -114,6 +124,7 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
   // =============================================================
   //                           Token
   // =============================================================
+
   /// @inheritdoc ISpaceDelegation
   function setRiverToken(address newToken) external onlyOwner {
     if (newToken == address(0))
@@ -134,6 +145,7 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
   // =============================================================
   //                      Mainnet Delegation
   // =============================================================
+
   /// @inheritdoc ISpaceDelegation
   function setMainnetDelegation(address newDelegation) external onlyOwner {
     if (newDelegation == address(0))
@@ -151,6 +163,7 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
   // =============================================================
   //                           Stake
   // =============================================================
+
   /// @inheritdoc ISpaceDelegation
   function getTotalDelegation(
     address operator
@@ -177,6 +190,7 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
   // =============================================================
   //                           Space Factory
   // =============================================================
+
   /// @inheritdoc ISpaceDelegation
   function setSpaceFactory(address spaceFactory) external onlyOwner {
     if (spaceFactory == address(0))
@@ -195,6 +209,34 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
   //                           Internal
   // =============================================================
 
+  /// @dev Sweeps the rewards in the space delegation to the operator if necessary
+  function _sweepSpaceRewardsIfNecessary(
+    address space,
+    address currentOperator
+  ) internal {
+    StakingRewards.Layout storage staking = RewardsDistributionStorage
+      .layout()
+      .staking;
+    StakingRewards.Treasure storage spaceTreasure = staking
+      .treasureByBeneficiary[space];
+
+    staking.updateGlobalReward();
+    staking.updateReward(spaceTreasure);
+
+    uint256 reward = spaceTreasure.unclaimedRewardSnapshot;
+    if (reward == 0) return;
+
+    // forfeit the rewards if the space has undelegated
+    if (currentOperator != address(0)) {
+      StakingRewards.Treasure storage operatorTreasure = staking
+        .treasureByBeneficiary[currentOperator];
+      operatorTreasure.unclaimedRewardSnapshot += reward;
+    }
+    spaceTreasure.unclaimedRewardSnapshot = 0;
+
+    emit SpaceRewardsSwept(space, currentOperator, reward);
+  }
+
   function _getTotalDelegation(
     address operator
   ) internal view returns (uint256) {
@@ -232,7 +274,7 @@ contract SpaceDelegationFacet is ISpaceDelegation, OwnableBase, Facet {
 
   function _isValidSpaceOwner(address space) internal view returns (bool) {
     return
-      IERC173(space).owner() == msg.sender &&
-      IArchitect(getSpaceFactory()).getTokenIdBySpace(space) > 0;
+      IArchitect(getSpaceFactory()).getTokenIdBySpace(space) > 0 &&
+      IERC173(space).owner() == msg.sender;
   }
 }

contracts/src/base/registry/facets/distribution/v2/RewardsDistribution.sol

@@ -304,7 +304,6 @@ contract RewardsDistribution is
   }
 
   /// @inheritdoc IRewardsDistribution
-  // TODO: transfer rewards when a space redelegates
   function claimReward(
     address beneficiary,
     address recipient

contracts/src/base/registry/facets/distribution/v2/RewardsDistributionBase.sol

@@ -112,7 +112,8 @@ abstract contract RewardsDistributionBase is IRewardsDistributionBase {
   /// @dev Sweeps the rewards in the space delegation to the operator if necessary
   /// @dev Must be called after `StakingRewards.updateGlobalReward`
   function _sweepSpaceRewardsIfNecessary(address space) internal {
-    if (!_isSpace(space)) return;
+    address operator = _getOperatorBySpace(space);
+    if (operator == address(0)) return;
 
     StakingRewards.Layout storage staking = RewardsDistributionStorage
       .layout()
@@ -124,7 +125,6 @@ abstract contract RewardsDistributionBase is IRewardsDistributionBase {
     uint256 scaledReward = spaceTreasure.unclaimedRewardSnapshot;
     if (scaledReward == 0) return;
 
-    address operator = _getOperatorBySpace(space);
     StakingRewards.Treasure storage operatorTreasure = staking
       .treasureByBeneficiary[operator];
 

contracts/test/base/registry/BaseRegistry.t.sol

@@ -125,18 +125,18 @@ abstract contract BaseRegistryTest is BaseSetup, IRewardsDistributionBase {
   /*                           SPACE                            */
   /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
-  function deploySpace() internal returns (address _space) {
+  function deploySpace(address _deployer) internal returns (address _space) {
     IArchitectBase.SpaceInfo memory spaceInfo = _createSpaceInfo(
       string(abi.encode(_randomUint256()))
     );
     spaceInfo.membership.settings.pricingModule = pricingModule;
-    vm.prank(deployer);
+    vm.prank(_deployer);
     _space = ICreateSpace(spaceFactory).createSpace(spaceInfo);
     space = _space;
   }
 
   modifier givenSpaceIsDeployed() {
-    deploySpace();
+    deploySpace(deployer);
     _;
   }
 
@@ -342,4 +342,31 @@ abstract contract BaseRegistryTest is BaseSetup, IRewardsDistributionBase {
       "expected reward"
     );
   }
+
+  function verifySweep(
+    address space,
+    address operator,
+    uint256 amount,
+    uint256 commissionRate,
+    uint256 timeLapse
+  ) internal view {
+    StakingState memory state = rewardsDistributionFacet.stakingState();
+    StakingRewards.Treasure memory spaceTreasure = rewardsDistributionFacet
+      .treasureByBeneficiary(space);
+
+    assertEq(spaceTreasure.earningPower, (amount * commissionRate) / 10000);
+    assertEq(
+      spaceTreasure.rewardPerTokenAccumulated,
+      state.rewardPerTokenAccumulated
+    );
+    assertEq(spaceTreasure.unclaimedRewardSnapshot, 0);
+
+    assertEq(
+      rewardsDistributionFacet
+        .treasureByBeneficiary(operator)
+        .unclaimedRewardSnapshot,
+      spaceTreasure.earningPower *
+        state.rewardRate.fullMulDiv(timeLapse, state.totalStaked)
+    );
+  }
 }

contracts/test/base/registry/RewardsDistributionV2.t.sol

@@ -1076,7 +1076,7 @@ contract RewardsDistributionV2Test is
     bridgeTokensForUser(address(this), 1 ether * uint256(count));
     river.approve(address(rewardsDistributionFacet), type(uint256).max);
     for (uint256 i; i < count; ++i) {
-      address _space = deploySpace();
+      address _space = deploySpace(deployer);
       pointSpaceToOperator(_space, OPERATOR);
       rewardsDistributionFacet.stake(1 ether, _space, address(this));
     }