⚡️ Optimize cdCompress and cdDecompress more

src/utils/LibZip.sol

@@ -183,17 +183,23 @@ library LibZip {
                 data := add(data, 1)
                 let c := byte(31, mload(data))
                 if iszero(c) {
-                    let z := 0
                     for {} 1 {} {
-                        let r := 0x20
-                        let x := mload(add(data, r))
-                        if x { r := countLeadingZeroBytes(x) }
-                        r := min(min(sub(end, data), r), sub(0x7f, z))
+                        let x := mload(add(data, 0x20))
+                        if iszero(x) {
+                            let r := min(sub(end, data), 0x20)
+                            r := min(sub(0x7f, c), r)
+                            data := add(data, r)
+                            c := add(c, r)
+                            if iszero(gt(r, 0x1f)) { break }
+                            continue
+                        }
+                        let r := countLeadingZeroBytes(x)
+                        r := min(sub(end, data), r)
                         data := add(data, r)
-                        z := add(z, r)
-                        if iszero(gt(r, 0x1f)) { break }
+                        c := add(c, r)
+                        break
                     }
-                    mstore(o, shl(240, z))
+                    mstore(o, shl(240, c))
                     o := add(o, 2)
                     continue
                 }
@@ -222,27 +228,47 @@ library LibZip {
     function cdDecompress(bytes memory data) internal pure returns (bytes memory result) {
         /// @solidity memory-safe-assembly
         assembly {
+            function countLeadingZeroBytes(x_) -> _r {
+                _r := shl(7, lt(0xffffffffffffffffffffffffffffffff, x_))
+                _r := or(_r, shl(6, lt(0xffffffffffffffff, shr(_r, x_))))
+                _r := or(_r, shl(5, lt(0xffffffff, shr(_r, x_))))
+                _r := or(_r, shl(4, lt(0xffff, shr(_r, x_))))
+                _r := xor(31, or(shr(3, _r), lt(0xff, shr(_r, x_))))
+            }
+            function min(x_, y_) -> _z {
+                _z := xor(x_, mul(xor(x_, y_), lt(y_, x_)))
+            }
             if mload(data) {
                 result := mload(0x40)
-                let o := add(result, 0x20)
                 let s := add(data, 4)
                 let v := mload(s)
-                let end := add(data, mload(data))
+                let end := add(add(0x20, data), mload(data))
+                let m := not(shl(7, div(not(iszero(end)), 255))) // `0x7f7f ...`.
+                let o := add(result, 0x20)
                 mstore(s, not(v)) // Bitwise negate the first 4 bytes.
-                for {} lt(data, end) {} {
-                    data := add(data, 1)
-                    let c := byte(31, mload(data))
-                    if iszero(c) {
-                        data := add(data, 1)
-                        let d := byte(31, mload(data))
-                        // Fill with either 0xff or 0x00.
-                        mstore(o, not(0))
-                        if iszero(gt(d, 0x7f)) { calldatacopy(o, calldatasize(), add(d, 1)) }
-                        o := add(o, add(and(d, 0x7f), 1))
+                for { let i := add(0x20, data) } 1 {} {
+                    let c := mload(i)
+                    if iszero(byte(0, c)) {
+                        c := add(byte(1, c), 1)
+                        if iszero(gt(c, 0x80)) {
+                            calldatacopy(o, calldatasize(), c) // Fill with 0x00.
+                            o := add(o, c)
+                            i := add(i, 2)
+                            if iszero(lt(i, end)) { break }
+                            continue
+                        }
+                        mstore(o, not(0)) // Fill with 0xff.
+                        o := add(o, sub(c, 0x80))
+                        i := add(i, 2)
+                        if iszero(lt(i, end)) { break }
                         continue
                     }
-                    mstore8(o, c)
-                    o := add(o, 1)
+                    mstore(o, c)
+                    c := not(or(or(add(and(c, m), m), c), m))
+                    c := add(iszero(c), countLeadingZeroBytes(c))
+                    o := add(min(sub(end, i), c), o)
+                    i := add(c, i)
+                    if iszero(lt(i, end)) { break }
                 }
                 mstore(s, v) // Restore the first 4 bytes.
                 mstore(result, sub(o, add(result, 0x20))) // Store the length.

test/LibZip.t.sol

@@ -38,6 +38,29 @@ contract LibZipTest is SoladyTest {
     bytes internal constant _CD_COMPRESS_INPUT =
         hex"00000000000000000000000000000000000000000000000000000000000ae11c0000000000000000000000000000000000000000000000000000002b9cdca0ab0000000000000000000000000000000000003961790f8baa365051889e4c367d00000000000000000000000000000000000026d85539440bc844167ac0cc42320000000000000000000000000000000000000000000000007b55939986433925";
 
+    bytes internal constant _CD_COMPRESS_OUTPUT =
+        hex"ffe3f51e1c001a2b9cdca0ab00113961790f8baa365051889e4c367d001126d85539440bc844167ac0cc423200177b55939986433925";
+
+    function testABCCdCompressAndDecompressGas() public {
+        bytes memory data = abi.encode(_A, _B, _C);
+        assertEq(LibZip.cdDecompress(LibZip.cdCompress(data)).length, data.length);
+    }
+
+    function testABCCdCompressAndDecompressOriginalGas() public {
+        bytes memory data = abi.encode(_A, _B, _C);
+        assertEq(_cdDecompressOriginal(_cdCompressOriginal(data)).length, data.length);
+    }
+
+    function testCdDecompressGas() public {
+        bytes memory data = _CD_COMPRESS_OUTPUT;
+        assertGt(LibZip.cdDecompress(data).length, data.length);
+    }
+
+    function testCdDecompressOriginalGas() public {
+        bytes memory data = _CD_COMPRESS_OUTPUT;
+        assertGt(_cdDecompressOriginal(data).length, data.length);
+    }
+
     function testCdCompressGas() public {
         bytes memory data = _CD_COMPRESS_INPUT;
         assertLt(LibZip.cdCompress(data).length, data.length);
@@ -48,15 +71,15 @@ contract LibZipTest is SoladyTest {
         assertLt(_cdCompressOriginal(data).length, data.length);
     }
 
-    function testStoreABCWithCdCompressGas() public {
+    function testABCStoreWithCdCompressGas() public {
         _bytesStorage.set(LibZip.cdCompress(abi.encode(_A, _B, _C)));
     }
 
-    function testStoreABCWithCdCompressOriginalGas() public {
+    function testABCStoreWithCdCompressOriginalGas() public {
         _bytesStorage.set(_cdCompressOriginal(abi.encode(_A, _B, _C)));
     }
 
-    function testStoreABCWithFlzCompressGas() public {
+    function testABCStoreWithFlzCompressGas() public {
         _bytesStorage.set(LibZip.flzCompress(abi.encode(_A, _B, _C)));
     }
 
@@ -73,11 +96,36 @@ contract LibZipTest is SoladyTest {
     }
 
     function testCdCompressDifferential(bytes32) public {
-        testCdCompressDifferential(_randomCd());
+        bytes memory data;
+        if (_randomChance(8)) data = _randomCd();
+        uint256 t = _randomUniform() % 4;
+        for (uint256 i; i < t; ++i) {
+            if (_randomChance(2)) data = abi.encodePacked(data, _random());
+            if (_randomChance(2)) data = abi.encodePacked(data, new bytes(_random() & 0x3ff));
+            if (_randomChance(2)) data = abi.encodePacked(data, _random());
+            if (_randomChance(32)) data = abi.encodePacked(data, _randomCd());
+        }
+        testCdCompressDifferential(data);
     }
 
     function testCdCompressDifferential(bytes memory data) public {
-        assertEq(LibZip.cdCompress(data), _cdCompressOriginal(data));
+        if (_randomChance(32)) _misalignFreeMemoryPointer();
+        if (_randomChance(32)) _brutalizeMemory();
+        bytes memory computed = LibZip.cdCompress(data);
+        assertEq(computed, _cdCompressOriginal(data));
+    }
+
+    function testCdDecompressDifferential(bytes32) public {
+        bytes memory data = _randomCd();
+        if (_randomChance(2)) {
+            testCdDecompressDifferential(LibZip.cdCompress(data));
+        } else {
+            testCdDecompressDifferential(data);
+        }
+    }
+
+    function testCdDecompressDifferential(bytes memory data) public {
+        assertEq(LibZip.cdDecompress(data), _cdDecompressOriginal(data));
     }
 
     function _cdCompressOriginal(bytes memory data) internal pure returns (bytes memory result) {
@@ -121,6 +169,39 @@ contract LibZipTest is SoladyTest {
         }
     }
 
+    function _cdDecompressOriginal(bytes memory data) internal pure returns (bytes memory result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            if mload(data) {
+                result := mload(0x40)
+                let o := add(result, 0x20)
+                let s := add(data, 4)
+                let v := mload(s)
+                let end := add(data, mload(data))
+                mstore(s, not(v)) // Bitwise negate the first 4 bytes.
+                for {} lt(data, end) {} {
+                    data := add(data, 1)
+                    let c := byte(31, mload(data))
+                    if iszero(c) {
+                        data := add(data, 1)
+                        let d := byte(31, mload(data))
+                        // Fill with either 0xff or 0x00.
+                        mstore(o, not(0))
+                        if iszero(gt(d, 0x7f)) { calldatacopy(o, calldatasize(), add(d, 1)) }
+                        o := add(o, add(and(d, 0x7f), 1))
+                        continue
+                    }
+                    mstore8(o, c)
+                    o := add(o, 1)
+                }
+                mstore(s, v) // Restore the first 4 bytes.
+                mstore(result, sub(o, add(result, 0x20))) // Store the length.
+                mstore(o, 0) // Zeroize the slot after the string.
+                mstore(0x40, add(o, 0x20)) // Allocate the memory.
+            }
+        }
+    }
+
     function testFlzCompressDecompress() public brutalizeMemory {
         assertEq(LibZip.flzCompress(""), "");
         assertEq(LibZip.flzDecompress(""), "");
@@ -211,7 +292,18 @@ contract LibZipTest is SoladyTest {
     function _randomCd() internal returns (bytes memory data) {
         uint256 n = _randomChance(8) ? _random() % 2048 : _random() % 256;
         data = new bytes(n);
-        if (_randomChance(2)) {
+        if (_randomChance(32)) {
+            uint256 r = _randomUniform();
+            /// @solidity memory-safe-assembly
+            assembly {
+                mstore(0x00, r)
+                for { let i := 0 } lt(i, n) { i := add(i, 0x20) } {
+                    mstore(0x20, xor("randomUniform", i))
+                    mstore(add(add(data, 0x20), i), keccak256(0x00, 0x40))
+                }
+            }
+        }
+        if (_randomChance(4)) {
             /// @solidity memory-safe-assembly
             assembly {
                 for { let i := 0 } lt(i, n) { i := add(i, 0x20) } {
@@ -224,16 +316,45 @@ contract LibZipTest is SoladyTest {
             /// @solidity memory-safe-assembly
             assembly {
                 mstore(0x00, r)
+                mstore(0x20, xor("mode", not(0)))
+                let mode := and(1, keccak256(0x00, 0x40))
                 for { let i := 0 } lt(i, n) { i := add(i, 0x20) } {
-                    mstore(0x20, i)
-                    if and(1, keccak256(0x00, 0x40)) { mstore(add(add(data, 0x20), i), 0) }
+                    mstore(0x20, xor("mode", i))
+                    mode := xor(mode, iszero(and(keccak256(0x00, 0x40), 7)))
+                    mstore(add(add(data, 0x20), i), mul(iszero(mode), not(0)))
                 }
             }
         }
-        if (n != 0) {
-            uint256 m = _random() % 8;
-            for (uint256 j; j < m; ++j) {
-                data[_random() % n] = bytes1(uint8(_random()));
+        if (_randomChance(16)) {
+            uint256 r = _randomUniform();
+            /// @solidity memory-safe-assembly
+            assembly {
+                mstore(0x00, r)
+                for { let i := 0 } lt(i, n) { i := add(i, 0x20) } {
+                    mstore(0x20, xor("0", i))
+                    let p := keccak256(0x00, 0x40)
+                    if and(0x01, p) { mstore(add(add(data, 0x20), i), 0) }
+                }
+            }
+        }
+        if (_randomChance(16)) {
+            uint256 r = _randomUniform();
+            /// @solidity memory-safe-assembly
+            assembly {
+                mstore(0x00, r)
+                for { let i := 0 } lt(i, n) { i := add(i, 0x20) } {
+                    mstore(0x20, xor("not(0)", i))
+                    let p := keccak256(0x00, 0x40)
+                    if and(0x10, p) { mstore(add(add(data, 0x20), i), not(0)) }
+                }
+            }
+        }
+        if (_randomChance(2)) {
+            if (n != 0) {
+                uint256 m = _random() % 8;
+                for (uint256 j; j < m; ++j) {
+                    data[_random() % n] = bytes1(uint8(_random()));
+                }
             }
         }
     }
@@ -393,4 +514,27 @@ contract LibZipTest is SoladyTest {
         }
         assertEq(a, b);
     }
+
+    function testCountLeadingNonZeroBytes(bytes32 s) public {
+        uint256 expected;
+        uint256 computed;
+        /// @solidity memory-safe-assembly
+        assembly {
+            let n := 0
+            for {} byte(n, s) { n := add(n, 1) } {} // Scan for '\0'.
+            expected := n
+            let m := 0x7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F
+            let x := not(or(or(add(and(s, m), m), s), m))
+            computed := 0x20
+            if x {
+                let r := shl(7, lt(0xffffffffffffffffffffffffffffffff, x))
+                r := or(r, shl(6, lt(0xffffffffffffffff, shr(r, x))))
+                r := or(r, shl(5, lt(0xffffffff, shr(r, x))))
+                r := or(r, shl(4, lt(0xffff, shr(r, x))))
+                r := xor(31, or(shr(3, r), lt(0xff, shr(r, x))))
+                computed := r
+            }
+        }
+        assertEq(computed, expected);
+    }
 }