Add TimelockControllerEnumerable

contracts/governance/TimelockControllerEnumerable.sol

@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {TimelockController} from "@openzeppelin/contracts/governance/TimelockController.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+
+/// @dev Extends the TimelockController to allow for enumerable operations
+contract TimelockControllerEnumerable is TimelockController {
+    using EnumerableSet for EnumerableSet.Bytes32Set;
+
+    /// @notice The operation struct
+    struct Operation {
+        address target;
+        uint256 value;
+        bytes data;
+        bytes32 predecessor;
+        bytes32 salt;
+        uint256 delay;
+    }
+
+    /// @notice The operation batch struct
+    struct OperationBatch {
+        address[] targets;
+        uint256[] values;
+        bytes[] payloads;
+        bytes32 predecessor;
+        bytes32 salt;
+        uint256 delay;
+    }
+
+    /// @notice The error when the operation index is not found
+    error OperationIndexNotFound(uint256 index);
+    /// @notice The error when the operation id is not found
+    error OperationIdNotFound(bytes32 id);
+    /// @notice The error when the operation batch index is not found
+    error OperationBatchIndexNotFound(uint256 index);
+    /// @notice The error when the operation batch id is not found
+    error OperationBatchIdNotFound(bytes32 id);
+
+    /// @notice The operations id set
+    EnumerableSet.Bytes32Set private _operationsIdSet;
+    /// @notice The operations map
+    mapping(bytes32 id => Operation operation) private _operationsMap;
+
+    /// @notice The operations batch id set
+    EnumerableSet.Bytes32Set private _operationsBatchIdSet;
+    /// @notice The operations batch map
+    mapping(bytes32 id => OperationBatch operationBatch) private _operationsBatchMap;
+
+    constructor(
+        uint256 minDelay,
+        address[] memory proposers,
+        address[] memory executors,
+        address canceller
+    ) TimelockController(minDelay, proposers, executors, canceller) {}
+
+    /// @inheritdoc TimelockController
+    /// @dev Store the operation the mapping and set
+    function schedule(
+        address target,
+        uint256 value,
+        bytes calldata data,
+        bytes32 predecessor,
+        bytes32 salt,
+        uint256 delay
+    ) public virtual override {
+        super.schedule(target, value, data, predecessor, salt, delay);
+        bytes32 id = hashOperation(target, value, data, predecessor, salt);
+        _operationsIdSet.add(id);
+        _operationsMap[id] = Operation({
+            target: target,
+            value: value,
+            data: data,
+            predecessor: predecessor,
+            salt: salt,
+            delay: delay
+        });
+    }
+
+    /// @inheritdoc TimelockController
+    /// @dev Store the operationBatch the mapping and set
+    function scheduleBatch(
+        address[] calldata targets,
+        uint256[] calldata values,
+        bytes[] calldata payloads,
+        bytes32 predecessor,
+        bytes32 salt,
+        uint256 delay
+    ) public virtual override {
+        super.scheduleBatch(targets, values, payloads, predecessor, salt, delay);
+        bytes32 id = hashOperationBatch(targets, values, payloads, predecessor, salt);
+        _operationsBatchIdSet.add(id);
+        _operationsBatchMap[id] = OperationBatch({
+            targets: targets,
+            values: values,
+            payloads: payloads,
+            predecessor: predecessor,
+            salt: salt,
+            delay: delay
+        });
+    }
+
+    /// @inheritdoc TimelockController
+    /// @dev Remove the operation from the mapping and set
+    function cancel(bytes32 id) public virtual override {
+        super.cancel(id);
+        if (_operationsIdSet.contains(id)) {
+            _operationsIdSet.remove(id);
+            delete _operationsMap[id];
+        }
+        if (_operationsBatchIdSet.contains(id)) {
+            _operationsBatchIdSet.remove(id);
+            delete _operationsBatchMap[id];
+        }
+    }
+
+    /// @dev Return the operations from the mapping and set
+    /// @return operations_ The operations array
+    function operations() public view returns (Operation[] memory operations_) {
+        uint256 operationsCount_ = _operationsIdSet.length();
+        operations_ = new Operation[](operationsCount_);
+        for (uint256 i = 0; i < operationsCount_; i++) {
+            operations_[i] = _operationsMap[_operationsIdSet.at(i)];
+        }
+        return operations_;
+    }
+
+    /// @dev Return the number of operations from the set
+    /// @return operationsCount_ The number of operations
+    function operationsCount() public view returns (uint256 operationsCount_) {
+        operationsCount_ = _operationsIdSet.length();
+        return operationsCount_;
+    }
+
+    /// @dev Return the operation at the given index from the mapping and set
+    /// @param index The index of the operation
+    /// @return operation_ The operation
+    function operation(uint256 index) public view returns (Operation memory operation_) {
+        if (index >= _operationsIdSet.length()) {
+            revert OperationIndexNotFound(index);
+        }
+        operation_ = _operationsMap[_operationsIdSet.at(index)];
+        return operation_;
+    }
+
+    /// @dev Return the operation with the given id from the mapping and set
+    /// @param id The id of the operation
+    /// @return operation_ The operation
+    function operation(bytes32 id) public view returns (Operation memory operation_) {
+        if (!_operationsIdSet.contains(id)) {
+            revert OperationIdNotFound(id);
+        }
+        operation_ = _operationsMap[id];
+        return operation_;
+    }
+
+    /// @dev Return the operationsBatch from the mapping and set
+    /// @return operationsBatch_ The operationsBatch array
+    function operationsBatch() public view returns (OperationBatch[] memory operationsBatch_) {
+        uint256 operationsBatchCount_ = _operationsBatchIdSet.length();
+        operationsBatch_ = new OperationBatch[](operationsBatchCount_);
+        for (uint256 i = 0; i < operationsBatchCount_; i++) {
+            operationsBatch_[i] = _operationsBatchMap[_operationsBatchIdSet.at(i)];
+        }
+        return operationsBatch_;
+    }
+
+    /// @dev Return the number of operationsBatch from the set
+    /// @return operationsBatchCount_ The number of operationsBatch
+    function operationsBatchCount() public view returns (uint256 operationsBatchCount_) {
+        operationsBatchCount_ = _operationsBatchIdSet.length();
+        return operationsBatchCount_;
+    }
+
+    /// @dev Return the operationsBatch at the given index from the mapping and set
+    /// @param index The index of the operationsBatch
+    /// @return operationBatch_ The operationsBatch
+    function operationBatch(uint256 index) public view returns (OperationBatch memory operationBatch_) {
+        if (index >= _operationsBatchIdSet.length()) {
+            revert OperationBatchIndexNotFound(index);
+        }
+        operationBatch_ = _operationsBatchMap[_operationsBatchIdSet.at(index)];
+        return operationBatch_;
+    }
+
+    /// @dev Return the operationsBatch with the given id from the mapping and set
+    /// @param id The id of the operationsBatch
+    /// @return operationBatch_ The operationsBatch
+    function operationBatch(bytes32 id) public view returns (OperationBatch memory operationBatch_) {
+        if (!_operationsBatchIdSet.contains(id)) {
+            revert OperationBatchIdNotFound(id);
+        }
+        operationBatch_ = _operationsBatchMap[id];
+        return operationBatch_;
+    }
+}

test/governance/TimelockControllerEnumerable.t.sol

@@ -0,0 +1,177 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.20;
+
+import {Test} from "forge-std/Test.sol";
+import {
+    TimelockControllerEnumerable
+} from "@openzeppelin/community-contracts/governance/TimelockControllerEnumerable.sol";
+
+contract TimelockControllerEnumerableTest is Test {
+    TimelockControllerEnumerable public timelockControllerEnumerable;
+
+    event Call();
+
+    function setUp() public {
+        address[] memory proposers = new address[](1);
+        address[] memory executors = new address[](1);
+        proposers[0] = address(this);
+        executors[0] = address(this);
+        uint256 minDelay = 1 days;
+        timelockControllerEnumerable = new TimelockControllerEnumerable(minDelay, proposers, executors, address(0));
+    }
+
+    function call() external {
+        emit Call();
+    }
+
+    function test_schedule() public {
+        timelockControllerEnumerable.schedule(
+            address(this),
+            0,
+            abi.encodeCall(this.call, ()),
+            bytes32(0),
+            bytes32(0),
+            1 days
+        );
+        assertEq(timelockControllerEnumerable.operationsCount(), 1);
+        TimelockControllerEnumerable.Operation memory operation = timelockControllerEnumerable.operation(uint256(0));
+        assertEq(operation.target, address(this));
+        assertEq(operation.value, 0);
+        assertEq(operation.data, abi.encodeCall(this.call, ()));
+        assertEq(operation.predecessor, bytes32(0));
+        assertEq(operation.salt, bytes32(0));
+        assertEq(operation.delay, 1 days);
+        bytes32 id = timelockControllerEnumerable.hashOperation(
+            address(this),
+            0,
+            abi.encodeCall(this.call, ()),
+            bytes32(0),
+            bytes32(0)
+        );
+        operation = timelockControllerEnumerable.operation(id);
+        assertEq(operation.target, address(this));
+        assertEq(operation.value, 0);
+        assertEq(operation.data, abi.encodeCall(this.call, ()));
+        assertEq(operation.predecessor, bytes32(0));
+        assertEq(operation.salt, bytes32(0));
+        assertEq(operation.delay, 1 days);
+    }
+
+    function test_schedule_execute() public {
+        test_schedule();
+        TimelockControllerEnumerable.Operation memory operation = timelockControllerEnumerable.operation(uint256(0));
+        bytes32 id = timelockControllerEnumerable.hashOperation(
+            operation.target,
+            operation.value,
+            operation.data,
+            operation.predecessor,
+            operation.salt
+        );
+        assertEq(timelockControllerEnumerable.isOperationPending(id), true);
+        vm.warp(block.timestamp + operation.delay);
+        timelockControllerEnumerable.execute(
+            operation.target,
+            operation.value,
+            operation.data,
+            operation.predecessor,
+            operation.salt
+        );
+        assertEq(timelockControllerEnumerable.isOperationPending(id), false);
+    }
+
+    function test_scheduleBatch() public {
+        address[] memory targets = new address[](1);
+        uint256[] memory values = new uint256[](1);
+        bytes[] memory payloads = new bytes[](1);
+        targets[0] = address(this);
+        values[0] = 0;
+        payloads[0] = abi.encodeCall(this.call, ());
+        timelockControllerEnumerable.scheduleBatch(targets, values, payloads, bytes32(0), bytes32(0), 1 days);
+        assertEq(timelockControllerEnumerable.operationsBatchCount(), 1);
+        TimelockControllerEnumerable.OperationBatch memory operationBatch = timelockControllerEnumerable.operationBatch(
+            uint256(0)
+        );
+        assertEq(operationBatch.targets[0], address(this));
+        assertEq(operationBatch.values[0], 0);
+        assertEq(operationBatch.payloads[0], abi.encodeCall(this.call, ()));
+        assertEq(operationBatch.predecessor, bytes32(0));
+        assertEq(operationBatch.salt, bytes32(0));
+        assertEq(operationBatch.delay, 1 days);
+        bytes32 id = timelockControllerEnumerable.hashOperationBatch(targets, values, payloads, bytes32(0), bytes32(0));
+        operationBatch = timelockControllerEnumerable.operationBatch(id);
+        assertEq(operationBatch.targets[0], address(this));
+        assertEq(operationBatch.values[0], 0);
+        assertEq(operationBatch.payloads[0], abi.encodeCall(this.call, ()));
+        assertEq(operationBatch.predecessor, bytes32(0));
+        assertEq(operationBatch.salt, bytes32(0));
+        assertEq(operationBatch.delay, 1 days);
+    }
+
+    function test_scheduleBatch_execute() public {
+        test_scheduleBatch();
+        TimelockControllerEnumerable.OperationBatch memory operationBatch = timelockControllerEnumerable.operationBatch(
+            uint256(0)
+        );
+        bytes32 id = timelockControllerEnumerable.hashOperationBatch(
+            operationBatch.targets,
+            operationBatch.values,
+            operationBatch.payloads,
+            operationBatch.predecessor,
+            operationBatch.salt
+        );
+        assertEq(timelockControllerEnumerable.isOperationPending(id), true);
+        vm.warp(block.timestamp + operationBatch.delay);
+        timelockControllerEnumerable.executeBatch(
+            operationBatch.targets,
+            operationBatch.values,
+            operationBatch.payloads,
+            operationBatch.predecessor,
+            operationBatch.salt
+        );
+        assertEq(timelockControllerEnumerable.isOperationPending(id), false);
+    }
+
+    function test_cancel_schedule() public {
+        timelockControllerEnumerable.schedule(
+            address(this),
+            0,
+            abi.encodeCall(this.call, ()),
+            bytes32(0),
+            bytes32(0),
+            1 days
+        );
+        assertEq(timelockControllerEnumerable.operationsCount(), 1);
+        bytes32 id = timelockControllerEnumerable.hashOperation(
+            address(this),
+            0,
+            abi.encodeCall(this.call, ()),
+            bytes32(0),
+            bytes32(0)
+        );
+        timelockControllerEnumerable.cancel(id);
+        assertEq(timelockControllerEnumerable.operationsCount(), 0);
+        vm.expectRevert(abi.encodeWithSelector(TimelockControllerEnumerable.OperationIdNotFound.selector, id));
+        timelockControllerEnumerable.operation(id);
+        vm.expectRevert(abi.encodeWithSelector(TimelockControllerEnumerable.OperationIndexNotFound.selector, 0));
+        timelockControllerEnumerable.operation(uint256(0));
+    }
+
+    function test_cancel_scheduleBatch() public {
+        address[] memory targets = new address[](1);
+        uint256[] memory values = new uint256[](1);
+        bytes[] memory payloads = new bytes[](1);
+        targets[0] = address(this);
+        values[0] = 0;
+        payloads[0] = abi.encodeCall(this.call, ());
+        timelockControllerEnumerable.scheduleBatch(targets, values, payloads, bytes32(0), bytes32(0), 1 days);
+        assertEq(timelockControllerEnumerable.operationsBatchCount(), 1);
+        bytes32 id = timelockControllerEnumerable.hashOperationBatch(targets, values, payloads, bytes32(0), bytes32(0));
+        timelockControllerEnumerable.cancel(id);
+        assertEq(timelockControllerEnumerable.operationsBatchCount(), 0);
+        vm.expectRevert(abi.encodeWithSelector(TimelockControllerEnumerable.OperationBatchIdNotFound.selector, id));
+        timelockControllerEnumerable.operationBatch(id);
+        vm.expectRevert(abi.encodeWithSelector(TimelockControllerEnumerable.OperationBatchIndexNotFound.selector, 0));
+        timelockControllerEnumerable.operationBatch(uint256(0));
+    }
+}