From 75413280a611d04a53ca6984278edcc47e000943 Mon Sep 17 00:00:00 2001 From: kyriediculous Date: Thu, 30 Nov 2023 20:11:00 +0100 Subject: [PATCH] feat: improve UnlockQueue API --- src/UnlockQueue.sol | 52 ++++++++++++++++++++++++------------ test/UnlockQueue.harness.sol | 16 +++++------ test/UnlockQueue.t.sol | 32 +++++++++++++--------- 3 files changed, 62 insertions(+), 38 deletions(-) diff --git a/src/UnlockQueue.sol b/src/UnlockQueue.sol index 2add175..b5279f9 100644 --- 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 index 3da87da..430c4be 100644 --- 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 index 7d36108..f6093e1 100644 --- 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); } }