From c3cb7a02954f6e7e8f4c2c97ed85b6462289c180 Mon Sep 17 00:00:00 2001
From: Hadrien Croubois <hadrien.croubois@gmail.com>
Date: Mon, 25 Nov 2024 18:47:55 +0100
Subject: [PATCH] Deduplicate logic in Votes.sol (#5314)

Co-authored-by: Arr00 <13561405+arr00@users.noreply.github.com>
---
 contracts/governance/utils/Votes.sol         | 21 ++++++++++----------
 contracts/governance/utils/VotesExtended.sol | 17 ++++------------
 2 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/contracts/governance/utils/Votes.sol b/contracts/governance/utils/Votes.sol
index bbbc2264f..976579719 100644
--- a/contracts/governance/utils/Votes.sol
+++ b/contracts/governance/utils/Votes.sol
@@ -71,6 +71,15 @@ abstract contract Votes is Context, EIP712, Nonces, IERC5805 {
         return "mode=blocknumber&from=default";
     }
 
+    /**
+     * @dev Validate that a timepoint is in the past, and return it as a uint48.
+     */
+    function _validateTimepoint(uint256 timepoint) internal view returns (uint48) {
+        uint48 currentTimepoint = clock();
+        if (timepoint >= currentTimepoint) revert ERC5805FutureLookup(timepoint, currentTimepoint);
+        return SafeCast.toUint48(timepoint);
+    }
+
     /**
      * @dev Returns the current amount of votes that `account` has.
      */
@@ -87,11 +96,7 @@ abstract contract Votes is Context, EIP712, Nonces, IERC5805 {
      * - `timepoint` must be in the past. If operating using block numbers, the block must be already mined.
      */
     function getPastVotes(address account, uint256 timepoint) public view virtual returns (uint256) {
-        uint48 currentTimepoint = clock();
-        if (timepoint >= currentTimepoint) {
-            revert ERC5805FutureLookup(timepoint, currentTimepoint);
-        }
-        return _delegateCheckpoints[account].upperLookupRecent(SafeCast.toUint48(timepoint));
+        return _delegateCheckpoints[account].upperLookupRecent(_validateTimepoint(timepoint));
     }
 
     /**
@@ -107,11 +112,7 @@ abstract contract Votes is Context, EIP712, Nonces, IERC5805 {
      * - `timepoint` must be in the past. If operating using block numbers, the block must be already mined.
      */
     function getPastTotalSupply(uint256 timepoint) public view virtual returns (uint256) {
-        uint48 currentTimepoint = clock();
-        if (timepoint >= currentTimepoint) {
-            revert ERC5805FutureLookup(timepoint, currentTimepoint);
-        }
-        return _totalCheckpoints.upperLookupRecent(SafeCast.toUint48(timepoint));
+        return _totalCheckpoints.upperLookupRecent(_validateTimepoint(timepoint));
     }
 
     /**
diff --git a/contracts/governance/utils/VotesExtended.sol b/contracts/governance/utils/VotesExtended.sol
index b68ee2712..27baaab84 100644
--- a/contracts/governance/utils/VotesExtended.sol
+++ b/contracts/governance/utils/VotesExtended.sol
@@ -31,7 +31,6 @@ import {SafeCast} from "../../utils/math/SafeCast.sol";
  * {ERC20Votes} and {ERC721Votes} follow this pattern and are thus safe to use with {VotesExtended}.
  */
 abstract contract VotesExtended is Votes {
-    using SafeCast for uint256;
     using Checkpoints for Checkpoints.Trace160;
     using Checkpoints for Checkpoints.Trace208;
 
@@ -47,11 +46,7 @@ abstract contract VotesExtended is Votes {
      * - `timepoint` must be in the past. If operating using block numbers, the block must be already mined.
      */
     function getPastDelegate(address account, uint256 timepoint) public view virtual returns (address) {
-        uint48 currentTimepoint = clock();
-        if (timepoint >= currentTimepoint) {
-            revert ERC5805FutureLookup(timepoint, currentTimepoint);
-        }
-        return address(_delegateCheckpoints[account].upperLookupRecent(timepoint.toUint48()));
+        return address(_delegateCheckpoints[account].upperLookupRecent(_validateTimepoint(timepoint)));
     }
 
     /**
@@ -63,11 +58,7 @@ abstract contract VotesExtended is Votes {
      * - `timepoint` must be in the past. If operating using block numbers, the block must be already mined.
      */
     function getPastBalanceOf(address account, uint256 timepoint) public view virtual returns (uint256) {
-        uint48 currentTimepoint = clock();
-        if (timepoint >= currentTimepoint) {
-            revert ERC5805FutureLookup(timepoint, currentTimepoint);
-        }
-        return _balanceOfCheckpoints[account].upperLookupRecent(timepoint.toUint48());
+        return _balanceOfCheckpoints[account].upperLookupRecent(_validateTimepoint(timepoint));
     }
 
     /// @inheritdoc Votes
@@ -82,10 +73,10 @@ abstract contract VotesExtended is Votes {
         super._transferVotingUnits(from, to, amount);
         if (from != to) {
             if (from != address(0)) {
-                _balanceOfCheckpoints[from].push(clock(), _getVotingUnits(from).toUint208());
+                _balanceOfCheckpoints[from].push(clock(), SafeCast.toUint208(_getVotingUnits(from)));
             }
             if (to != address(0)) {
-                _balanceOfCheckpoints[to].push(clock(), _getVotingUnits(to).toUint208());
+                _balanceOfCheckpoints[to].push(clock(), SafeCast.toUint208(_getVotingUnits(to)));
             }
         }
     }