When health changes update

contracts/Account.sol

@@ -194,7 +194,12 @@ contract Account is UUPSOwnableUpgradeable, UsingRegistryUpgradeable, Managed, I
      * @notice Scheduling transfer was not successfull since
      * total amount of "from" and "to" are not the same.
      */
-    error TransferAmountMisalignment();
+    error TransferAmountMisalignment(uint256 totalFromVotes, uint256 totalToVotes);
+
+    /**
+     * @notice When scheduling transfer and the from group doesn't have enough CELO.
+     */
+    error NotEnoughCeloInGroup(address group, uint256 expectedAmount, uint256 realAmount);
 
     /**
      * @notice Empty constructor for proxy implementation, `initializer` modifer ensures the
@@ -288,7 +293,8 @@ contract Account is UUPSOwnableUpgradeable, UsingRegistryUpgradeable, Managed, I
         for (uint256 i = 0; i < fromGroups.length; i++) {
             uint256 celoAvailableForGroup = getCeloForGroup(fromGroups[i]);
 
-            if (celoAvailableForGroup < fromVotes[i]) revert TransferAmountMisalignment();
+            if (celoAvailableForGroup < fromVotes[i])
+                revert NotEnoughCeloInGroup(fromGroups[i], fromVotes[i], celoAvailableForGroup);
             getAndUpdateToVoteAndToRevoke(fromGroups[i], 0, fromVotes[i]);
             totalFromVotes += fromVotes[i];
         }
@@ -299,7 +305,7 @@ contract Account is UUPSOwnableUpgradeable, UsingRegistryUpgradeable, Managed, I
         }
 
         if (totalFromVotes != totalToVotes) {
-            revert TransferAmountMisalignment();
+            revert TransferAmountMisalignment(totalFromVotes, totalToVotes);
         }
     }
 

contracts/DefaultStrategy.sol

@@ -291,12 +291,16 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
      * @return finalGroups The groups that were chosen for distribution.
      * @return finalVotes The votes of chosen finalGroups.
      */
-    function generateDepositVoteDistribution(uint256 celoAmount, address depositGroupToIgnore)
+    function generateDepositVoteDistribution(
+        uint256 celoAmount,
+        uint256 stCeloAmount,
+        address depositGroupToIgnore
+    )
         external
         managerOrStrategy
         returns (address[] memory finalGroups, uint256[] memory finalVotes)
     {
-        return _generateDepositVoteDistribution(celoAmount, depositGroupToIgnore);
+        return _generateDepositVoteDistribution(celoAmount, stCeloAmount, depositGroupToIgnore);
     }
 
     /**
@@ -694,6 +698,7 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
             _updateGroupStCelo(group, groupTotalStCeloVotes, false);
             _generateDepositVoteDistribution(
                 IManager(manager).toCelo(groupTotalStCeloVotes),
+                groupTotalStCeloVotes,
                 address(0)
             );
         }
@@ -709,10 +714,11 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
      * @return finalGroups The groups that were chosen for distribution.
      * @return finalVotes The votes of chosen finalGroups.
      */
-    function _generateDepositVoteDistribution(uint256 celoAmount, address depositGroupToIgnore)
-        private
-        returns (address[] memory finalGroups, uint256[] memory finalVotes)
-    {
+    function _generateDepositVoteDistribution(
+        uint256 celoAmount,
+        uint256 stCeloAmount,
+        address depositGroupToIgnore
+    ) private returns (address[] memory finalGroups, uint256[] memory finalVotes) {
         if (activeGroups.getNumElements() == 0) {
             revert NoActiveGroups();
         }
@@ -735,11 +741,9 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
             votes[groupsIndex] = Math.min(receivableVotes, celoAmount);
             groups[groupsIndex] = votedGroup;
             celoAmount -= votes[groupsIndex];
-            _updateGroupStCelo(
-                votedGroup,
-                IManager(manager).toStakedCelo(votes[groupsIndex]),
-                true
-            );
+            uint256 stCelo = IManager(manager).toStakedCelo(votes[groupsIndex]);
+            stCeloAmount = stCeloAmount >= stCelo ? stCeloAmount - stCelo : 0;
+            _updateGroupStCelo(votedGroup, stCelo, true);
             trySort(votedGroup, stCeloInGroup[votedGroup], true);
 
             if (sorted) {
@@ -754,6 +758,11 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
             revert NotAbleToDistributeVotes();
         }
 
+        if (stCeloAmount != 0) {
+            _updateGroupStCelo(votedGroup, stCeloAmount, true);
+            trySort(votedGroup, stCeloInGroup[votedGroup], true);
+        }
+
         finalGroups = new address[](groupsIndex);
         finalVotes = new uint256[](groupsIndex);
 

contracts/Manager.sol

@@ -532,6 +532,8 @@ contract Manager is Errors, UUPSOwnableUpgradeable, UsingRegistryUpgradeable, Pa
             return (celoToRemove - celoScheduled, 0);
         }
 
+        // TODO: check if we could rebalance to unhealthy group if needed
+
         actualCelo = celoScheduled - celoToRemove;
 
         bool isSpecificGroupStrategy = !specificGroupStrategy.isBlockedGroup(group);
@@ -635,6 +637,7 @@ contract Manager is Errors, UUPSOwnableUpgradeable, UsingRegistryUpgradeable, Pa
         } else {
             (finalGroups, finalVotes) = defaultStrategy.generateDepositVoteDistribution(
                 votes,
+                stCeloAmount,
                 address(0)
             );
         }

contracts/SpecificGroupStrategy.sol

@@ -333,14 +333,14 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
 
             celoAmount -= votesToBeScheduledForSpecificGroup;
             if (celoAmount > 0) {
+                uint256 stCeloAmountOverflow = Math.min(
+                    IManager(manager).toStakedCelo(celoAmount),
+                    stCeloAmount
+                );
                 // overflow
                 (address[] memory groups, uint256[] memory votesForGroups) = defaultStrategy
-                    .generateDepositVoteDistribution(celoAmount, group);
-                updateOverflowGroup(
-                    group,
-                    Math.min(IManager(manager).toStakedCelo(celoAmount), stCeloAmount),
-                    true
-                );
+                    .generateDepositVoteDistribution(celoAmount, stCeloAmountOverflow, group);
+                updateOverflowGroup(group, stCeloAmountOverflow, true);
                 finalGroups = new address[](groups.length + 1);
                 finalVotes = new uint256[](groups.length + 1);
                 for (uint256 i = 0; i < groups.length; i++) {
@@ -358,6 +358,7 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
         } else {
             (finalGroups, finalVotes) = defaultStrategy.generateDepositVoteDistribution(
                 celoAmount,
+                stCeloAmount,
                 group
             );
             updateUnhealthyGroupStCelo(group, stCeloAmount, true);
@@ -391,10 +392,10 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
             }
 
             uint256 overflow = stCeloInGroupOverflowed[group];
-            uint256 toMove = totalStCeloInGroup - unhealthyStCelo - overflow;
+            uint256 toMoveStCelo = totalStCeloInGroup - unhealthyStCelo - overflow;
 
-            transferToDefaultStrategy(group, toMove);
-            updateUnhealthyGroupStCelo(group, toMove, true);
+            transferToDefaultStrategy(group, toMoveStCelo);
+            updateUnhealthyGroupStCelo(group, toMoveStCelo, true);
         }
     }
 
@@ -602,13 +603,16 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
      * @param stCeloToMove StCelo amount to be moved.
      */
     function transferToDefaultStrategy(address group, uint256 stCeloToMove) private {
-        uint256 toMoveCelo = IManager(manager).toCelo(stCeloToMove);
+        uint256 toMoveCelo = Math.min(
+            account.getCeloForGroup(group),
+            IManager(manager).toCelo(stCeloToMove)
+        );
         address[] memory fromGroups = new address[](1);
         uint256[] memory fromVotes = new uint256[](1);
         fromGroups[0] = group;
         fromVotes[0] = toMoveCelo;
         (address[] memory toGroups, uint256[] memory toVotes) = defaultStrategy
-            .generateDepositVoteDistribution(toMoveCelo, address(0));
+            .generateDepositVoteDistribution(toMoveCelo, stCeloToMove, address(0));
         IManager(manager).scheduleTransferWithinStrategy(fromGroups, toGroups, fromVotes, toVotes);
     }
 

contracts/interfaces/IDefaultStrategy.sol

@@ -2,9 +2,11 @@
 pragma solidity 0.8.11;
 
 interface IDefaultStrategy {
-    function generateDepositVoteDistribution(uint256 celoAmount, address depositGroupToIgnore)
-        external
-        returns (address[] memory finalGroups, uint256[] memory finalVotes);
+    function generateDepositVoteDistribution(
+        uint256 celoAmount,
+        uint256 stCeloAmount,
+        address depositGroupToIgnore
+    ) external returns (address[] memory finalGroups, uint256[] memory finalVotes);
 
     function generateWithdrawalVoteDistribution(uint256 celoAmount)
         external

test-ts/account.test.ts

@@ -520,15 +520,15 @@ describe("Account", () => {
         account
           .connect(manager)
           .scheduleTransfer(groupAddresses, [100, 30, 80], groupAddresses, [30, 70, 100])
-      ).revertedWith(`TransferAmountMisalignment()`);
+      ).revertedWith(`'NotEnoughCeloInGroup("${groupAddresses[0]}", 100, 0)`);
     });
 
     it("should revert when incorrect vote sum 2", async () => {
       await expect(
         account
           .connect(manager)
           .scheduleTransfer(groupAddresses, [100, 30, 60], groupAddresses, [30, 70, 100])
-      ).revertedWith(`TransferAmountMisalignment()`);
+      ).revertedWith(`NotEnoughCeloInGroup("${groupAddresses[0]}", 100, 0)`);
     });
 
     describe("When group has activated votes", () => {
@@ -645,7 +645,7 @@ describe("Account", () => {
               [groupAddresses[1]],
               [originalGroupAmount * 2]
             )
-        ).revertedWith(`TransferAmountMisalignment()`);
+        ).revertedWith(`NotEnoughCeloInGroup("${groupAddresses[0]}", 200, 100)`);
       });
     });
 

test-ts/default-strategy.test.ts

@@ -1250,7 +1250,7 @@ describe("DefaultStrategy", () => {
       await expect(
         defaultStrategyContract
           .connect(nonManager)
-          .generateDepositVoteDistribution(10, ADDRESS_ZERO)
+          .generateDepositVoteDistribution(10, 10, ADDRESS_ZERO)
       ).revertedWith(`CallerNotManagerNorStrategy("${nonManager.address}")`);
     });
   });