fix(contract): add reward sweeping in SpaceDelegationFacet

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/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);
     _;
   }
 

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

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

@@ -0,0 +1,130 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.23;
+
+// interfaces
+import {IOwnableBase} from "contracts/src/diamond/facets/ownable/IERC173.sol";
+import {ISpaceDelegationBase} from "contracts/src/base/registry/facets/delegation/ISpaceDelegation.sol";
+
+// libraries
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+
+// contracts
+import {SpaceDelegationFacet} from "contracts/src/base/registry/facets/delegation/SpaceDelegationFacet.sol";
+import {BaseRegistryTest} from "./BaseRegistry.t.sol";
+
+contract SpaceDelegationTest is
+  BaseRegistryTest,
+  IOwnableBase,
+  ISpaceDelegationBase
+{
+  using EnumerableSet for EnumerableSet.AddressSet;
+
+  SpaceDelegationFacet internal spaceDelegation;
+  EnumerableSet.AddressSet internal spaceSet;
+  EnumerableSet.AddressSet internal operatorSet;
+
+  function setUp() public override {
+    super.setUp();
+    spaceDelegation = SpaceDelegationFacet(baseRegistry);
+  }
+
+  /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+  /*                       ADD DELEGATION                       */
+  /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+  function test_addSpaceDelegation_revertIf_invalidSpace() public {
+    vm.expectRevert(SpaceDelegation__InvalidSpace.selector);
+    spaceDelegation.addSpaceDelegation(address(this), address(0));
+  }
+
+  function test_fuzz_addSpaceDelegation(
+    address operator
+  ) public givenOperator(operator, 0) returns (address space) {
+    space = deploySpace(deployer);
+
+    vm.prank(deployer);
+    spaceDelegation.addSpaceDelegation(space, operator);
+
+    address assignedOperator = spaceDelegation.getSpaceDelegation(space);
+    assertEq(assignedOperator, operator, "Space delegation failed");
+  }
+
+  /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+  /*                      REMOVE DELEGATION                     */
+  /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+  function test_removeSpaceDelegation_revertIf_invalidSpace() public {
+    vm.expectRevert(SpaceDelegation__InvalidSpace.selector);
+    spaceDelegation.removeSpaceDelegation(address(0));
+  }
+
+  function test_fuzz_removeSpaceDelegation(
+    address operator
+  ) public givenOperator(operator, 0) {
+    address space = test_fuzz_addSpaceDelegation(operator);
+
+    vm.prank(deployer);
+    spaceDelegation.removeSpaceDelegation(space);
+
+    address afterRemovalOperator = spaceDelegation.getSpaceDelegation(space);
+    assertEq(afterRemovalOperator, address(0), "Space removal failed");
+  }
+
+  /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+  /*                           GETTERS                          */
+  /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+  function test_fuzz_getSpaceDelegationsByOperator(address operator) public {
+    address space1 = test_fuzz_addSpaceDelegation(operator);
+    address space2 = test_fuzz_addSpaceDelegation(operator);
+
+    address[] memory spaces = spaceDelegation.getSpaceDelegationsByOperator(
+      operator
+    );
+
+    assertEq(spaces.length, 2);
+    assertEq(spaces[0], space1);
+    assertEq(spaces[1], space2);
+  }
+
+  /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+  /*                           SETTERS                          */
+  /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+  function test_setRiverToken_revertIf_notOwner() public {
+    vm.expectRevert(
+      abi.encodeWithSelector(Ownable__NotOwner.selector, address(this))
+    );
+    spaceDelegation.setRiverToken(address(0));
+  }
+
+  function test_fuzz_setRiverToken(address newToken) public {
+    vm.assume(newToken != address(0));
+
+    vm.expectEmit(address(spaceDelegation));
+    emit RiverTokenChanged(newToken);
+
+    vm.prank(deployer);
+    spaceDelegation.setRiverToken(newToken);
+
+    address retrievedToken = spaceDelegation.riverToken();
+    assertEq(retrievedToken, newToken);
+  }
+
+  function test_fuzz_setSpaceFactory_revertIf_notOwner() public {
+    vm.expectRevert(
+      abi.encodeWithSelector(Ownable__NotOwner.selector, address(this))
+    );
+    spaceDelegation.setSpaceFactory(address(0));
+  }
+
+  function test_fuzz_setSpaceFactory(address newSpaceFactory) public {
+    vm.assume(newSpaceFactory != address(0));
+
+    vm.prank(deployer);
+    spaceDelegation.setSpaceFactory(newSpaceFactory);
+
+    address retrievedFactory = spaceDelegation.getSpaceFactory();
+    assertEq(retrievedFactory, newSpaceFactory);
+  }
+}