From b382a1e90c9403e613bb48126b4031372a048004 Mon Sep 17 00:00:00 2001
From: saucepoint <saucepoint@protonmail.com>
Date: Wed, 26 Jun 2024 12:33:37 -0400
Subject: [PATCH] save 600 gas, cleaner code when moving caller delta to
 tokensOwed

---
 .../autocompound_excessFeesCredit.snap        |  2 +-
 contracts/base/BaseLiquidityManagement.sol    | 14 ++---
 .../BalanceDeltaExtensionLibrary.sol          | 53 +++++++++++++++++++
 test/position-managers/Gas.t.sol              |  2 +-
 4 files changed, 62 insertions(+), 9 deletions(-)
 create mode 100644 contracts/libraries/BalanceDeltaExtensionLibrary.sol

diff --git a/.forge-snapshots/autocompound_excessFeesCredit.snap b/.forge-snapshots/autocompound_excessFeesCredit.snap
index e6c27d24..bcf9757d 100644
--- a/.forge-snapshots/autocompound_excessFeesCredit.snap
+++ b/.forge-snapshots/autocompound_excessFeesCredit.snap
@@ -1 +1 @@
-279673
\ No newline at end of file
+279016
\ No newline at end of file
diff --git a/contracts/base/BaseLiquidityManagement.sol b/contracts/base/BaseLiquidityManagement.sol
index dda04d7a..bc9ab1da 100644
--- a/contracts/base/BaseLiquidityManagement.sol
+++ b/contracts/base/BaseLiquidityManagement.sol
@@ -23,6 +23,7 @@ import {FeeMath} from "../libraries/FeeMath.sol";
 import {LiquiditySaltLibrary} from "../libraries/LiquiditySaltLibrary.sol";
 import {IBaseLiquidityManagement} from "../interfaces/IBaseLiquidityManagement.sol";
 import {PositionLibrary} from "../libraries/Position.sol";
+import {BalanceDeltaExtensionLibrary} from "../libraries/BalanceDeltaExtensionLibrary.sol";
 
 import "forge-std/console2.sol";
 
@@ -38,6 +39,7 @@ contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallback {
     using SafeCast for uint256;
     using LiquiditySaltLibrary for IHooks;
     using PositionLibrary for IBaseLiquidityManagement.Position;
+    using BalanceDeltaExtensionLibrary for BalanceDelta;
 
     mapping(address owner => mapping(LiquidityRangeId rangeId => Position)) public positions;
 
@@ -208,17 +210,15 @@ contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallback {
         BalanceDelta thisDelta
     ) private returns (BalanceDelta, BalanceDelta, BalanceDelta) {
         // credit the excess tokens to the position's tokensOwed
-        tokensOwed = useAmount0
-            ? toBalanceDelta(callerDelta.amount0(), tokensOwed.amount1())
-            : toBalanceDelta(tokensOwed.amount0(), callerDelta.amount1());
+        tokensOwed =
+            useAmount0 ? tokensOwed.setAmount0(callerDelta.amount0()) : tokensOwed.setAmount1(callerDelta.amount1());
 
         // this contract is responsible for custodying the excess tokens
-        thisDelta = useAmount0
-            ? thisDelta + toBalanceDelta(callerDelta.amount0(), 0)
-            : thisDelta + toBalanceDelta(0, callerDelta.amount1());
+        thisDelta =
+            useAmount0 ? thisDelta.addAmount0(callerDelta.amount0()) : thisDelta.addAmount1(callerDelta.amount1());
 
         // the caller is not expected to collect the excess tokens
-        callerDelta = useAmount0 ? toBalanceDelta(0, callerDelta.amount1()) : toBalanceDelta(callerDelta.amount0(), 0);
+        callerDelta = useAmount0 ? callerDelta.setAmount0(0) : callerDelta.setAmount1(0);
 
         return (tokensOwed, callerDelta, thisDelta);
     }
diff --git a/contracts/libraries/BalanceDeltaExtensionLibrary.sol b/contracts/libraries/BalanceDeltaExtensionLibrary.sol
new file mode 100644
index 00000000..e8b3a7f0
--- /dev/null
+++ b/contracts/libraries/BalanceDeltaExtensionLibrary.sol
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+
+library BalanceDeltaExtensionLibrary {
+    function setAmount0(BalanceDelta a, int128 amount0) internal pure returns (BalanceDelta) {
+        assembly {
+            // set the upper 128 bits of a to amount0
+            a := or(shl(128, amount0), and(sub(shl(128, 1), 1), a))
+        }
+        return a;
+    }
+
+    function setAmount1(BalanceDelta a, int128 amount1) internal pure returns (BalanceDelta) {
+        assembly {
+            // set the lower 128 bits of a to amount1
+            a := or(and(shl(128, sub(shl(128, 1), 1)), a), amount1)
+        }
+        return a;
+    }
+
+    function addAmount0(BalanceDelta a, int128 amount0) internal pure returns (BalanceDelta) {
+        assembly {
+            let a0 := sar(128, a)
+            let res0 := add(a0, amount0)
+            a := or(shl(128, res0), and(sub(shl(128, 1), 1), a))
+        }
+        return a;
+    }
+
+    function addAmount1(BalanceDelta a, int128 amount1) internal pure returns (BalanceDelta) {
+        assembly {
+            let a1 := signextend(15, a)
+            let res1 := add(a1, amount1)
+            a := or(and(shl(128, sub(shl(128, 1), 1)), a), res1)
+        }
+        return a;
+    }
+
+    function addAndAssign(BalanceDelta a, BalanceDelta b) internal pure returns (BalanceDelta) {
+        assembly {
+            let a0 := sar(128, a)
+            let a1 := signextend(15, a)
+            let b0 := sar(128, b)
+            let b1 := signextend(15, b)
+            let res0 := add(a0, b0)
+            let res1 := add(a1, b1)
+            a := or(shl(128, res0), and(sub(shl(128, 1), 1), res1))
+        }
+        return a;
+    }
+}
diff --git a/test/position-managers/Gas.t.sol b/test/position-managers/Gas.t.sol
index a6980005..fe2005e2 100644
--- a/test/position-managers/Gas.t.sol
+++ b/test/position-managers/Gas.t.sol
@@ -214,7 +214,7 @@ contract GasTest is Test, Deployers, GasSnapshot {
 
         // alice will use half of her fees to increase liquidity
         (uint256 token0Owed, uint256 token1Owed) = lpm.feesOwed(tokenIdAlice);
-        
+
         (uint160 sqrtPriceX96,,,) = StateLibrary.getSlot0(manager, range.poolKey.toId());
         uint256 liquidityDelta = LiquidityAmounts.getLiquidityForAmounts(
             sqrtPriceX96,