Twamm fixes

.forge-snapshots/TWAMM executTWAMMOrders 1 interval.snap

@@ -0,0 +1 @@
+484776

.forge-snapshots/TWAMM executTWAMMOrders 2 intervals.snap

@@ -0,0 +1 @@
+590266

.forge-snapshots/TWAMM executTWAMMOrders 3 intervals.snap

@@ -0,0 +1 @@
+687728

.forge-snapshots/TWAMM executTWAMMOrders singleSell 1 interval.snap

@@ -0,0 +1 @@
+256690

.forge-snapshots/TWAMM executTWAMMOrders singleSell 2 intervals.snap

@@ -0,0 +1 @@
+288800

.forge-snapshots/TWAMMSubmitOrder.snap

@@ -1 +1 @@
-122043
+122174

contracts/hooks/examples/TWAMM.sol

@@ -145,12 +145,16 @@ contract TWAMM is BaseHook, ITWAMM {
         PoolId poolId = key.toId();
         (uint160 sqrtPriceX96,,,) = manager.getSlot0(poolId);
         State storage twamm = twammStates[poolId];
+        if (twamm.lastVirtualOrderTimestamp == 0) revert NotInitialized();
 
         (bool zeroForOne, uint160 sqrtPriceLimitX96) =
             _executeTWAMMOrders(twamm, manager, key, PoolParamsOnExecute(sqrtPriceX96, manager.getLiquidity(poolId)));
 
         if (sqrtPriceLimitX96 != 0 && sqrtPriceLimitX96 != sqrtPriceX96) {
-            manager.unlock(abi.encode(key, IPoolManager.SwapParams(zeroForOne, type(int256).max, sqrtPriceLimitX96)));
+            // we trade to the sqrtPriceLimitX96, but v3 math inherently has small imprecision, must set swapAmountLimit
+            // to balance in case the trade needs more wei than is left in the contract
+            int256 swapAmountLimit = -int256(zeroForOne ? key.currency0.balanceOfSelf() : key.currency1.balanceOfSelf());
+            manager.unlock(abi.encode(key, IPoolManager.SwapParams(zeroForOne, swapAmountLimit, sqrtPriceLimitX96)));
         }
     }
 
@@ -256,13 +260,16 @@ contract TWAMM is BaseHook, ITWAMM {
         if (amountDelta != 0 && orderKey.expiration <= block.timestamp) revert CannotModifyCompletedOrder(orderKey);
 
         unchecked {
-            uint256 earningsFactor = orderPool.earningsFactorCurrent - order.earningsFactorLast;
-            buyTokensOwed = (earningsFactor * order.sellRate) >> FixedPoint96.RESOLUTION;
-            earningsFactorLast = orderPool.earningsFactorCurrent;
-            order.earningsFactorLast = earningsFactorLast;
+            earningsFactorLast = orderKey.expiration <= block.timestamp
+                ? orderPool.earningsFactorAtInterval[orderKey.expiration]
+                : orderPool.earningsFactorCurrent;
+            buyTokensOwed =
+                ((earningsFactorLast - order.earningsFactorLast) * order.sellRate) >> FixedPoint96.RESOLUTION;
 
             if (orderKey.expiration <= block.timestamp) {
                 delete self.orders[_orderId(orderKey)];
+            } else {
+                order.earningsFactorLast = earningsFactorLast;
             }
 
             if (amountDelta != 0) {

test/shared/implementation/TWAMMImplementation.sol

@@ -5,6 +5,8 @@ import {BaseHook} from "../../../contracts/BaseHook.sol";
 import {TWAMM} from "../../../contracts/hooks/examples/TWAMM.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {PoolId} from "@uniswap/v4-core/src/types/PoolId.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
 
 contract TWAMMImplementation is TWAMM {
     constructor(IPoolManager poolManager, uint256 interval, TWAMM addressToEtch) TWAMM(poolManager, interval) {
@@ -13,4 +15,22 @@ contract TWAMMImplementation is TWAMM {
 
     // make this a no-op in testing
     function validateHookAddress(BaseHook _this) internal pure override {}
+
+    function getOrderPoolEarningsFactorAtInterval(PoolId id, bool zeroForOne, uint256 timestamp)
+        external
+        view
+        returns (uint256 earningsFactor)
+    {
+        if (zeroForOne) return twammStates[id].orderPool0For1.earningsFactorAtInterval[timestamp];
+        else return twammStates[id].orderPool1For0.earningsFactorAtInterval[timestamp];
+    }
+
+    function isCrossingInitializedTick(
+        PoolParamsOnExecute memory pool,
+        IPoolManager poolManager,
+        PoolKey memory poolKey,
+        uint160 nextSqrtPriceX96
+    ) external view returns (bool crossingInitializedTick, int24 nextTickInit) {
+        return _isCrossingInitializedTick(pool, poolManager, poolKey, nextSqrtPriceX96);
+    }
 }