Merge pull request #21 from morpho-labs/test/bucket-bit-manipulation

test: formal verification of bit manipulation in bucket computation

Author: MathisGD

.github/workflows/ci-foundry.yml

@@ -37,3 +37,9 @@ jobs:
 
       - name: Run tests
         run: forge test
+
+      - name: Install symtest
+        run: python3 -m pip install --upgrade symtest
+
+      - name: Run symtest
+        run: symtest --function testProve --loop 4

src/LogarithmicBuckets.sol

@@ -39,17 +39,17 @@ library LogarithmicBuckets {
 
         if (value == 0) {
             if (_newValue == 0) revert ZeroValue();
-            _insert(_buckets, _id, _computeBucket(_newValue), _head);
+            _insert(_buckets, _id, computeBucket(_newValue), _head);
             return;
         }
 
-        uint256 currentBucket = _computeBucket(value);
+        uint256 currentBucket = computeBucket(value);
         if (_newValue == 0) {
             _remove(_buckets, _id, currentBucket);
             return;
         }
 
-        uint256 newBucket = _computeBucket(_newValue);
+        uint256 newBucket = computeBucket(_newValue);
         if (newBucket != currentBucket) {
             _remove(_buckets, _id, currentBucket);
             _insert(_buckets, _id, newBucket, _head);
@@ -71,7 +71,7 @@ library LogarithmicBuckets {
         view
         returns (BucketDLL.List storage)
     {
-        uint256 bucket = _computeBucket(_value);
+        uint256 bucket = computeBucket(_value);
         return _buckets.lists[bucket];
     }
 
@@ -82,13 +82,13 @@ library LogarithmicBuckets {
     function getMatch(BucketList storage _buckets, uint256 _value) internal view returns (address) {
         uint256 bucketsMask = _buckets.bucketsMask;
         if (bucketsMask == 0) return address(0);
-        uint256 lowerMask = _setLowerBits(_value);
+        uint256 lowerMask = setLowerBits(_value);
 
-        uint256 next = _nextBucket(lowerMask, bucketsMask);
+        uint256 next = nextBucket(lowerMask, bucketsMask);
 
         if (next != 0) return _buckets.lists[next].getHead();
 
-        uint256 prev = _prevBucket(lowerMask, bucketsMask);
+        uint256 prev = prevBucket(lowerMask, bucketsMask);
 
         return _buckets.lists[prev].getHead();
     }
@@ -126,14 +126,14 @@ library LogarithmicBuckets {
     }
 
     /// @notice Returns the bucket in which the given value would fall.
-    function _computeBucket(uint256 _value) private pure returns (uint256) {
-        uint256 lowerMask = _setLowerBits(_value);
+    function computeBucket(uint256 _value) internal pure returns (uint256) {
+        uint256 lowerMask = setLowerBits(_value);
         return lowerMask ^ (lowerMask >> 1);
     }
 
     /// @notice Sets all the bits lower than (or equal to) the highest bit in the input.
     /// @dev This is the same as rounding the input the nearest upper value of the form `2 ** n - 1`.
-    function _setLowerBits(uint256 x) private pure returns (uint256 y) {
+    function setLowerBits(uint256 x) internal pure returns (uint256 y) {
         assembly {
             x := or(x, shr(1, x))
             x := or(x, shr(2, x))
@@ -148,8 +148,8 @@ library LogarithmicBuckets {
 
     /// @notice Returns the following bucket which contains greater values.
     /// @dev The bucket returned is the lowest that is in `bucketsMask` and not in `lowerMask`.
-    function _nextBucket(uint256 lowerMask, uint256 bucketsMask)
-        private
+    function nextBucket(uint256 lowerMask, uint256 bucketsMask)
+        internal
         pure
         returns (uint256 bucket)
     {
@@ -161,8 +161,8 @@ library LogarithmicBuckets {
 
     /// @notice Returns the preceding bucket which contains smaller values.
     /// @dev The bucket returned is the highest that is in both `bucketsMask` and `lowerMask`.
-    function _prevBucket(uint256 lowerMask, uint256 bucketsMask) private pure returns (uint256) {
-        uint256 lowerBucketsMask = _setLowerBits(lowerMask & bucketsMask);
+    function prevBucket(uint256 lowerMask, uint256 bucketsMask) internal pure returns (uint256) {
+        uint256 lowerBucketsMask = setLowerBits(lowerMask & bucketsMask);
         return lowerBucketsMask ^ (lowerBucketsMask >> 1);
     }
 }

test/TestLogarithmicBuckets.t.sol

@@ -115,4 +115,54 @@ contract TestLogarithmicBuckets is LogarithmicBucketsMock, Test {
         assertEq(bucketList.getMatch(16), accounts[0], "head equal");
         assertEq(bucketList.getMatch(32), accounts[0], "head above");
     }
+
+    function isPowerOfTwo(uint256 x) public pure returns (bool) {
+        unchecked {
+            return x != 0 && (x & (x - 1)) == 0;
+        }
+    }
+
+    function testProveComputeBucket(uint256 _value) public {
+        uint256 bucket = LogarithmicBuckets.computeBucket(_value);
+        unchecked {
+            // cross-check that bucket == 2^{floor(log_2 value)}, or 0 if value == 0
+            assertTrue(bucket == 0 || isPowerOfTwo(bucket));
+            assertTrue(bucket <= _value);
+            assertTrue(_value <= 2 * bucket - 1); // abusing overflow when bucket == 2**255
+        }
+    }
+
+    function testProveNextBucket(uint256 _value) public {
+        uint256 curr = LogarithmicBuckets.computeBucket(_value);
+        uint256 next = nextBucket(_value);
+        uint256 bucketsMask = bucketList.bucketsMask;
+        // check that `next` is a strictly higer non-empty bucket, or zero
+        assertTrue(next == 0 || isPowerOfTwo(next));
+        assertTrue(next == 0 || next > curr);
+        assertTrue(next == 0 || bucketsMask & next != 0);
+        unchecked {
+            // check that `next` is the lowest one among such higher non-empty buckets, if exist
+            // note: this also checks that all the higher buckets are empty when `next` == 0
+            for (uint256 i = curr << 1; i != next; i <<= 1) {
+                assertTrue(bucketsMask & i == 0);
+            }
+        }
+    }
+
+    function testProvePrevBucket(uint256 _value) public {
+        uint256 curr = LogarithmicBuckets.computeBucket(_value);
+        uint256 prev = prevBucket(_value);
+        uint256 bucketsMask = bucketList.bucketsMask;
+        // check that `prev` is a non-empty bucket that is lower than or equal to `curr`; or zero
+        assertTrue(prev == 0 || isPowerOfTwo(prev));
+        assertTrue(prev <= curr);
+        assertTrue(prev == 0 || bucketsMask & prev != 0);
+        unchecked {
+            // check that `prev` is the highest one among such lower non-empty buckets, if exist
+            // note: this also checks that all the lower buckets are empty when `prev` == 0
+            for (uint256 i = curr; i > prev; i >>= 1) {
+                assertTrue(bucketsMask & i == 0);
+            }
+        }
+    }
 }

test/mocks/LogarithmicBucketsMock.sol

@@ -21,21 +21,8 @@ contract LogarithmicBucketsMock {
         return bucketList.getValueOf(_id);
     }
 
-    function _setLowerBits(uint256 x) private pure returns (uint256 y) {
-        assembly {
-            x := or(x, shr(1, x))
-            x := or(x, shr(2, x))
-            x := or(x, shr(4, x))
-            x := or(x, shr(8, x))
-            x := or(x, shr(16, x))
-            x := or(x, shr(32, x))
-            x := or(x, shr(64, x))
-            y := or(x, shr(128, x))
-        }
-    }
-
     function maxBucket() public view returns (uint256) {
-        uint256 lowerMask = _setLowerBits(bucketList.bucketsMask);
+        uint256 lowerMask = LogarithmicBuckets.setLowerBits(bucketList.bucketsMask);
         return lowerMask ^ (lowerMask >> 1);
     }
 
@@ -60,4 +47,16 @@ contract LogarithmicBucketsMock {
         }
         return true;
     }
+
+    function nextBucket(uint256 _value) internal view returns (uint256) {
+        uint256 bucketsMask = bucketList.bucketsMask;
+        uint256 lowerMask = LogarithmicBuckets.setLowerBits(_value);
+        return LogarithmicBuckets.nextBucket(lowerMask, bucketsMask);
+    }
+
+    function prevBucket(uint256 _value) internal view returns (uint256) {
+        uint256 bucketsMask = bucketList.bucketsMask;
+        uint256 lowerMask = LogarithmicBuckets.setLowerBits(_value);
+        return LogarithmicBuckets.prevBucket(lowerMask, bucketsMask);
+    }
 }