refactor(core): optimize Balance storage library via Yul

.gas-snapshot

@@ -1,50 +1,50 @@
-AquaBalancesTest:testBalancesReturnsCorrectAmounts() (gas: 105025)
+AquaBalancesTest:testBalancesReturnsCorrectAmounts() (gas: 104917)
 AquaBalancesTest:testBalancesReturnsZeroForNonExistentStrategy() (gas: 14729)
-AquaBalancesTest:testBalancesReturnsZeroForTokenNotInStrategy() (gas: 50610)
-AquaBalancesTest:testSafeBalancesReturnsCorrectAmountsForActiveStrategy() (gas: 75494)
-AquaBalancesTest:testSafeBalancesRevertsAfterDock() (gas: 54700)
+AquaBalancesTest:testBalancesReturnsZeroForTokenNotInStrategy() (gas: 50580)
+AquaBalancesTest:testSafeBalancesReturnsCorrectAmountsForActiveStrategy() (gas: 75237)
+AquaBalancesTest:testSafeBalancesRevertsAfterDock() (gas: 54670)
 AquaBalancesTest:testSafeBalancesRevertsForNonExistentStrategy() (gas: 20643)
-AquaBalancesTest:testSafeBalancesRevertsIfFirstTokenNotInStrategy() (gas: 79639)
-AquaBalancesTest:testSafeBalancesRevertsIfSecondTokenNotInStrategy() (gas: 80304)
-AquaBalancesTest:testSafeBalancesTracksChangesFromPushPull() (gas: 138451)
-AquaEventsTest:testDockEmitsDockedEvent() (gas: 51825)
-AquaEventsTest:testPullEmitsPulledEvent() (gas: 86959)
-AquaEventsTest:testPushEmitsPushedEvent() (gas: 72320)
-AquaEventsTest:testShipEmitsPushedEventForEachToken() (gas: 81869)
-AquaEventsTest:testShipEmitsShippedEvent() (gas: 48377)
-AquaLifecycleTest:testFullLifecycle() (gas: 144801)
-AquaLifecycleTest:testMultipleStrategiesSameTokens() (gas: 169589)
-AquaPushPullTest:testPullAfterDockRevertsOnUnderflow() (gas: 52346)
-AquaPushPullTest:testPullFromNonExistentStrategy() (gas: 18345)
-AquaPushPullTest:testPullRevertsOnInsufficientBalance() (gas: 47845)
-AquaPushPullTest:testPushAccumulatesBalance() (gas: 70867)
-AquaPushPullTest:testPushFailsAfterDock() (gas: 54853)
-AquaPushPullTest:testPushOnlyForShippedTokens() (gas: 54189)
+AquaBalancesTest:testSafeBalancesRevertsIfFirstTokenNotInStrategy() (gas: 79570)
+AquaBalancesTest:testSafeBalancesRevertsIfSecondTokenNotInStrategy() (gas: 80047)
+AquaBalancesTest:testSafeBalancesTracksChangesFromPushPull() (gas: 138051)
+AquaEventsTest:testDockEmitsDockedEvent() (gas: 51795)
+AquaEventsTest:testPullEmitsPulledEvent() (gas: 86856)
+AquaEventsTest:testPushEmitsPushedEvent() (gas: 72220)
+AquaEventsTest:testShipEmitsPushedEventForEachToken() (gas: 81800)
+AquaEventsTest:testShipEmitsShippedEvent() (gas: 48347)
+AquaLifecycleTest:testFullLifecycle() (gas: 144589)
+AquaLifecycleTest:testMultipleStrategiesSameTokens() (gas: 169381)
+AquaPushPullTest:testPullAfterDockRevertsOnUnderflow() (gas: 52246)
+AquaPushPullTest:testPullFromNonExistentStrategy() (gas: 18275)
+AquaPushPullTest:testPullRevertsOnInsufficientBalance() (gas: 47745)
+AquaPushPullTest:testPushAccumulatesBalance() (gas: 70767)
+AquaPushPullTest:testPushFailsAfterDock() (gas: 54823)
+AquaPushPullTest:testPushOnlyForShippedTokens() (gas: 54159)
 AquaPushPullTest:testPushRequiresActiveStrategy() (gas: 20933)
-AquaPushPullTest:testPushRevertsOnBalanceOverflow() (gas: 51865)
-AquaPushPullTest:testPushRevertsOnUint248Overflow() (gas: 49656)
-AquaPushPullTest:testShipRevertsOnUint248AmountOverflow() (gas: 22935)
-AquaRouterTest:test_RescueFundsERC20() (gas: 1820965)
-AquaShipDockTest:testDockAlreadyDockedStrategyReverts() (gas: 53109)
+AquaPushPullTest:testPushRevertsOnBalanceOverflow() (gas: 51768)
+AquaPushPullTest:testPushRevertsOnUint248Overflow() (gas: 49611)
+AquaPushPullTest:testShipRevertsOnUint248AmountOverflow() (gas: 22869)
+AquaRouterTest:test_RescueFundsERC20() (gas: 1764392)
+AquaShipDockTest:testDockAlreadyDockedStrategyReverts() (gas: 53079)
 AquaShipDockTest:testDockNonExistentStrategyReverts() (gas: 19526)
-AquaShipDockTest:testDockRequiresAllTokensFromShip() (gas: 76759)
-AquaShipDockTest:testDockRequiresCorrectTokenCount() (gas: 79171)
-AquaShipDockTest:testDockRequiresExactTokensFromShip() (gas: 81174)
-AquaShipDockTest:testFuzzShipAmountDoesNotOverwriteTokenCount(uint248,uint248) (runs: 256, μ: 80910, ~: 80941)
-AquaShipDockTest:testShipCannotBeCalledTwiceForSameStrategy() (gas: 52652)
-AquaShipDockTest:testShipCannotHaveDuplicateTokens() (gas: 46705)
-AquaShipDockTest:testShipDockShipSameStrategyReverts() (gas: 93592)
+AquaShipDockTest:testDockRequiresAllTokensFromShip() (gas: 76690)
+AquaShipDockTest:testDockRequiresCorrectTokenCount() (gas: 79102)
+AquaShipDockTest:testDockRequiresExactTokensFromShip() (gas: 81105)
+AquaShipDockTest:testFuzzShipAmountDoesNotOverwriteTokenCount(uint248,uint248) (runs: 256, μ: 80838, ~: 80872)
+AquaShipDockTest:testShipCannotBeCalledTwiceForSameStrategy() (gas: 52632)
+AquaShipDockTest:testShipCannotHaveDuplicateTokens() (gas: 46676)
+AquaShipDockTest:testShipDockShipSameStrategyReverts() (gas: 93533)
 AquaShipDockTest:testShipRevertsWhenTokenCountEquals255() (gas: 128466)
 AquaShipDockTest:testShipRevertsWhenTokenCountExceeds255() (gas: 128460)
-AquaShipDockTest:testShipSameStrategyHashDifferentTokens() (gas: 111346)
-AquaShipDockTest:testShipSameStrategyHashPartiallyOverlappingTokensReverts() (gas: 83316)
-AquaShipDockTest:testShipWithLargeAmountsPreservesTokenCount() (gas: 102687)
-AquaShipDockTest:testShipWithZeroAmounts() (gas: 102277)
-AquaStorageTest:testDock1Token() (gas: 50947)
-AquaStorageTest:testDock2Tokens() (gas: 80108)
-AquaStorageTest:testDock3Tokens() (gas: 109385)
-AquaStorageTest:testPullSingleSloadSstore() (gas: 67696)
-AquaStorageTest:testPushSingleSloadSstore() (gas: 70616)
-AquaStorageTest:testShip1Token() (gas: 46337)
-AquaStorageTest:testShip2Tokens() (gas: 74317)
-AquaStorageTest:testShip3Tokens() (gas: 102276)
+AquaShipDockTest:testShipSameStrategyHashDifferentTokens() (gas: 111247)
+AquaShipDockTest:testShipSameStrategyHashPartiallyOverlappingTokensReverts() (gas: 83257)
+AquaShipDockTest:testShipWithLargeAmountsPreservesTokenCount() (gas: 102548)
+AquaShipDockTest:testShipWithZeroAmounts() (gas: 102138)
+AquaStorageTest:testDock1Token() (gas: 50917)
+AquaStorageTest:testDock2Tokens() (gas: 80039)
+AquaStorageTest:testDock3Tokens() (gas: 109277)
+AquaStorageTest:testPullSingleSloadSstore() (gas: 67593)
+AquaStorageTest:testPushSingleSloadSstore() (gas: 70516)
+AquaStorageTest:testShip1Token() (gas: 46307)
+AquaStorageTest:testShip2Tokens() (gas: 74248)
+AquaStorageTest:testShip3Tokens() (gas: 102168)

src/Aqua.sol

@@ -24,17 +24,23 @@ contract Aqua is IAqua {
                 mapping(address token => Balance)))) private _balances; // aka makers' allowances
 
     function rawBalances(address maker, address app, bytes32 strategyHash, address token) external view returns (uint248 balance, uint8 tokensCount) {
-        return _balances[maker][app][strategyHash][token].load();
+        uint256 rawData;
+        (rawData, tokensCount) = _balances[maker][app][strategyHash][token].load();
+
+        balance = uint248(rawData);
     }
 
     function safeBalances(address maker, address app, bytes32 strategyHash, address token0, address token1) external view returns (uint256 balance0, uint256 balance1) {
-        (uint248 amount0, uint8 tokensCount0) = _balances[maker][app][strategyHash][token0].load();
+        mapping(address => Balance) storage strategyBalances = _balances[maker][app][strategyHash];
+
+        (uint256 rawData0, uint8 tokensCount0) = strategyBalances[token0].load();
         require(tokensCount0 > 0 && tokensCount0 != _DOCKED, SafeBalancesForTokenNotInActiveStrategy(maker, app, strategyHash, token0));
-        balance0 = amount0;
 
-        (uint248 amount1, uint8 tokensCount1) = _balances[maker][app][strategyHash][token1].load();
+        (uint256 rawData1, uint8 tokensCount1) = strategyBalances[token1].load();
         require(tokensCount1 > 0 && tokensCount1 != _DOCKED, SafeBalancesForTokenNotInActiveStrategy(maker, app, strategyHash, token1));
-        balance1 = amount1;
+
+        balance0 = uint248(rawData0);
+        balance1 = uint248(rawData1);
     }
 
     function ship(address app, bytes calldata strategy, address[] calldata tokens, uint256[] calldata amounts) external returns(bytes32 strategyHash) {
@@ -43,10 +49,15 @@ contract Aqua is IAqua {
         require(tokensCount != _DOCKED, MaxNumberOfTokensExceeded(tokensCount, _DOCKED - 1));
 
         emit Shipped(msg.sender, app, strategyHash, strategy);
+
         for (uint256 i = 0; i < tokens.length; i++) {
             Balance storage balance = _balances[msg.sender][app][strategyHash][tokens[i]];
-            require(balance.tokensCount == 0, StrategiesMustBeImmutable(app, strategyHash));
-            balance.store(amounts[i].toUint248(), tokensCount);
+
+            (, uint8 currentTokensCount) = balance.load();
+            require(currentTokensCount == 0, StrategiesMustBeImmutable(app, strategyHash));
+
+            balance.set(amounts[i], tokensCount);
+
             emit Pushed(msg.sender, app, strategyHash, tokens[i], amounts[i]);
         }
     }
@@ -55,25 +66,29 @@ contract Aqua is IAqua {
         for (uint256 i = 0; i < tokens.length; i++) {
             Balance storage balance = _balances[msg.sender][app][strategyHash][tokens[i]];
             require(balance.tokensCount == tokens.length, DockingShouldCloseAllTokens(app, strategyHash));
-            balance.store(0, _DOCKED);
+            balance.set(0, _DOCKED);
         }
         emit Docked(msg.sender, app, strategyHash);
     }
 
     function pull(address maker, bytes32 strategyHash, address token, uint256 amount, address to) external {
         Balance storage balance = _balances[maker][msg.sender][strategyHash][token];
-        (uint248 prevBalance, uint8 tokensCount) = balance.load();
-        balance.store(prevBalance - amount.toUint248(), tokensCount);
+        (uint256 rawData,) = balance.load();
+
+        balance.decrease(rawData, amount);
 
         IERC20(token).safeTransferFrom(maker, to, amount);
         emit Pulled(maker, msg.sender, strategyHash, token, amount);
     }
 
     function push(address maker, address app, bytes32 strategyHash, address token, uint256 amount) external {
         Balance storage balance = _balances[maker][app][strategyHash][token];
-        (uint248 prevBalance, uint8 tokensCount) = balance.load();
-        require(tokensCount > 0 && tokensCount != _DOCKED, PushToNonActiveStrategyPrevented(maker, app, strategyHash, token));
-        balance.store(prevBalance + amount.toUint248(), tokensCount);
+        (uint256 rawData, uint8 tokensCount) = balance.load();
+        require(
+            tokensCount > 0 && tokensCount != _DOCKED, PushToNonActiveStrategyPrevented(maker, app, strategyHash, token)
+        );
+
+        balance.increase(rawData, amount);
 
         IERC20(token).safeTransferFrom(msg.sender, maker, amount);
         emit Pushed(maker, app, strategyHash, token, amount);

src/libs/Balance.sol

@@ -16,28 +16,74 @@ struct Balance {
 /// @title BalanceLib - Gas-optimized balance storage operations
 /// @notice Provides single-SLOAD/SSTORE operations for packed Balance struct
 library BalanceLib {
-    /// @notice Loads balance data from storage using exactly 1 SLOAD
-    /// @dev Assembly implementation ensures optimal gas usage
+    /// @notice Loads packed balance and token count data from storage using exactly 1 SLOAD
+    /// @dev Assembly implementation avoids bitmasking `amount` to save gas, returning the unmasked slot word
     /// @param balance The storage pointer to the Balance struct
-    /// @return amount The token balance amount
-    /// @return tokensCount The number of tokens in the strategy
-    function load(Balance storage balance) internal view returns (uint248 amount, uint8 tokensCount) {
+    /// @return rawSlot The full raw 256-bit storage word (lower 248 bits contain the balance amount)
+    /// @return tokensCount The number of tokens in the strategy (extracted from the highest 8 bits)
+    function load(Balance storage balance) internal view returns (uint256 rawSlot, uint8 tokensCount) {
         assembly ("memory-safe") {
-            let packed := sload(balance.slot)
-            amount := and(packed, 0x00ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff)
-            tokensCount := shr(248, packed)
+            rawSlot := sload(balance.slot)
+            tokensCount := shr(248, rawSlot)
         }
     }
 
-    /// @notice Stores balance data to storage using exactly 1 SSTORE
-    /// @dev Assembly implementation ensures optimal gas usage
-    /// @param balance The storage pointer to the Balance struct
-    /// @param amount The token balance amount to store
-    /// @param tokensCount The number of tokens in the strategy
-    function store(Balance storage balance, uint248 amount, uint8 tokensCount) internal {
+    /// @notice Stores a new balance and token count, ensuring safe 248-bit boundaries
+    /// @dev Replaces SafeCast library with native Yul panic for gas efficiency
+    function set(Balance storage balance, uint256 amount, uint8 tokensCount) internal {
         assembly ("memory-safe") {
+            let max248 := shr(8, not(0))
+
+            if gt(amount, max248) {
+                mstore(0x00, shl(224, 0x4e487b71)) // Panic(0x11)
+                mstore(0x04, 0x11)
+                revert(0x00, 0x24)
+            }
+
             let packed := or(amount, shl(248, tokensCount))
             sstore(balance.slot, packed)
         }
     }
+
+    /// @notice Safely increases the packed balance by a specified amount using in-place addition.
+    /// @dev Utilizes Yul assembly to bypass bitmasking and repacking overhead.
+    /// Reverts with standard Solidity arithmetic Panic (0x11) if the new balance exceeds the 248-bit limit.
+    /// @param balance The storage pointer to the packed Balance struct.
+    /// @param rawSlot The current unmasked 256-bit storage word (lower 248 bits: balance, upper 8 bits: tokensCount).
+    /// @param amount The token amount to add to the current balance.
+    function increase(Balance storage balance, uint256 rawSlot, uint256 amount) internal {
+        assembly ("memory-safe") {
+            let max248 := shr(8, not(0))
+            let prevBalance := and(rawSlot, max248)
+
+            if gt(amount, sub(max248, prevBalance)) {
+                mstore(0x00, shl(224, 0x4e487b71)) // Panic(0x11)
+                mstore(0x04, 0x11)
+                revert(0x00, 0x24)
+            }
+
+            sstore(balance.slot, add(rawSlot, amount))
+        }
+    }
+
+    /// @notice Safely decreases the packed balance by a specified amount using in-place subtraction.
+    /// @dev Utilizes Yul assembly to bypass bitmasking and repacking overhead.
+    /// Reverts with standard Solidity arithmetic Panic (0x11) if the amount exceeds the current balance.
+    /// @param balance The storage pointer to the packed Balance struct.
+    /// @param rawSlot The current unmasked 256-bit storage word (lower 248 bits: balance, upper 8 bits: tokensCount).
+    /// @param amount The token amount to subtract from the current balance.
+    function decrease(Balance storage balance, uint256 rawSlot, uint256 amount) internal {
+        assembly ("memory-safe") {
+            let max248 := shr(8, not(0))
+            let prevBalance := and(rawSlot, max248)
+
+            if gt(amount, prevBalance) {
+                mstore(0x00, shl(224, 0x4e487b71)) // Panic(0x11)
+                mstore(0x04, 0x11)
+                revert(0x00, 0x24)
+            }
+
+            sstore(balance.slot, sub(rawSlot, amount))
+        }
+    }
 }