fix: minGreaterThan logic

clober-dex · Jan 12, 2024 · ec1868a · ec1868a
1 parent 651638e
commit ec1868a
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 13 deletions.
diff --git a/contracts/Heap.sol b/contracts/Heap.sol
@@ -12,6 +12,7 @@ library Heap {
     uint256 public constant EMPTY_ERROR = 0;
     uint256 public constant ALREADY_EXISTS_ERROR = 1;
     uint256 public constant B0_BITMAP_KEY = uint256(keccak256("heap"));
+    uint256 public constant MAX_UINT_256_MINUS_1 = 0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffe;
 
     function has(mapping(uint256 => uint256) storage heap, uint24 value) internal view returns (bool) {
         (uint256 b0b1, uint256 b2) = _split(value);
@@ -84,22 +85,22 @@ library Heap {
 
     function minGreaterThan(mapping(uint256 => uint256) storage heap, uint24 value) internal view returns (uint24) {
         (uint256 b0b1, uint256 b2) = _split(value);
-        uint256 b2Bitmap = ((type(uint256).max >> b2) << b2) & heap[b0b1];
+        uint256 b2Bitmap = (MAX_UINT_256_MINUS_1 << b2) & heap[b0b1];
         if (b2Bitmap == 0) {
             uint256 b0 = b0b1 >> 8;
-            uint256 b1 = b0b1 & 0xff;
-            uint256 b1Bitmap = ((type(uint256).max >> b1) << b1) & heap[~b0];
+            uint256 b1Bitmap = (MAX_UINT_256_MINUS_1 << (b0b1 & 0xff)) & heap[~b0];
             if (b1Bitmap == 0) {
-                uint256 b0Bitmap = ((type(uint256).max >> b0) << b0) & heap[B0_BITMAP_KEY];
+                uint256 b0Bitmap = (MAX_UINT_256_MINUS_1 << b0) & heap[B0_BITMAP_KEY];
                 if (b0Bitmap == 0) {
                     revert HeapError(EMPTY_ERROR);
                 }
                 b0 = b0Bitmap.leastSignificantBit();
+                b1Bitmap = heap[~b0];
             }
-            b1 = b1Bitmap.leastSignificantBit();
-            b0b1 = (b0 << 8) | b1;
+            b0b1 = (b0 << 8) | b1Bitmap.leastSignificantBit();
+            b2Bitmap = heap[b0b1];
         }
-        b2 = b0b1.leastSignificantBit();
+        b2 = b2Bitmap.leastSignificantBit();
         return uint24((b0b1 << 8) | b2);
     }
 }
diff --git a/contracts/interfaces/CloberMinBitMap.sol → contracts/interfaces/CloberHeap.sol b/contracts/interfaces/CloberMinBitMap.sol → contracts/interfaces/CloberHeap.sol
diff --git a/contracts/mocks/HeapWrapper.sol b/contracts/mocks/HeapWrapper.sol
@@ -4,7 +4,7 @@
 pragma solidity ^0.8.0;
 
 import "../Heap.sol";
-import "../interfaces/CloberMinBitMap.sol";
+import "../interfaces/CloberHeap.sol";
 
 contract HeapWrapper is CloberHeap {
     using Heap for mapping(uint256 => uint256);
@@ -31,7 +31,7 @@ contract HeapWrapper is CloberHeap {
         return _heap.root();
     }
 
-    function getNextMinValue(uint24 value) external view returns (uint24) {
+    function minGreaterThan(uint24 value) external view returns (uint24) {
         return _heap.minGreaterThan(value);
     }
 }
diff --git a/test/foundry/Heap.t.sol b/test/foundry/Heap.t.sol
@@ -18,28 +18,41 @@ contract HeapTest is Test {
         _min = type(uint24).max;
     }
 
-    function _push(uint24[] calldata numbers) private {
+    function _push(uint24[] memory numbers) private returns (uint256 length) {
         for (uint32 i = 0; i < numbers.length; ++i) {
             uint24 number = numbers[i];
             if (testWrapper.has(number)) continue;
             if (number < _min) _min = number;
 
             assertFalse(testWrapper.has(number), "BEFORE_PUSH");
             testWrapper.push(number);
+            length += 1;
             assertTrue(testWrapper.has(number), "AFTER_PUSH");
             assertEq(testWrapper.root(), _min, "ASSERT_MIN");
         }
     }
 
-    function testPop(uint24[] calldata numbers) public {
+    function testPopxx(uint24[] calldata numbers) public {
         vm.assume(1 <= numbers.length && numbers.length <= _MAX_HEAP_SIZE);
         assertTrue(testWrapper.isEmpty(), "HAS_TO_BE_EMPTY");
-        _push(numbers);
+        uint256 length = _push(numbers);
         assertFalse(testWrapper.isEmpty(), "HAS_TO_BE_OCCUPIED");
         while (!testWrapper.isEmpty()) {
             _min = testWrapper.root();
-            assertTrue(testWrapper.has(_min), "BITMAP_HAS_ROOT");
+            assertTrue(testWrapper.has(_min), "HEAP_HAS_ROOT");
+            uint256 min;
+            if (length == 1) {
+                vm.expectRevert(abi.encodeWithSelector(Heap.HeapError.selector, Heap.EMPTY_ERROR));
+                testWrapper.minGreaterThan(_min);
+            } else {
+                min = testWrapper.minGreaterThan(_min);
+            }
             testWrapper.pop();
+            length -= 1;
+            if (length > 0) {
+                assertTrue(testWrapper.root() == min, "WRONG_MIN");
+            }
+
             assertFalse(testWrapper.has(_min), "ROOT_HAS_BEEN_POPPED");
             if (testWrapper.isEmpty()) break;
             assertGt(testWrapper.root(), _min, "ROOT_HAS_TO_BE_MIN");