From 2939740f363308b9bd60771fb752f86e81ab4f7a Mon Sep 17 00:00:00 2001
From: Jean-Grimal <83286814+Jean-Grimal@users.noreply.github.com>
Date: Mon, 21 Oct 2024 13:48:56 +0200
Subject: [PATCH] feat: include balance in voting power slot

---
 src/ERC20DelegatesUpgradeable.sol |  6 +++---
 test/MorphoToken.t.sol            | 23 ++++++++++++++++++++++-
 test/helpers/BaseTest.sol         |  1 +
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/ERC20DelegatesUpgradeable.sol b/src/ERC20DelegatesUpgradeable.sol
index 0c1e99a..ff013b7 100644
--- a/src/ERC20DelegatesUpgradeable.sol
+++ b/src/ERC20DelegatesUpgradeable.sol
@@ -50,7 +50,7 @@ abstract contract ERC20DelegatesUpgradeable is
     /// @dev Returns the delegate that `account` has chosen.
     function delegates(address account) public view returns (address) {
         ERC20DelegatesStorage storage $ = _getERC20DelegatesStorage();
-        return $._delegatee[account];
+        return $._delegatee[account] == address(0) ? account : $._delegatee[account];
     }
 
     /* EXTERNAL */
@@ -58,7 +58,7 @@ abstract contract ERC20DelegatesUpgradeable is
     /// @dev Returns the current amount of votes that `account` has.
     function getVotes(address account) external view returns (uint256) {
         ERC20DelegatesStorage storage $ = _getERC20DelegatesStorage();
-        return delegates(account) == address(0) ? $._votingPower[account] + balanceOf(account) : $._votingPower[account];
+        return $._votingPower[account];
     }
 
     /// @dev Delegates votes from the sender to `delegatee`.
@@ -89,7 +89,7 @@ abstract contract ERC20DelegatesUpgradeable is
         $._delegatee[account] = delegatee;
 
         emit DelegateChanged(account, oldDelegate, delegatee);
-        _moveDelegateVotes(oldDelegate, delegatee, _getVotingUnits(account));
+        _moveDelegateVotes(oldDelegate, delegatee == address(0) ? account : delegatee, _getVotingUnits(account));
     }
 
     /// @dev Must return the voting units held by an account.
diff --git a/test/MorphoToken.t.sol b/test/MorphoToken.t.sol
index aad3851..147e9f8 100644
--- a/test/MorphoToken.t.sol
+++ b/test/MorphoToken.t.sol
@@ -1,13 +1,19 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
+import {console} from "lib/forge-std/src/Test.sol";
 import {BaseTest} from "./helpers/BaseTest.sol";
 import {SigUtils} from "./helpers/SigUtils.sol";
 import {MorphoToken} from "../src/MorphoToken.sol";
 import {ERC1967Proxy} from
     "lib/openzeppelin-contracts-upgradeable/lib/openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import {IERC20} from
+    "lib/openzeppelin-contracts-upgradeable/lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
 
 contract MorphoTokenTest is BaseTest {
+    bytes32 private constant ERC20DelegatesStorageLocation =
+        0x1dc92b2c6e971ab6e08dfd7dcec0e9496d223ced663ba2a06543451548549500;
+
     function testInitilizeZeroAddress(address randomAddress) public {
         vm.assume(randomAddress != address(0));
 
@@ -37,7 +43,7 @@ contract MorphoTokenTest is BaseTest {
         newMorpho.upgradeToAndCall(newImplem, hex"");
     }
 
-    function testOwnDelegation(address delegator, uint256 amount) public {
+    function testSelfDelegate(address delegator, uint256 amount) public {
         vm.assume(delegator != address(0));
         vm.assume(delegator != MORPHO_DAO);
         amount = bound(amount, MIN_TEST_AMOUNT, MAX_TEST_AMOUNT);
@@ -252,4 +258,19 @@ contract MorphoTokenTest is BaseTest {
             keccak256(abi.encode(uint256(keccak256("morpho.storage.ERC20Delegates")) - 1)) & ~bytes32(uint256(0xff));
         assertEq(expected, 0x1dc92b2c6e971ab6e08dfd7dcec0e9496d223ced663ba2a06543451548549500);
     }
+
+    function _getVotingPowerSlot(address account) internal pure returns (bytes32) {
+        return keccak256(abi.encode(account, uint256(ERC20DelegatesStorageLocation) + 1));
+    }
+
+    function deal(address token, address to, uint256 give) internal virtual override {
+        uint256 previousBalance = IERC20(address(newMorpho)).balanceOf(to);
+        if (address(newMorpho) == token) {
+            bytes32 votingPowerSlot = _getVotingPowerSlot(to);
+            uint256 previousVotingPower = uint256(vm.load(to, votingPowerSlot));
+            uint256 delegatedVotingPower = previousVotingPower - previousBalance;
+            vm.store(address(newMorpho), votingPowerSlot, bytes32(delegatedVotingPower + give));
+        }
+        super.deal(token, to, give);
+    }
 }
diff --git a/test/helpers/BaseTest.sol b/test/helpers/BaseTest.sol
index 2714e21..502c0a5 100644
--- a/test/helpers/BaseTest.sol
+++ b/test/helpers/BaseTest.sol
@@ -40,6 +40,7 @@ contract BaseTest is Test {
         for (uint256 i = 0; i < addresses.length; i++) {
             vm.assume(addresses[i] != address(0));
             vm.assume(addresses[i] != MORPHO_DAO);
+            assumeNotPrecompile(addresses[i]);
             for (uint256 j = i + 1; j < addresses.length; j++) {
                 vm.assume(addresses[i] != addresses[j]);
             }