⚡️ Optimize DynamicArrayLib (#1084)

Vectorized · Sep 19, 2024 · 8f3d6ce · 8f3d6ce
1 parent 1312da6
commit 8f3d6ce
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 16 deletions.
diff --git a/src/utils/DynamicArrayLib.sol b/src/utils/DynamicArrayLib.sol
@@ -28,10 +28,9 @@ library DynamicArrayLib {
     function malloc(uint256 n) internal pure returns (uint256[] memory result) {
         /// @solidity memory-safe-assembly
         assembly {
-            result := mload(0x40)
+            result := or(sub(0, shr(32, n)), mload(0x40))
             mstore(result, n)
             mstore(0x40, add(add(result, 0x20), shl(5, n)))
-            if iszero(lt(n, 0xffffffff)) { invalid() }
         }
     }
 

diff --git a/test/DynamicArrayLib.t.sol b/test/DynamicArrayLib.t.sol
@@ -214,6 +214,36 @@ contract DynamicArrayLibTest is SoladyTest {
         }
     }
 
+    function testUint256ArrayPopulate() public {
+        unchecked {
+            uint256 n = 100;
+            uint256[] memory a = DynamicArrayLib.malloc(n);
+            for (uint256 i; i != n; ++i) {
+                a.set(i, i);
+            }
+            uint256 sum;
+            for (uint256 i; i != n; ++i) {
+                sum += a.get(i);
+            }
+            assertEq(sum, 4950);
+        }
+    }
+
+    function testUint256ArrayPopulateOriginal() public {
+        unchecked {
+            uint256 n = 100;
+            uint256[] memory a = new uint256[](n);
+            for (uint256 i; i != n; ++i) {
+                a[i] = i;
+            }
+            uint256 sum;
+            for (uint256 i; i != n; ++i) {
+                sum += a[i];
+            }
+            assertEq(sum, 4950);
+        }
+    }
+
     function testUint256ArrayOperations(uint256 n, uint256 r) public {
         unchecked {
             n = _bound(n, 0, 50);
@@ -281,6 +311,15 @@ contract DynamicArrayLibTest is SoladyTest {
         }
     }
 
+    function testDynamicArrayFree(uint256 n) public {
+        DynamicArrayLib.DynamicArray memory a;
+        uint256 m = _freeMemoryPointer();
+        n = _bound(n, 0, 50);
+        if (_randomChance(16)) a.reserve(n);
+        a.free();
+        assertEq(m, _freeMemoryPointer());
+    }
+
     function _sliceOriginal(uint256[] memory a, uint256 start, uint256 end)
         internal
         pure

diff --git a/test/DynamicBufferLib.t.sol b/test/DynamicBufferLib.t.sol
@@ -211,20 +211,6 @@ contract DynamicBufferLibTest is SoladyTest {
         assertEq(_freeMemoryPointer(), m);
     }
 
-    function _incrementFreeMemoryPointer() internal pure {
-        /// @solidity memory-safe-assembly
-        assembly {
-            mstore(0x40, add(mload(0x40), 0x20))
-        }
-    }
-
-    function _freeMemoryPointer() internal pure returns (uint256 m) {
-        /// @solidity memory-safe-assembly
-        assembly {
-            m := mload(0x40)
-        }
-    }
-
     function _bufferLocation(DynamicBufferLib.DynamicBuffer memory buffer)
         internal
         pure

diff --git a/test/utils/TestPlus.sol b/test/utils/TestPlus.sol
@@ -565,6 +565,23 @@ contract TestPlus is Brutalizer {
         }
     }
 
+    /// @dev Returns the free memory pointer.
+    function _freeMemoryPointer() internal pure returns (uint256 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            result := mload(0x40)
+        }
+    }
+
+    /// @dev Increments the free memory pointer by a world.
+    function _incrementFreeMemoryPointer() internal pure {
+        uint256 word = 0x20;
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x40, add(mload(0x40), word))
+        }
+    }
+
     /// @dev Adapted from `bound`:
     /// https://github.com/foundry-rs/forge-std/blob/ff4bf7db008d096ea5a657f2c20516182252a3ed/src/StdUtils.sol#L10
     /// Differentially fuzzed tested against the original implementation.