feat: change total voting power

morpho-org · Oct 18, 2024 · a8046ac · a8046ac
1 parent caa2c37
commit a8046ac
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 14 deletions.
diff --git a/src/ERC20DelegatesUpgradeable.sol b/src/ERC20DelegatesUpgradeable.sol
@@ -63,6 +63,12 @@ abstract contract ERC20DelegatesUpgradeable is
         return $._votingPower[account];
     }
 
+    /// @dev Returns the current total voting power.
+    function getTotalVotes() external view returns (uint256) {
+        ERC20DelegatesStorage storage $ = _getERC20DelegatesStorage();
+        return $._totalVotingPower;
+    }
+
     /// @dev Delegates votes from the sender to `delegatee`.
     function delegate(address delegatee) external {
         address account = _msgSender();
@@ -94,19 +100,6 @@ abstract contract ERC20DelegatesUpgradeable is
         _moveDelegateVotes(oldDelegate, delegatee, _getVotingUnits(account));
     }
 
-    /// @dev Transfers, mints, or burns voting units. To register a mint, `from` should be zero. To register a burn, `to`
-    /// should be zero. Total supply of voting units will be adjusted with mints and burns.
-    function _transferVotingUnits(address from, address to, uint256 amount) internal {
-        ERC20DelegatesStorage storage $ = _getERC20DelegatesStorage();
-        if (from == address(0)) {
-            $._totalVotingPower += amount;
-        }
-        if (to == address(0)) {
-            $._totalVotingPower -= amount;
-        }
-        _moveDelegateVotes(delegates(from), delegates(to), amount);
-    }
-
     /// @dev Must return the voting units held by an account.
     function _getVotingUnits(address account) internal view returns (uint256) {
         return balanceOf(account);
@@ -117,7 +110,7 @@ abstract contract ERC20DelegatesUpgradeable is
     function _update(address from, address to, uint256 value) internal virtual override {
         super._update(from, to, value);
         // No check of supply cap here like in OZ implementation as MORPHO has a 1B total supply cap.
-        _transferVotingUnits(from, to, value);
+        _moveDelegateVotes(delegates(from), delegates(to), value);
     }
 
     /* PRIVATE */
@@ -131,12 +124,16 @@ abstract contract ERC20DelegatesUpgradeable is
                 uint256 newValue = oldValue - amount;
                 $._votingPower[from] = newValue;
                 emit DelegateVotesChanged(from, oldValue, newValue);
+            } else {
+                $._totalVotingPower += amount;
             }
             if (to != address(0)) {
                 uint256 oldValue = $._votingPower[to];
                 uint256 newValue = oldValue + amount;
                 $._votingPower[to] = newValue;
                 emit DelegateVotesChanged(to, oldValue, newValue);
+            } else {
+                $._totalVotingPower -= amount;
             }
         }
     }

diff --git a/test/MorphoToken.t.sol b/test/MorphoToken.t.sol
@@ -58,6 +58,8 @@ contract MorphoTokenTest is BaseTest {
         _validateAddresses(addresses);
         amount = bound(amount, MIN_TEST_AMOUNT, MAX_TEST_AMOUNT);
 
+        assertEq(newMorpho.getTotalVotes(), 0);
+
         deal(address(newMorpho), delegator, amount);
 
         vm.prank(delegator);
@@ -66,6 +68,14 @@ contract MorphoTokenTest is BaseTest {
         assertEq(newMorpho.delegates(delegator), delegatee);
         assertEq(newMorpho.getVotes(delegator), 0);
         assertEq(newMorpho.getVotes(delegatee), amount);
+        assertEq(newMorpho.getTotalVotes(), amount);
+
+        vm.prank(delegator);
+        newMorpho.delegate(address(0));
+
+        assertEq(newMorpho.delegates(delegator), address(0));
+        assertEq(newMorpho.getVotes(delegator), 0);
+        assertEq(newMorpho.getTotalVotes(), 0);
     }
 
     function testDelegateBySigExpired(SigUtils.Delegation memory delegation, uint256 privateKey, uint256 expiry)
@@ -142,6 +152,7 @@ contract MorphoTokenTest is BaseTest {
         assertEq(newMorpho.getVotes(delegator), 0);
         assertEq(newMorpho.getVotes(delegation.delegatee), amount);
         assertEq(newMorpho.nonces(delegator), 1);
+        assertEq(newMorpho.getTotalVotes(), amount);
     }
 
     function testMultipleDelegations(