(Node Todo) Add checksum to check for JIT

contracts/docs/payload-types.md

@@ -171,10 +171,12 @@ enum RewardsUpdate {
     MultiTick {
         start_tick: i24,
         start_liquidity: u128,
-        quantities: List<u128>
+        quantities: List<u128>,
+        reward_checksum: u160
     },
     CurrentOnly {
-        amount: u128
+        amount: u128,
+        expected_liquidity: u128
     }
 }
 ```
@@ -191,9 +193,8 @@ as the loop progresses to ensure consistency of reward distribution.
 |Field|Description |
 |-----|-----------|
 |`start_tick: i24`| When `below = true` the current tick: the first tick **above** the first tick to donate to. <br> When rewarding above: just the first tick actually being donated to. |
-|`start_liquidity: u128`|The current liquidity if `start_tick` were the current tick.|
-|`quantities: List<u128>`|The reward for each initialized tick range *including* the current tick in
-`asset0` base units.|
+|`start_liquidity: u128`|The current liquidity if the first tick to donate to were the current tick.|
+|`quantities: List<u128>`|The reward for each initialized tick range *including* the current tick in `asset0` base units.|
 
 **Reward Update Internals**
 
@@ -223,12 +224,14 @@ def update_rewards(
     start_tick: Tick,
     quantities: list[int],
     liquidity: int,
+    expected_checksum: int,
     below: bool
 ):
     cumulative_reward_growth: float = 0
 
     end_tick: Tick = get_current_tick()
     ticks: list[Tick] = initialized_tick_range(start_tick, end_tick, include_end=below)
+    reward_checksum: bytes32 = b'\x00' * 32
 
     for tick, quantity in zip(ticks, quantities):
         cumulative_reward_growth += quantity / liquidity
@@ -239,7 +242,15 @@ def update_rewards(
         else:
             liquidity -= tick.net_liquidity
 
+        reward_checksum = keccak256(abi_encode_packed(
+            (reward_checksum, 'bytes32'),
+            (liquidity, 'uint128'),
+            (tick, 'int24')
+        ))
+
     assert len(quantities) == len(ticks) + 1, 'Unused quantities'
+    checksum_bits = int.from_bytes(reward_checksum, 'big') >> (256 - 160)
+    assert checksum_bits == expected_checksum, 'Invalid checksum'
 
     current_tick_reward: int = quantities[len(ticks)]
     cumulative_reward_growth += current_tick_reward / liquidity

contracts/src/modules/GrowthOutsideUpdater.sol

@@ -18,12 +18,14 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
     using TickLib for uint256;
 
     error WrongEndLiquidity(uint128 endLiquidity, uint128 actualCurrentLiquidity);
+    error JustInTimeLiquidityChange();
 
     // Stack too deep shenanigan.
     struct RewardParams {
         PoolId id;
         int24 tickSpacing;
         int24 currentTick;
+        uint256 rewardChecksum;
     }
 
     function _decodeAndReward(
@@ -37,9 +39,12 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
         if (currentOnly) {
             uint128 amount;
             (reader, amount) = reader.readU128();
+            uint128 expectedLiquidity;
+            (reader, expectedLiquidity) = reader.readU128();
+            uint128 pooLiquidity = UNI_V4.getPoolLiquidity(id);
+            if (expectedLiquidity != pooLiquidity) revert JustInTimeLiquidityChange();
             unchecked {
-                poolRewards_.globalGrowth +=
-                    X128MathLib.flatDivX128(amount, UNI_V4.getPoolLiquidity(id));
+                poolRewards_.globalGrowth += X128MathLib.flatDivX128(amount, pooLiquidity);
             }
 
             return (reader, amount);
@@ -58,7 +63,7 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
         PoolRewards storage poolRewards = poolRewards_;
 
         uint256 total;
-        RewardParams memory pool = RewardParams(id, tickSpacing, currentTick);
+        RewardParams memory pool = RewardParams(id, tickSpacing, currentTick, 0);
         (newReader, total, cumulativeGrowth, endLiquidity) = startTick <= pool.currentTick
             ? _rewardBelow(poolRewards.rewardGrowthOutside, startTick, newReader, liquidity, pool)
             : _rewardAbove(poolRewards.rewardGrowthOutside, startTick, newReader, liquidity, pool);
@@ -77,6 +82,14 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
             revert WrongEndLiquidity(endLiquidity, currentLiquidity);
         }
 
+        {
+            uint160 expectedRewardChecksum;
+            (newReader, expectedRewardChecksum) = newReader.readU160();
+            if (expectedRewardChecksum != pool.rewardChecksum >> 96) {
+                revert JustInTimeLiquidityChange();
+            }
+        }
+
         unchecked {
             poolRewards.globalGrowth += cumulativeGrowth;
         }
@@ -94,6 +107,7 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
         bool initialized = true;
         uint256 total = 0;
         uint256 cumulativeGrowth = 0;
+        uint256 rewardChecksum = 0;
 
         do {
             if (initialized) {
@@ -108,10 +122,19 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
 
                 (, int128 netLiquidity) = UNI_V4.getTickLiquidity(pool.id, rewardTick);
                 liquidity = MixedSignLib.add(liquidity, netLiquidity);
+
+                assembly ("memory-safe") {
+                    mstore(0x13, rewardTick)
+                    mstore(0x10, liquidity)
+                    mstore(0x00, rewardChecksum)
+                    rewardChecksum := keccak256(0x00, 0x33)
+                }
             }
             (initialized, rewardTick) = UNI_V4.getNextTickGt(pool.id, rewardTick, pool.tickSpacing);
         } while (rewardTick <= pool.currentTick);
 
+        pool.rewardChecksum = rewardChecksum;
+
         return (reader, total, cumulativeGrowth, liquidity);
     }
 
@@ -125,6 +148,7 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
         bool initialized = true;
         uint256 total = 0;
         uint256 cumulativeGrowth = 0;
+        uint256 rewardChecksum = 0;
 
         do {
             if (initialized) {
@@ -139,10 +163,19 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
 
                 (, int128 netLiquidity) = UNI_V4.getTickLiquidity(pool.id, rewardTick);
                 liquidity = MixedSignLib.sub(liquidity, netLiquidity);
+
+                assembly ("memory-safe") {
+                    mstore(0x13, rewardTick)
+                    mstore(0x10, liquidity)
+                    mstore(0x00, rewardChecksum)
+                    rewardChecksum := keccak256(0x00, 0x33)
+                }
             }
             (initialized, rewardTick) = UNI_V4.getNextTickLt(pool.id, rewardTick, pool.tickSpacing);
         } while (rewardTick > pool.currentTick);
 
+        pool.rewardChecksum = rewardChecksum;
+
         return (reader, total, cumulativeGrowth, liquidity);
     }
 }

contracts/src/periphery/AngstromView.sol

@@ -14,6 +14,7 @@ library AngstromView {
     uint256 constant UNLOCKED_FEE_PACKED_SET_SLOT = 2;
 
     uint256 constant LAST_BLOCK_CONFIG_STORE_SLOT = 3;
+    uint256 constant LAST_BLOCK_BIT_OFFSET = 64;
     uint256 constant STORE_BIT_OFFSET = 64;
 
     uint256 constant BALANCES_SLOT = 5;
@@ -22,6 +23,10 @@ library AngstromView {
         return address(uint160(self.extsload(CONTROLLER_SLOT)));
     }
 
+    function lastBlockUpdated(IAngstromAuth self) internal view returns (uint64) {
+        return uint64(self.extsload(LAST_BLOCK_CONFIG_STORE_SLOT) >> LAST_BLOCK_BIT_OFFSET);
+    }
+
     function configStore(IAngstromAuth self) internal view returns (PoolConfigStore addr) {
         uint256 value = self.extsload(LAST_BLOCK_CONFIG_STORE_SLOT);
         return PoolConfigStore.wrap(address(uint160(value >> STORE_BIT_OFFSET)));

contracts/src/periphery/ControllerV1.sol

@@ -114,7 +114,9 @@ contract ControllerV1 is Ownable2Step {
             entry = configStore.getWithDefaultEmpty(key, poolIndex);
             if (!entry.isEmpty()) break;
             poolIndex++;
-            if (poolIndex >= totalEntries) revert NonexistentPool(asset0, asset1);
+            if (poolIndex >= totalEntries) {
+                revert NonexistentPool(asset0, asset1);
+            }
         }
 
         emit PoolRemoved(asset0, asset1, entry.tickSpacing(), entry.bundleFee());

contracts/src/types/CalldataReader.sol

@@ -131,6 +131,14 @@ library CalldataReaderLib {
         return (self, addr);
     }
 
+    function readU160(CalldataReader self) internal pure returns (CalldataReader, uint160 value) {
+        assembly ("memory-safe") {
+            value := shr(96, calldataload(self))
+            self := add(self, 20)
+        }
+        return (self, value);
+    }
+
     function readU256(CalldataReader self) internal pure returns (CalldataReader, uint256 value) {
         assembly ("memory-safe") {
             value := calldataload(self)

contracts/test/_helpers/RewardLib.sol

@@ -49,9 +49,14 @@ library RewardLib {
         return rewards.length;
     }
 
-    function CurrentOnly(uint128 amount) internal pure returns (RewardsUpdate memory update) {
+    function CurrentOnly(IPoolManager uni, PoolId id, uint128 amount)
+        internal
+        view
+        returns (RewardsUpdate memory update)
+    {
         update.onlyCurrent = true;
         update.onlyCurrentQuantity = amount;
+        update.startLiquidity = uni.getPoolLiquidity(id);
     }
 
     function toUpdates(TickReward[] memory rewards, IPoolManager uni, PoolId id, int24 tickSpacing)
@@ -192,7 +197,7 @@ library RewardLib {
 
         if (initializedTicks.length == 0) {
             require(rewards.length == 1, "expected rewards length 1");
-            return CurrentOnly(rewards[0].amount);
+            return CurrentOnly(uni, id, rewards[0].amount);
         }
 
         update.quantities = new uint128[](initializedTicks.length + 1);
@@ -226,6 +231,20 @@ library RewardLib {
         update.startTick = int24(uint24(initializedTicks.get(0)));
         uint128 poolLiq = getLiquidityAtTick(uni, id, currentTick, tickSpacing);
         update.startLiquidity = MixedSignLib.sub(poolLiq, cumulativeNetLiquidity);
+
+        {
+            PoolId id2 = id;
+            IPoolManager uni2 = uni;
+            bytes32 rewardChecksum;
+            uint128 liquidity = update.startLiquidity;
+            for (uint256 i = 0; i < initializedTicks.length; i++) {
+                int24 tick = int24(int256(initializedTicks.get(i)));
+                (, int128 netLiquidity) = uni2.getTickLiquidity(id2, tick);
+                liquidity = MixedSignLib.add(liquidity, netLiquidity);
+                rewardChecksum = keccak256(abi.encodePacked(rewardChecksum, liquidity, tick));
+            }
+            update.rewardChecksum = uint160(uint256(rewardChecksum) >> 96);
+        }
     }
 
     function _createRewardUpdateAbove(
@@ -260,7 +279,7 @@ library RewardLib {
         if (initializedTicks.length == 0) {
             console.log("WARNING\nWARNING: Above somehow called with donate to current only???");
             require(rewards.length == 1, "Expected exact one reward");
-            return CurrentOnly(rewards[0].amount);
+            return CurrentOnly(uni, id, rewards[0].amount);
         }
 
         update.startTick = startTick;
@@ -294,6 +313,20 @@ library RewardLib {
 
         uint128 poolLiq = getLiquidityAtTick(uni, id, currentTick, tickSpacing);
         update.startLiquidity = MixedSignLib.add(poolLiq, cumulativeNetLiquidity);
+
+        {
+            PoolId id2 = id;
+            IPoolManager uni2 = uni;
+            bytes32 rewardChecksum;
+            uint128 liquidity = update.startLiquidity;
+            for (uint256 i = 0; i < initializedTicks.length; i++) {
+                int24 tick = int24(int256(initializedTicks.get(i)));
+                (, int128 netLiquidity) = uni2.getTickLiquidity(id2, tick);
+                liquidity = MixedSignLib.sub(liquidity, netLiquidity);
+                rewardChecksum = keccak256(abi.encodePacked(rewardChecksum, liquidity, tick));
+            }
+            update.rewardChecksum = uint160(uint256(rewardChecksum) >> 96);
+        }
     }
 
     function getLiquidityAtTick(IPoolManager uni, PoolId id, int24 futureTick, int24 tickSpacing)

contracts/test/_helpers/tests/RewardLib.t.sol

@@ -26,7 +26,9 @@ contract FakeUni is PoolManager {
 
     function setCurrentTick(int24 tick) public {
         _pools[id].slot0 = _pools[id].slot0.setTick(tick);
-        _pools[id].liquidity = tickLiq[tick.normalizeUnchecked(TICK_SPACING)];
+        int24 norm = tick.normalizeUnchecked(TICK_SPACING);
+        uint128 newLiq = tickLiq[norm];
+        _pools[id].liquidity = newLiq;
     }
 
     function liq() public view returns (uint128) {
@@ -49,6 +51,10 @@ contract FakeUni is PoolManager {
         }
     }
 
+    function liqNet(int24 tick) public view returns (int128) {
+        return _pools[id].ticks[tick].liquidityNet;
+    }
+
     function _initialize(int24 tick) internal {
         (int16 wordPos, uint8 bitPos) = TickLib.position(TickLib.compress(tick, TICK_SPACING));
         _pools[id].tickBitmap[wordPos] |= (1 << uint256(bitPos));
@@ -205,9 +211,10 @@ contract RewardLibTest is BaseTest {
         uint128 amount3
     ) public {
         uni.setCurrentTick(tick = -30);
+        TickReward[] memory rewards =
+            re(TickReward(120, amount1), TickReward(0, amount2), TickReward(-120, amount3));
         assertCreatesUpdates(
-            re(TickReward(120, amount1), TickReward(0, amount2), TickReward(-120, amount3)),
-            MultiTickReward(120, uni.tickLiq(120), [amount1, 0, amount2, amount3].into())
+            rewards, MultiTickReward(120, uni.tickLiq(120), [amount1, 0, amount2, amount3].into())
         );
     }
 
@@ -273,7 +280,8 @@ contract RewardLibTest is BaseTest {
             onlyCurrentQuantity: 0,
             startTick: startTick,
             startLiquidity: startLiquidity,
-            quantities: quantities
+            quantities: quantities,
+            rewardChecksum: 0
         });
     }
 

contracts/test/_reference/PoolUpdate.sol

@@ -19,6 +19,7 @@ struct RewardsUpdate {
     int24 startTick;
     uint128 startLiquidity;
     uint128[] quantities;
+    uint160 rewardChecksum;
 }
 
 using PoolUpdateLib for PoolUpdate global;
@@ -58,7 +59,7 @@ library PoolUpdateLib {
 
     function encode(RewardsUpdate memory self) internal pure returns (bytes memory) {
         if (self.onlyCurrent) {
-            return bytes.concat(bytes16(self.onlyCurrentQuantity));
+            return bytes.concat(bytes16(self.onlyCurrentQuantity), bytes16(self.startLiquidity));
         }
         bytes memory encodedQuantities;
         for (uint256 i = 0; i < self.quantities.length; i++) {
@@ -68,7 +69,10 @@ library PoolUpdateLib {
             bytes.concat(bytes3(encodedQuantities.length.toUint24()), encodedQuantities);
 
         return bytes.concat(
-            bytes3(uint24(self.startTick)), bytes16(self.startLiquidity), encodedQuantities
+            bytes3(uint24(self.startTick)),
+            bytes16(self.startLiquidity),
+            encodedQuantities,
+            bytes20(self.rewardChecksum)
         );
     }
 

contracts/test/benchmark/FullBundle.b.sol

@@ -15,6 +15,7 @@ import {UserOrder, UserOrderLib} from "../_reference/UserOrder.sol";
 import {TickMath} from "v4-core/src/libraries/TickMath.sol";
 import {PriceAB} from "src/types/Price.sol";
 import {RouterActor} from "test/_mocks/RouterActor.sol";
+import {RewardLib} from "../_helpers/RewardLib.sol";
 
 import {console} from "forge-std/console.sol";
 
@@ -107,7 +108,7 @@ contract FullBundleBenchmark is BaseTest {
             update.assetIn = asset1;
             update.assetOut = asset0;
             update.amountIn = 100.0e18;
-            update.rewardUpdate.onlyCurrent = true;
+            update.rewardUpdate = RewardLib.CurrentOnly(uni, poolId(angstrom, asset0, asset1), 0);
         }
 
         bundle.userOrders = new UserOrder[](total);
@@ -228,7 +229,7 @@ contract FullBundleBenchmark is BaseTest {
             update.assetIn = asset1;
             update.assetOut = asset0;
             update.amountIn = 100.0e18;
-            update.rewardUpdate.onlyCurrent = true;
+            update.rewardUpdate = RewardLib.CurrentOnly(uni, poolId(angstrom, asset0, asset1), 0);
         }
 
         bundle.userOrders = new UserOrder[](total);