Lock Manager unlock primitives WIP

src/LockManager.sol

@@ -5,54 +5,167 @@ import {IDAO} from "@aragon/osx/core/dao/IDAO.sol";
 import {DaoAuthorizable} from "@aragon/osx/core/plugin/dao-authorizable/DaoAuthorizable.sol";
 import {ILockManager} from "./interfaces/ILockManager.sol";
 import {ILockToVote} from "./interfaces/ILockToVote.sol";
-import {IVotesUpgradeable} from "@openzeppelin/contracts-upgradeable/governance/utils/IVotesUpgradeable.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 contract LockManager is ILockManager, DaoAuthorizable {
-    enum LockMode {
+    /// @notice Defines whether locked funds can be unlocked at any time or not
+    enum UnlockMode {
         STRICT,
-        FLEXIBLE
+        EARLY
     }
+
+    /// @notice The struct containing the LockManager helper settings
     struct Settings {
-        LockMode lockMode;
+        /// @param lockMode The mode defining whether funds can be unlocked at any time or not
+        UnlockMode unlockMode;
+        /// @param plugin The address of the lock to vote plugin to use
+        ILockToVote plugin;
+        /// @param token The address of the token contract
+        IERC20 token;
     }
 
+    /// @notice The current LockManager settings
     Settings public settings;
 
-    error InvalidLockMode();
+    /// @notice Keeps track of the amount of tokens locked by address
+    mapping(address => uint256) lockedBalance;
+
+    /// @notice Keeps a list of the known active proposal ID's
+    /// @dev Executed proposals will be actively reported, but defeated proposals will need to be garbage collected over time.
+    uint256[] knownProposalIds;
+
+    /// @notice Thrown when trying to assign an invalid lock mode
+    error InvalidUnlockMode();
+
+    /// @notice Raised when the caller holds no tokens or didn't lock any tokens
+    error NoBalance();
+
+    /// @notice Raised when attempting to unlock while active votes are cast in strict mode
+    error LocksStillActive();
 
     constructor(IDAO _dao, Settings memory _settings) DaoAuthorizable(_dao) {
         if (
-            _settings.lockMode != LockMode.STRICT &&
-            _settings.lockMode != LockMode.FLEXIBLE
+            _settings.unlockMode != UnlockMode.STRICT &&
+            _settings.unlockMode != UnlockMode.EARLY
         ) {
-            revert InvalidLockMode();
+            revert InvalidUnlockMode();
         }
 
-        settings.lockMode = _settings.lockMode;
+        settings.unlockMode = _settings.unlockMode;
+        settings.plugin = _settings.plugin;
+        settings.token = _settings.token;
     }
 
     /// @inheritdoc ILockManager
     function lock() public {
+        // Register the token if not present
         //
     }
 
     /// @inheritdoc ILockManager
-    function lockAndVote(ILockToVote plugin, uint256 proposalId) public {
+    function lockAndVote(uint256 proposalId) public {
         //
     }
 
     /// @inheritdoc ILockManager
-    function vote(ILockToVote plugin, uint256 proposalId) public {
-        //
+    function vote(uint256 proposalId) public {
+        uint256 newVotingPower = lockedBalance[msg.sender];
+        settings.plugin.vote(proposalId, msg.sender, newVotingPower);
     }
 
     /// @inheritdoc ILockManager
     function unlock() public {
-        //
+        if (lockedBalance[msg.sender] == 0) {
+            revert NoBalance();
+        }
+
+        if (settings.unlockMode == UnlockMode.STRICT) {
+            if (hasActiveLocks()) revert LocksStillActive();
+        } else {
+            withdrawActiveVotingPower();
+        }
+
+        // All votes clear
+
+        // Refund
+        uint256 refundBalance = lockedBalance[msg.sender];
+        lockedBalance[msg.sender] = 0;
+
+        settings.token.transfer(msg.sender, refundBalance);
     }
 
     /// @inheritdoc ILockManager
     function releaseLock(uint256 proposalId) public {
         //
     }
+
+    // Internal
+
+    function hasActiveLocks() internal returns (bool) {
+        uint256 _proposalCount = knownProposalIds.length;
+        for (uint256 _i; _i < _proposalCount; ) {
+            (bool open, ) = settings.plugin.getProposal(knownProposalIds[_i]);
+            if (!open) {
+                cleanKnownProposalId(_i);
+
+                // Are we at the last item?
+                /// @dev Comparing to `_proposalCount` instead of `_proposalCount - 1`, because the array is now shorter
+                if (_i == _proposalCount) {
+                    return false;
+                }
+
+                // Recheck the same index (now, another proposal)
+                continue;
+            }
+
+            if (settings.plugin.hasVoted(msg.sender)) {
+                return true;
+            }
+
+            unchecked {
+                _i++;
+            }
+        }
+    }
+
+    function withdrawActiveVotingPower() internal {
+        uint256 _proposalCount = knownProposalIds.length;
+        for (uint256 _i; _i < _proposalCount; ) {
+            (bool open, ) = settings.plugin.getProposal(knownProposalIds[_i]);
+            if (!open) {
+                cleanKnownProposalId(_i);
+
+                // Are we at the last item?
+                /// @dev Comparing to `_proposalCount` instead of `_proposalCount - 1`, because the array is now shorter
+                if (_i == _proposalCount) {
+                    return;
+                }
+
+                // Recheck the same index (now, another proposal)
+                continue;
+            }
+
+            uint votedBalance = settings.plugin.votedBalance(msg.sender);
+            if (votedBalance > 0) {
+                settings.plugin.clearVote(knownProposalIds[_i], msg.sender);
+            }
+
+            unchecked {
+                _i++;
+            }
+        }
+    }
+
+    /// @dev Cleaning up ended proposals, otherwise they would pile up and make unlocks more and more gas costly over time
+    function cleanKnownProposalId(uint _arrayIndex) internal {
+        // Swap the current item with the last, if needed
+        if (_arrayIndex < knownProposalIds.length - 1) {
+            knownProposalIds[_arrayIndex] = knownProposalIds[
+                knownProposalIds.length - 1
+            ];
+        }
+
+        // Trim the array's last item
+        knownProposalIds.length -= 1;
+    }
 }

src/interfaces/ILockManager.sol

@@ -12,14 +12,12 @@ interface ILockManager {
     function lock() external;
 
     /// @notice Locks the balance currently allowed by msg.sender on this contract and registers a vote on the target plugin
-    /// @param plugin The address of the lock to vote plugin where the lock will be used
     /// @param proposalId The ID of the proposal where the vote will be registered
-    function lockAndVote(ILockToVote plugin, uint256 proposalId) external;
+    function lockAndVote(uint256 proposalId) external;
 
     /// @notice Uses the locked balance to place a vote on the given proposal for the given plugin
-    /// @param plugin The address of the lock to vote plugin where the locked balance will be used
     /// @param proposalId The ID of the proposal where the vote will be registered
-    function vote(ILockToVote plugin, uint256 proposalId) external;
+    function vote(uint256 proposalId) external;
 
     /// @notice If the mode allows it, releases all active locks placed on active proposals and transfers msg.sender's locked balance back. Depending on the current mode, it withdraws only if no locks are being used in active proposals.
     function unlock() external;

src/interfaces/ILockToVote.sol

@@ -35,28 +35,38 @@ interface ILockToVote {
     /// - was executed, or
     /// - the voter doesn't have any tokens locked.
     /// @param proposalId The proposal Id.
-    /// @param account The account address to be checked.
+    /// @param voter The account address to be checked.
     /// @return Returns true if the account is allowed to vote.
     /// @dev The function assumes that the queried proposal exists.
-    function canVeto(
+    function canVote(
         uint256 proposalId,
-        address account
+        address voter
     ) external view returns (bool);
 
     /// @notice Registers an approval vote for the given proposal.
     /// @param proposalId The ID of the proposal to vote on.
-    function vote(uint256 proposalId) external;
+    /// @param voter The address of the account whose vote will be registered
+    /// @param newVotingPower The new balance that should be allocated to the voter. It can only be bigger.
+    /// @dev newVotingPower updates any prior voting power, it does not add to the existing amount.
+    function vote(
+        uint256 proposalId,
+        address voter,
+        uint newVotingPower
+    ) external;
 
-    function clearVote(uint256 proposalId) external;
+    /// @notice Reverts the existing voter's vote, if any.
+    /// @param proposalId The ID of the proposal.
+    /// @param voter The voter's address.
+    function clearVote(uint256 proposalId, address voter) external;
 
     /// @notice Returns whether the account has voted for the proposal.
     /// @param proposalId The ID of the proposal.
-    /// @param account The account address to be checked.
-    /// @return The whether the given account has voted for the given proposal to pass.
-    function hasVoted(
+    /// @param voter The account address to be checked.
+    /// @return The amount of balance that has been allocated for to the proposal by the given account.
+    function votedBalance(
         uint256 proposalId,
-        address account
-    ) external view returns (bool);
+        address voter
+    ) external view returns (uint256);
 
     /// @notice Checks if the amount of locked votes for the given proposal is greater than the approval threshold.
     /// @param proposalId The ID of the proposal.
@@ -74,7 +84,7 @@ interface ILockToVote {
     /// @param proposalId The ID of the proposal to execute.
     function execute(uint256 proposalId) external;
 
-    /// @notice If the given proposal is no longer active, it notifies the manager so that the active locks no longer track it.
+    /// @notice If the given proposal is no longer active, it allows to notify the manager.
     /// @param proposalId The ID of the proposal to clean up for.
     function releaseLock(uint256 proposalId) external;
 }