feat: improve UnlockQueue API

Tenderize · Nov 30, 2023 · 7541328 · 7541328
1 parent 078bcb9
commit 7541328
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 38 deletions.
diff --git a/src/UnlockQueue.sol b/src/UnlockQueue.sol
@@ -34,27 +34,45 @@ library UnlockQueue {
     }
 
     struct Data {
-        uint256 head; // oldest element
-        uint256 tail; // newest element
+        uint256 _head; // oldest element
+        uint256 _tail; // newest element
         mapping(uint256 index => Node) nodes; // elements as a map
     }
 
+    /**
+     * @notice Get the oldest element in the queue
+     * @param q The queue to query
+     * @return The oldest element in the queue
+     */
+    function head(UnlockQueue.Data storage q) internal view returns (Node memory) {
+        return q.nodes[q._head];
+    }
+
+    /**
+     * @notice Get the newest element in the queue
+     * @param q The queue to query
+     * @return The newest element in the queue
+     */
+    function tail(UnlockQueue.Data storage q) internal view returns (Node memory) {
+        return q.nodes[q._tail];
+    }
+
     /**
      * @notice Pop the oldest element from the queue
      * @param q The queue to pop from
      */
-    function popFront(UnlockQueue.Data storage q) internal returns (Item memory unlock) {
-        uint256 head = q.head;
+    function popHead(UnlockQueue.Data storage q) internal returns (Node memory node) {
+        uint256 head = q._head;
         if (head == 0) revert QueueEmpty();
 
-        unlock = q.nodes[head].data;
+        node = q.nodes[head];
 
         uint256 next = q.nodes[head].next;
         if (next == 0) {
-            q.head = 0;
-            q.tail = 0;
+            q._head = 0;
+            q._tail = 0;
         } else {
-            q.head = next;
+            q._head = next;
             q.nodes[next].prev = 0;
         }
 
@@ -65,18 +83,18 @@ library UnlockQueue {
      * @notice Pop the newest element from the queue
      * @param q The queue to pop from
      */
-    function popBack(UnlockQueue.Data storage q) internal returns (Item memory unlock) {
-        uint256 tail = q.tail;
+    function popTail(UnlockQueue.Data storage q) internal returns (Node memory node) {
+        uint256 tail = q._tail;
         if (tail == 0) revert QueueEmpty();
 
-        unlock = q.nodes[tail].data;
+        node = q.nodes[tail];
 
         uint256 prev = q.nodes[tail].prev;
         if (prev == 0) {
-            q.head = 0;
-            q.tail = 0;
+            q._head = 0;
+            q._tail = 0;
         } else {
-            q.tail = prev;
+            q._tail = prev;
             q.nodes[prev].next = 0;
         }
 
@@ -89,18 +107,18 @@ library UnlockQueue {
      * @param unlock The unlock data to push
      */
     function push(UnlockQueue.Data storage q, Item memory unlock) internal {
-        uint256 tail = q.tail;
+        uint256 tail = q._tail;
         uint256 newTail = unlock.id;
 
         q.nodes[newTail].data = unlock;
         q.nodes[newTail].prev = tail;
 
         if (tail == 0) {
-            q.head = newTail;
+            q._head = newTail;
         } else {
             q.nodes[tail].next = newTail;
         }
 
-        q.tail = newTail;
+        q._tail = newTail;
     }
 }
diff --git a/test/UnlockQueue.harness.sol b/test/UnlockQueue.harness.sol
@@ -24,19 +24,19 @@ contract UnlockQueueHarness {
         queue.push(item);
     }
 
-    function exposed_popFront() public returns (UnlockQueue.Item memory item) {
-        item = queue.popFront();
+    function exposed_popHead() public returns (UnlockQueue.Node memory node) {
+        node = queue.popHead();
     }
 
-    function exposed_popBack() public returns (UnlockQueue.Item memory item) {
-        item = queue.popBack();
+    function exposed_popTail() public returns (UnlockQueue.Node memory node) {
+        node = queue.popTail();
     }
 
-    function exposed_head() public view returns (UnlockQueue.Item memory item) {
-        item = queue.nodes[queue.head].data;
+    function exposed_head() public view returns (UnlockQueue.Node memory node) {
+        node = queue.nodes[queue._head];
     }
 
-    function exposed_tail() public view returns (UnlockQueue.Item memory item) {
-        item = queue.nodes[queue.tail].data;
+    function exposed_tail() public view returns (UnlockQueue.Node memory node) {
+        node = queue.nodes[queue._tail];
     }
 }
diff --git a/test/UnlockQueue.t.sol b/test/UnlockQueue.t.sol
@@ -30,28 +30,34 @@ contract UnlockQueueTest is Test {
         queue.exposed_push(item1);
 
         // assert
-        assertEq(queue.exposed_head().id, 1);
-        assertEq(queue.exposed_tail().id, 1);
+        assertEq(queue.exposed_head().data.id, 1);
+        assertEq(queue.exposed_tail().data.id, 1);
 
         queue.exposed_push(item2);
         // assert
-        assertEq(queue.exposed_head().id, 1);
-        assertEq(queue.exposed_tail().id, 2);
+        assertEq(queue.exposed_head().data.id, 1);
+        assertEq(queue.exposed_head().prev, 0);
+        assertEq(queue.exposed_head().next, 2);
+        assertEq(queue.exposed_tail().data.id, 2);
 
         queue.exposed_push(item3);
         // assert
-        assertEq(queue.exposed_head().id, 1);
-        assertEq(queue.exposed_tail().id, 3);
+        assertEq(queue.exposed_head().data.id, 1);
+        assertEq(queue.exposed_head().prev, 0);
+        assertEq(queue.exposed_head().next, 2);
+        assertEq(queue.exposed_tail().data.id, 3);
+        assertEq(queue.exposed_tail().prev, 2);
+        assertEq(queue.exposed_tail().next, 0);
 
         // pop front
-        UnlockQueue.Item memory popped = queue.exposed_popFront();
-        assertEq(popped.id, 1);
-        assertEq(queue.exposed_head().id, 2);
+        UnlockQueue.Node memory popped = queue.exposed_popHead();
+        assertEq(popped.data.id, 1);
+        assertEq(queue.exposed_head().data.id, 2);
 
         // pop back
-        popped = queue.exposed_popBack();
-        assertEq(popped.id, 3);
-        assertEq(queue.exposed_head().id, 2);
-        assertEq(queue.exposed_tail().id, 2);
+        popped = queue.exposed_popTail();
+        assertEq(popped.data.id, 3);
+        assertEq(queue.exposed_head().data.id, 2);
+        assertEq(queue.exposed_tail().data.id, 2);
     }
 }