update RewardsDistributor

src/RewardsDistributor.sol

@@ -76,16 +76,29 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
         address token,
         uint256 emissionRate
     ) external onlyOwner {
+        require(receiver != address(0), "Invalid receiver");
         require(token != address(0), "No native rewards allowed");
-        require(emissionRate > 0, "Emission rate must be greater than 0");
+
+        if (rewardConfigurations[receiver][token].emissionRate == 0) {
+            rewardTokens[receiver].push(token);
+        }
 
         rewardConfigurations[receiver][token] = RewardConfiguration(
             emissionRate,
             block.timestamp
         );
 
-        // TODO check if token already exists
-        rewardTokens[receiver].push(token);
+        if (emissionRate == 0) {
+            // remove the token from the list
+            address[] storage tokens = rewardTokens[receiver];
+            for (uint256 i = 0; i < tokens.length; i++) {
+                if (tokens[i] == token) {
+                    tokens[i] = tokens[tokens.length - 1];
+                    tokens.pop();
+                    break;
+                }
+            }
+        }
 
         emit RewardConfigurationSet(receiver, token, emissionRate);
     }

src/Staking.sol

@@ -121,12 +121,27 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         // Distribute rewards
         rewardsDistributor.distributeRewards();
 
-        address[] rewardTokens = rewardsDistributor.rewardTokens();
-
         if (caller != address(0)) {
+            // If the caller has no assets or is the zero address, skip compound
+            if (assetsBefore != 0) {
+                // Calculate new assets after distributing rewards
+                uint256 assetsAfter = convertToAssets(balanceOf(caller));
+
+                // Calculate the difference in assets
+                uint256 newAssets = assetsAfter - assetsBefore;
+
+                // Convert the difference in assets to shares
+                uint256 shares = convertToShares(newAssets);
+
+                // Mint new shares based on the difference in assets
+                _mint(caller, shares);
+            }
+
+            address[] rewardTokens = rewardsDistributor.rewardTokens();
+
             for (uint256 i = 0; i < rewardTokens.length; i++) {
                 address token = rewardTokens[i];
-                // ignore staking token as it will be auto compounded
+                // ignore staking token as it was compounded above
                 if (token == address(stakingToken)) {
                     continue;
                 }
@@ -142,24 +157,6 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
             }
         }
 
-        // If the caller has no assets or is the zero address, skip compound
-        if (assetsBefore == 0 || caller == address(0)) {
-            _;
-            return;
-        }
-
-        // Calculate new assets after distributing rewards
-        uint256 assetsAfter = convertToAssets(balanceOf(caller));
-
-        // Calculate the difference in assets
-        uint256 newAssets = assetsAfter - assetsBefore;
-
-        // Convert the difference in assets to shares
-        uint256 shares = convertToShares(newAssets);
-
-        // Mint new shares based on the difference in assets
-        _mint(caller, shares);
-
         _;
     }