Feat/upgradeable task registry (#131)

* feat: upgradeable taskRegistry * chore: fmt
Aut-Labs · Oct 23, 2024 · 8ea35ba · 8ea35ba
1 parent 08f7492
commit 8ea35ba
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 14 deletions.
diff --git a/contracts/tasks/TaskRegistry.sol b/contracts/tasks/TaskRegistry.sol
@@ -3,22 +3,45 @@ pragma solidity ^0.8.20;
 
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {Task, ITaskRegistry} from "./interfaces/ITaskRegistry.sol";
-import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 
-contract TaskRegistry is ITaskRegistry, Ownable {
+contract TaskRegistry is ITaskRegistry, OwnableUpgradeable {
     using EnumerableSet for EnumerableSet.Bytes32Set;
 
-    EnumerableSet.Bytes32Set private _taskIds;
-    mapping(bytes32 => Task) private _tasks;
+    struct TaskRegistryStorage {
+        EnumerableSet.Bytes32Set taskIds;
+        mapping(bytes32 => Task) tasks;
+        address approved;
+    }
+
+    // keccak256(abi.encode(uint256(keccak256("aut.storage.TaskRegistry")) - 1))
+    bytes32 private constant TaskRegistryStorageLocation =
+        0xc1f8b0ede512a55dbbda242577f6b0c5214dd214df36b3a4a3854a073de90410;
+
+    function _getTaskRegistryStorage() private pure returns (TaskRegistryStorage storage $) {
+        assembly {
+            $.slot := TaskRegistryStorageLocation
+        }
+    }
 
-    address public approved;
+    constructor() {
+        _disableInitializers();
+    }
 
-    constructor() Ownable(msg.sender) {}
+    function initialize() external initializer {
+        __Ownable_init(msg.sender);
+    }
 
     // Access control to register tasks
 
+    function approved() external view returns (address) {
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        return $.approved;
+    }
+
     function setApproved(address _approved) external onlyOwner {
-        approved = _approved;
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        $.approved = _approved;
     }
 
     error NotApproved();
@@ -29,7 +52,8 @@ contract TaskRegistry is ITaskRegistry, Ownable {
 
     /// @dev returns true if the approved address set to 0 or the msg.sender is the approved address
     function _isApproved(address who) internal view returns (bool) {
-        address approved_ = approved; // gas
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        address approved_ = $.approved; // gas
         return (approved_ == address(0) || approved_ == msg.sender);
     }
 
@@ -50,11 +74,13 @@ contract TaskRegistry is ITaskRegistry, Ownable {
     }
 
     function _registerTask(Task memory task) internal returns (bytes32) {
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+
         bytes32 taskId = calcTaskId(task);
 
-        if (!_taskIds.add(taskId)) revert TaskAlreadyRegistered();
+        if (!$.taskIds.add(taskId)) revert TaskAlreadyRegistered();
 
-        _tasks[taskId] = task;
+        $.tasks[taskId] = task;
 
         emit RegisterTask(taskId, msg.sender, task.uri);
 
@@ -63,23 +89,27 @@ contract TaskRegistry is ITaskRegistry, Ownable {
 
     /// @inheritdoc ITaskRegistry
     function getTaskById(bytes32 taskId) external view returns (Task memory) {
-        return _tasks[taskId];
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        return $.tasks[taskId];
     }
 
     /// @inheritdoc ITaskRegistry
     function getTaskByIdEncoded(bytes32 taskId) external view returns (bytes memory) {
-        Task memory task = _tasks[taskId];
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        Task memory task = $.tasks[taskId];
         return encodeTask(task);
     }
 
     /// @inheritdoc ITaskRegistry
     function taskIds() external view returns (bytes32[] memory) {
-        return _taskIds.values();
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        return $.taskIds.values();
     }
 
     /// @inheritdoc ITaskRegistry
     function isTaskId(bytes32 taskId) public view returns (bool) {
-        return _taskIds.contains(taskId);
+        TaskRegistryStorage storage $ = _getTaskRegistryStorage();
+        return $.taskIds.contains(taskId);
     }
 
     /// @inheritdoc ITaskRegistry

diff --git a/script/DeployAll.s.sol b/script/DeployAll.s.sol
@@ -104,6 +104,9 @@ contract DeployAll is Script {
         // init hubDomainsRegistry now that hubRegistry is deployed
         hubDomainsRegistry.initialize(address(hubRegistry), "Hub Domains Registry", "HDR");
 
+        // other inits
+        taskRegistry.initialize();
+
         // Setup initial tasks
         Task[] memory tasks = new Task[](3);
         // open tasks