✨ Add uint256 array functions to DynamicArrayLib (#1083)

src/utils/DynamicArrayLib.sol

@@ -16,7 +16,90 @@ library DynamicArrayLib {
     }
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
-    /*                         OPERATIONS                         */
+    /*                  UINT256 ARRAY OPERATIONS                  */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    // Low level minimalist uint256 array operations.
+    // If you don't need syntax sugar, it's recommended to use these.
+    // Some of these functions returns the same array for function chaining.
+    // `e.g. `array.set(0, 1).set(1, 2)`.
+
+    /// @dev Returns a uint256 array with `n` elements. The elements are not zeroized.
+    function malloc(uint256 n) internal pure returns (uint256[] memory result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := mload(0x40)
+            mstore(result, n)
+            mstore(0x40, add(add(result, 0x20), shl(5, n)))
+            if iszero(lt(n, 0xffffffff)) { invalid() }
+        }
+    }
+
+    /// @dev Zeroizes all the elements of `array`.
+    function zeroize(uint256[] memory array) internal pure returns (uint256[] memory result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := array
+            codecopy(add(result, 0x20), codesize(), shl(5, mload(result)))
+        }
+    }
+
+    /// @dev Returns the element at `array[i]`, without bounds checking.
+    function get(uint256[] memory array, uint256 i) internal pure returns (uint256 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := mload(add(add(array, 0x20), shl(5, i)))
+        }
+    }
+
+    /// @dev Sets `array[i]` to `data`, without bounds checking.
+    function set(uint256[] memory array, uint256 i, uint256 data)
+        internal
+        pure
+        returns (uint256[] memory result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := array
+            mstore(add(add(result, 0x20), shl(5, i)), data)
+        }
+    }
+
+    /// @dev Reduces the size of `array` to `n`.
+    /// If `n` is greater than the size of `array`, this will be a no-op.
+    function truncate(uint256[] memory array, uint256 n)
+        internal
+        pure
+        returns (uint256[] memory result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := array
+            mstore(mul(lt(n, mload(result)), result), n)
+        }
+    }
+
+    /// @dev Clears the array and attempts to free the memory if possible.
+    function free(uint256[] memory array) internal pure returns (uint256[] memory result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := array
+            let n := mload(result)
+            mstore(shl(6, lt(iszero(n), eq(add(shl(5, add(1, n)), result), mload(0x40)))), result)
+            mstore(result, 0)
+        }
+    }
+
+    /// @dev Equivalent to `keccak256(abi.encodePacked(array))`.
+    function hash(uint256[] memory array) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := keccak256(add(array, 0x20), shl(5, mload(array)))
+        }
+    }
+
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                  DYNAMIC ARRAY OPERATIONS                  */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
     // Some of these functions returns the same array for function chaining.
@@ -49,8 +132,8 @@ library DynamicArrayLib {
                 let cap := mload(sub(arrData, 0x20))
                 // Extract `cap`, initializing it to zero if it is not a multiple of `prime`.
                 cap := mul(div(cap, prime), iszero(mod(cap, prime)))
-                // If the memory is contiguous, we can free it.
-                if eq(mload(0x40), add(arrData, add(0x20, cap))) {
+                // If `cap` is non-zero and the memory is contiguous, we can free it.
+                if lt(iszero(cap), eq(mload(0x40), add(arrData, add(0x20, cap)))) {
                     mstore(0x40, sub(arrData, 0x20))
                 }
                 mstore(result, 0x60)
@@ -396,7 +479,7 @@ library DynamicArrayLib {
         result = array;
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(add(add(mload(array), 0x20), shl(5, i)), data)
+            mstore(add(add(mload(result), 0x20), shl(5, i)), data)
         }
     }
 
@@ -410,7 +493,7 @@ library DynamicArrayLib {
         result = array;
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(add(add(mload(array), 0x20), shl(5, i)), shr(96, shl(96, data)))
+            mstore(add(add(mload(result), 0x20), shl(5, i)), shr(96, shl(96, data)))
         }
     }
 
@@ -424,7 +507,7 @@ library DynamicArrayLib {
         result = array;
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(add(add(mload(array), 0x20), shl(5, i)), iszero(iszero(data)))
+            mstore(add(add(mload(result), 0x20), shl(5, i)), iszero(iszero(data)))
         }
     }
 
@@ -438,7 +521,7 @@ library DynamicArrayLib {
         result = array;
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(add(add(mload(array), 0x20), shl(5, i)), data)
+            mstore(add(add(mload(result), 0x20), shl(5, i)), data)
         }
     }
 

test/DynamicArrayLib.t.sol

@@ -6,6 +6,7 @@ import {DynamicArrayLib} from "../src/utils/DynamicArrayLib.sol";
 
 contract DynamicArrayLibTest is SoladyTest {
     using DynamicArrayLib for DynamicArrayLib.DynamicArray;
+    using DynamicArrayLib for uint256[];
 
     function testDynamicArrayPushAndPop() public {
         uint256 n = 100;
@@ -213,6 +214,73 @@ contract DynamicArrayLibTest is SoladyTest {
         }
     }
 
+    function testUint256ArrayOperations(uint256 n, uint256 r) public {
+        unchecked {
+            n = _bound(n, 0, 50);
+            uint256[] memory a = DynamicArrayLib.malloc(n);
+            assertEq(a.length, n);
+            _checkMemory(a);
+            for (uint256 i; i != n; ++i) {
+                a.set(i, i ^ r);
+            }
+            for (uint256 i; i != n; ++i) {
+                assertEq(a.get(i), i ^ r);
+            }
+            if (_randomChance(32)) {
+                DynamicArrayLib.DynamicArray memory b;
+                for (uint256 i; i != n; ++i) {
+                    b.p(i ^ r);
+                }
+                assertEq(b.hash(), a.hash());
+                if (_randomChance(2)) {
+                    assertEq(b.resize(0).resize(n).hash(), a.zeroize().hash());
+                }
+            }
+            if (n > 5 && _randomChance(8)) {
+                a.set(0, 1).set(1, 2);
+                assertEq(a.get(0), 1);
+                assertEq(a.get(1), 2);
+            }
+            uint256 lengthBefore = n;
+            n = _bound(_random(), 0, 50);
+            if (n < lengthBefore) {
+                assertEq(a.truncate(n).length, n);
+            } else {
+                assertEq(a.truncate(n).length, lengthBefore);
+            }
+            if (_randomChance(2)) {
+                assertEq(a.free().length, 0);
+                _checkMemory(a);
+            }
+        }
+    }
+
+    function testDynamicArraySetAndGet(bytes32, uint256 i, uint256 n) public {
+        DynamicArrayLib.DynamicArray memory a;
+        n = _bound(n, 1, 5);
+        a.resize(n);
+        {
+            i = _bound(i, 0, n - 1);
+            address data = _randomHashedAddress();
+            assertEq(a.set(i, data).getAddress(i), data);
+        }
+        {
+            i = _bound(i, 0, n - 1);
+            bool data = _randomChance(2);
+            assertEq(a.set(i, data).getBool(i), data);
+        }
+        {
+            i = _bound(i, 0, n - 1);
+            bytes32 data = bytes32(_random());
+            assertEq(a.set(i, data).getBytes32(i), data);
+        }
+        {
+            i = _bound(i, 0, n - 1);
+            uint256 data = _random();
+            assertEq(a.set(i, data).get(i), data);
+        }
+    }
+
     function _sliceOriginal(uint256[] memory a, uint256 start, uint256 end)
         internal
         pure