Optimize DynamicArrayLib

Vectorized · Sep 18, 2024 · 6cd2a73 · 6cd2a73
1 parent 0234e7d
commit 6cd2a73
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/src/utils/DynamicArrayLib.sol b/src/utils/DynamicArrayLib.sol
@@ -30,8 +30,7 @@ library DynamicArrayLib {
         assembly {
             result := mload(0x40)
             mstore(result, n)
-            mstore(0x40, add(add(result, 0x20), shl(5, n)))
-            if iszero(lt(n, 0xffffffff)) { invalid() }
+            mstore(or(sub(0, shr(32, n)), 0x40), add(add(result, 0x20), shl(5, n)))
         }
     }
 
@@ -205,8 +204,7 @@ library DynamicArrayLib {
         result = array;
         /// @solidity memory-safe-assembly
         assembly {
-            if iszero(lt(minimum, 0xffffffff)) { invalid() } // For extra safety.
-            for { let arrData := mload(array) } 1 {} {
+            for { let arrData := mload(or(sub(0, shr(32, minimum)), array)) } 1 {} {
                 // Some random prime number to multiply `cap`, so that
                 // we know that the `cap` is for a dynamic array.
                 // Selected to be larger than any memory pointer realistically.

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);