Add remove function to the PriorityQueue (#147)

* Added `remove` function to the PriorityQueue

* Added `remove` function to the PriorityQueue

* optimize remove()

---------

Co-authored-by: Artem Chystiakov <artem.ch31@gmail.com>

Author: KyrylR

contracts/libs/data-structures/PriorityQueue.sol

@@ -8,6 +8,7 @@ import {TypeCaster} from "../utils/TypeCaster.sol";
  *
  * Courtesy of heap property,
  * add() and removeTop() operations are O(log(n)) complex
+ * remove() operation is O(n) for search + O(log(n)) for removal
  * top(), topValue() operations are O(1)
  *
  * The library might be useful to implement priority withdrawals/purchases, reputation based systems, and similar logic.
@@ -60,6 +61,17 @@ library PriorityQueue {
         _removeTop(queue._queue);
     }
 
+    /**
+     * @notice The function to remove a specific element from the uint256 queue. O(n) complex for search + O(log(n)) for removal
+     * In case of duplicate elements, removes the one with the highest priority.
+     * @param queue self
+     * @param value_ the element value to remove
+     * @return true if element was found and removed, false otherwise
+     */
+    function remove(UintQueue storage queue, uint256 value_) internal returns (bool) {
+        return _remove(queue._queue, bytes32(value_));
+    }
+
     /**
      * @notice The function to read the value of the element with the highest priority. O(1) complex
      * @param queue self
@@ -136,6 +148,14 @@ library PriorityQueue {
         _removeTop(queue._queue);
     }
 
+    /**
+     * @notice The function to remove a specific element from the bytes32 queue. O(n) complex for search + O(log(n)) for removal
+     * In case of duplicate elements, removes the one with the highest priority.
+     */
+    function remove(Bytes32Queue storage queue, bytes32 value_) internal returns (bool) {
+        return _remove(queue._queue, value_);
+    }
+
     /**
      * @notice The function to read the value of the element with the highest priority. O(1) complex
      */
@@ -200,6 +220,14 @@ library PriorityQueue {
         _removeTop(queue._queue);
     }
 
+    /**
+     * @notice The function to remove a specific element from the address queue. O(n) complex for search + O(log(n)) for removal
+     * In case of duplicate elements, removes the one with the highest priority.
+     */
+    function remove(AddressQueue storage queue, address value_) internal returns (bool) {
+        return _remove(queue._queue, bytes32(uint256(uint160(value_))));
+    }
+
     /**
      * @notice The function to read the value of the element with the highest priority. O(1) complex
      */
@@ -273,6 +301,44 @@ library PriorityQueue {
         _shiftDown(queue, 0);
     }
 
+    function _remove(Queue storage queue, bytes32 value_) private returns (bool) {
+        uint256 length_ = _length(queue);
+
+        if (length_ == 0) {
+            return false;
+        }
+
+        uint256 indexToRemove_ = type(uint256).max;
+
+        for (uint256 i = 0; i < length_; ++i) {
+            if (queue._values[i] == value_) {
+                indexToRemove_ = i;
+                break;
+            }
+        }
+
+        if (indexToRemove_ == type(uint256).max) {
+            return false;
+        }
+
+        if (indexToRemove_ == length_ - 1) {
+            queue._values.pop();
+            queue._priorities.pop();
+
+            return true;
+        }
+
+        queue._values[indexToRemove_] = queue._values[length_ - 1];
+        queue._priorities[indexToRemove_] = queue._priorities[length_ - 1];
+
+        queue._values.pop();
+        queue._priorities.pop();
+
+        _shiftDown(queue, indexToRemove_);
+
+        return true;
+    }
+
     function _topValue(Queue storage queue) private view returns (bytes32) {
         _requireNotEmpty(queue);
 

contracts/mock/libs/data-structures/PriorityQueueMock.sol

@@ -21,6 +21,10 @@ contract PriorityQueueMock {
         _uintQueue.removeTop();
     }
 
+    function removeUint(uint256 value_) external returns (bool) {
+        return _uintQueue.remove(value_);
+    }
+
     function topValueUint() external view returns (uint256) {
         return _uintQueue.topValue();
     }
@@ -51,6 +55,10 @@ contract PriorityQueueMock {
         _bytes32Queue.removeTop();
     }
 
+    function removeBytes32(bytes32 value_) external returns (bool) {
+        return _bytes32Queue.remove(value_);
+    }
+
     function topValueBytes32() external view returns (bytes32) {
         return _bytes32Queue.topValue();
     }
@@ -81,6 +89,10 @@ contract PriorityQueueMock {
         _addressQueue.removeTop();
     }
 
+    function removeAddress(address value_) external returns (bool) {
+        return _addressQueue.remove(value_);
+    }
+
     function topValueAddress() external view returns (address) {
         return _addressQueue.topValue();
     }

test/libs/data-structures/PriorityQueue.test.ts

@@ -96,7 +96,7 @@ describe("PriorityQueue", () => {
     });
   });
 
-  describe("remove()", async () => {
+  describe("removeTop()", async () => {
     describe("uint", () => {
       it("should add and remove the top elements", async () => {
         await mock.addUint(1, 1);
@@ -190,4 +190,153 @@ describe("PriorityQueue", () => {
       });
     });
   });
+
+  describe("remove()", () => {
+    describe("uint", () => {
+      it("should remove a specific element from the middle", async () => {
+        await mock.addUint(1, 1);
+        await mock.addUint(2, 2);
+        await mock.addUint(3, 3);
+        await mock.addUint(4, 4);
+        await mock.addUint(5, 5);
+
+        expect(await mock.lengthUint()).to.equal(5n);
+        expect(await mock.topValueUint()).to.equal(5n);
+
+        // Remove element with value 3
+        await mock.removeUint(3);
+
+        expect(await mock.lengthUint()).to.equal(4n);
+        expect(await mock.topValueUint()).to.equal(5n);
+        expect(await mock.valuesUint()).to.deep.equal([5n, 4n, 2n, 1n]);
+      });
+
+      it("should remove the top element", async () => {
+        await mock.addUint(1, 1);
+        await mock.addUint(2, 2);
+        await mock.addUint(3, 3);
+
+        // Remove the top element (value 3, priority 3)
+        await mock.removeUint(3);
+
+        expect(await mock.lengthUint()).to.equal(2n);
+        expect(await mock.topValueUint()).to.equal(2n);
+      });
+
+      it("should remove the last element", async () => {
+        await mock.addUint(1, 1);
+        await mock.addUint(2, 2);
+        await mock.addUint(3, 3);
+
+        // Remove element with value 1 (should be at the end)
+        await mock.removeUint(1);
+
+        expect(await mock.lengthUint()).to.equal(2n);
+        expect(await mock.topValueUint()).to.equal(3n);
+      });
+
+      it("should remove the element with the same priority", async () => {
+        await mock.addUint(1, ethers.MaxUint256);
+        await mock.addUint(2, ethers.MaxUint256);
+        await mock.addUint(3, ethers.MaxUint256);
+
+        await mock.removeUint(2);
+
+        expect(await mock.lengthUint()).to.equal(2n);
+        expect(await mock.topValueUint()).to.equal(1n);
+
+        await mock.removeUint(3);
+
+        expect(await mock.lengthUint()).to.equal(1n);
+        expect(await mock.topValueUint()).to.equal(1n);
+      });
+
+      it("should return false when removing non-existent element", async () => {
+        await mock.addUint(1, 1);
+        await mock.addUint(2, 2);
+
+        // Check the return value using staticCall (view function call)
+        expect(await mock.removeUint.staticCall(99)).to.equal(false);
+        expect(await mock.lengthUint()).to.equal(2n);
+      });
+
+      it("should return false when removing from empty queue", async () => {
+        expect(await mock.removeUint.staticCall(1)).to.equal(false);
+        expect(await mock.lengthUint()).to.equal(0n);
+      });
+
+      it("should handle removing elements with same priority", async () => {
+        await mock.addUint(1, 5);
+        await mock.addUint(2, 5);
+        await mock.addUint(3, 5);
+
+        await mock.removeUint(2);
+        expect(await mock.lengthUint()).to.equal(2n);
+      });
+
+      it("should handle removing all elements one by one", async () => {
+        await mock.addUint(1, 1);
+        await mock.addUint(2, 2);
+        await mock.addUint(3, 3);
+
+        await mock.removeUint(2);
+        expect(await mock.lengthUint()).to.equal(2n);
+
+        await mock.removeUint(1);
+        expect(await mock.lengthUint()).to.equal(1n);
+
+        await mock.removeUint(3);
+        expect(await mock.lengthUint()).to.equal(0n);
+      });
+    });
+
+    describe("bytes32", () => {
+      it("should remove a specific bytes32 element", async () => {
+        const val1 = ethers.keccak256("0x01");
+        const val2 = ethers.keccak256("0x02");
+        const val3 = ethers.keccak256("0x03");
+
+        await mock.addBytes32(val1, 1);
+        await mock.addBytes32(val2, 2);
+        await mock.addBytes32(val3, 3);
+
+        await mock.removeBytes32(val2);
+        expect(await mock.lengthBytes32()).to.equal(2n);
+        expect(await mock.topValueBytes32()).to.equal(val3);
+      });
+
+      it("should return false when removing non-existent bytes32 element", async () => {
+        const val1 = ethers.keccak256("0x01");
+        const val2 = ethers.keccak256("0x02");
+
+        await mock.addBytes32(val1, 1);
+
+        expect(await mock.removeBytes32.staticCall(val2)).to.equal(false);
+        expect(await mock.lengthBytes32()).to.equal(1n);
+      });
+    });
+
+    describe("address", () => {
+      it("should remove a specific address element", async () => {
+        const [FIRST, SECOND, THIRD] = await ethers.getSigners();
+
+        await mock.addAddress(FIRST.address, 1);
+        await mock.addAddress(SECOND.address, 2);
+        await mock.addAddress(THIRD.address, 3);
+
+        await mock.removeAddress(SECOND.address);
+        expect(await mock.lengthAddress()).to.equal(2n);
+        expect(await mock.topValueAddress()).to.equal(THIRD.address);
+      });
+
+      it("should return false when removing non-existent address element", async () => {
+        const [FIRST, SECOND] = await ethers.getSigners();
+
+        await mock.addAddress(FIRST.address, 1);
+
+        expect(await mock.removeAddress.staticCall(SECOND.address)).to.equal(false);
+        expect(await mock.lengthAddress()).to.equal(1n);
+      });
+    });
+  });
 });