feat: add flush to delegate set

defi-wonderland · Jan 5, 2024 · df91918 · df91918
1 parent d4e993f
commit df91918
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/solidity/contracts/governance/utils/DelegateSet.sol b/solidity/contracts/governance/utils/DelegateSet.sol
@@ -155,4 +155,15 @@ library DelegateSet {
   function values(Set storage set) internal view returns (IWonderVotes.Delegate[] memory) {
     return set._delegates;
   }
+
+  /**
+   * @dev Removes all delegates from a set. O(n).
+   */
+  function flush(Set storage set) internal {
+    for (uint256 i = 0; i < set._delegates.length; i++) {
+      delete set._positions[set._delegates[i].account];
+    }
+
+    delete set._delegates;
+  }
 }
diff --git a/solidity/contracts/governance/utils/WonderVotes.sol b/solidity/contracts/governance/utils/WonderVotes.sol
@@ -277,8 +277,9 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
     if (_weightSum != _weightNormalizer()) revert VotesInvalidWeightSum(_weightSum);
 
     Delegate[] memory _oldDelegates = delegates(account, proposalType);
-    for (uint256 i = 0; i < _oldDelegates.length; i++) {
-      _delegatees[account][proposalType].remove(_oldDelegates[i]);
+
+    if (_oldDelegates.length > 0) {
+      _delegatees[account][proposalType].flush();
     }
 
     for (uint256 i = 0; i < delegatees.length; i++) {