From df91918bfa347e95b8adb0b6b5ad2217a873c31b Mon Sep 17 00:00:00 2001
From: agusduha <agusnduha@gmail.com>
Date: Fri, 5 Jan 2024 10:16:33 -0300
Subject: [PATCH] feat: add flush to delegate set

---
 solidity/contracts/governance/utils/DelegateSet.sol | 11 +++++++++++
 solidity/contracts/governance/utils/WonderVotes.sol |  5 +++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/solidity/contracts/governance/utils/DelegateSet.sol b/solidity/contracts/governance/utils/DelegateSet.sol
index 02759df..8512170 100644
--- 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
index cc672c5..62e96f4 100644
--- 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++) {