Fix JIT reward

contracts/docs/payload-types.md

@@ -171,7 +171,8 @@ enum RewardsUpdate {
     MultiTick {
         start_tick: i24,
         start_liquidity: u128,
-        quantities: List<u128>
+        quantities: List<u128>,
+        reward_checksum: u80
     },
     CurrentOnly {
         amount: u128
@@ -194,6 +195,7 @@ as the loop progresses to ensure consistency of reward distribution.
 |`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.|
+|`reward_checksum: u80`|The upper 80-bits of the "reward checksum", the hash chain of the liquidity and ticks traversed for reward|
 
 **Reward Update Internals**
 
@@ -223,13 +225,16 @@ def update_rewards(
     start_tick: Tick,
     quantities: list[int],
     liquidity: int,
+    expected_checksum_bits: 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: bytes = b'\x00' * 32
+
     for tick, quantity in zip(ticks, quantities):
         cumulative_reward_growth += quantity / liquidity
         tick.reward_growth_outside += cumulative_reward_growth
@@ -239,7 +244,15 @@ def update_rewards(
         else:
             liquidity -= tick.net_liquidity
 
+        reward_checksum = keccak_256(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 - 80)
+    assert checksum_bits == expected_checksum_bits, '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 RewardChecksumMismatch();
 
     // Stack too deep shenanigan.
     struct RewardParams {
         PoolId id;
         int24 tickSpacing;
         int24 currentTick;
+        uint256 rewardChecksum;
     }
 
     function _decodeAndReward(
@@ -58,7 +60,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);
@@ -72,6 +74,14 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
 
         newReader.requireAtEndOf(amountsEnd);
 
+        {
+            uint256 upperCheckSumBits;
+            (newReader, upperCheckSumBits) = newReader.readU80();
+            if (upperCheckSumBits != pool.rewardChecksum >> 176) {
+                revert RewardChecksumMismatch();
+            }
+        }
+
         uint128 currentLiquidity = UNI_V4.getPoolLiquidity(pool.id);
         if (endLiquidity != currentLiquidity) {
             revert WrongEndLiquidity(endLiquidity, currentLiquidity);
@@ -94,6 +104,7 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
         bool initialized = true;
         uint256 total = 0;
         uint256 cumulativeGrowth = 0;
+        uint256 rewardChecksum = 0;
 
         do {
             if (initialized) {
@@ -108,10 +119,18 @@ 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, add(32, add(16, 3)))
+                }
             }
             (initialized, rewardTick) = UNI_V4.getNextTickGt(pool.id, rewardTick, pool.tickSpacing);
         } while (rewardTick <= pool.currentTick);
 
+        pool.rewardChecksum = rewardChecksum;
         return (reader, total, cumulativeGrowth, liquidity);
     }
 
@@ -125,6 +144,7 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
         bool initialized = true;
         uint256 total = 0;
         uint256 cumulativeGrowth = 0;
+        uint256 rewardChecksum = 0;
 
         do {
             if (initialized) {
@@ -139,10 +159,18 @@ abstract contract GrowthOutsideUpdater is UniConsumer {
 
                 (, int128 netLiquidity) = UNI_V4.getTickLiquidity(pool.id, rewardTick);
                 liquidity = MixedSignLib.sub(liquidity, netLiquidity);
+
+                assembly ("memory-safe") {
+                    mstore(0x20, rewardTick)
+                    mstore(0x1d, liquidity)
+                    mstore(0x0d, rewardChecksum)
+                    rewardChecksum := keccak256(0x0d, add(32, add(16, 3)))
+                }
             }
             (initialized, rewardTick) = UNI_V4.getNextTickLt(pool.id, rewardTick, pool.tickSpacing);
         } while (rewardTick > pool.currentTick);
 
+        pool.rewardChecksum = rewardChecksum;
         return (reader, total, cumulativeGrowth, liquidity);
     }
 }

contracts/src/types/CalldataReader.sol

@@ -115,6 +115,14 @@ library CalldataReaderLib {
         return (self, value);
     }
 
+    function readU80(CalldataReader self) internal pure returns (CalldataReader, uint80 value) {
+        assembly ("memory-safe") {
+            value := shr(176, calldataload(self))
+            self := add(self, 10)
+        }
+        return (self, value);
+    }
+
     function readU128(CalldataReader self) internal pure returns (CalldataReader, uint128 value) {
         assembly ("memory-safe") {
             value := shr(128, calldataload(self))

contracts/test/_helpers/PoolGate.sol

@@ -169,8 +169,8 @@ contract PoolGate is IUnlockCallback, CommonBase, BaseTest {
         bytes32 delta1Slot = keccak256(abi.encode(sender, asset1));
         bytes32 rawDelta0 = UNI_V4.exttload(delta0Slot);
         bytes32 rawDelta1 = UNI_V4.exttload(delta1Slot);
-        delta = delta
-            + toBalanceDelta(int128(int256(uint256(rawDelta0))), int128(int256(uint256(rawDelta1))));
+        delta =
+            toBalanceDelta(int128(int256(uint256(rawDelta0))), int128(int256(uint256(rawDelta1))));
 
         require(delta.amount0() >= 0 && delta.amount1() >= 0, "losing money for removing liquidity");
         _clear(asset0, asset1, delta);

contracts/test/_helpers/RewardLib.sol

@@ -226,6 +226,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 = uint80(uint256(rewardChecksum) >> 176);
+        }
     }
 
     function _createRewardUpdateAbove(
@@ -294,6 +308,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 = uint80(uint256(rewardChecksum) >> 176);
+        }
     }
 
     function getLiquidityAtTick(IPoolManager uni, PoolId id, int24 futureTick, int24 tickSpacing)

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

@@ -268,12 +268,22 @@ contract RewardLibTest is BaseTest {
         pure
         returns (RewardsUpdate memory)
     {
+        return MultiTickReward(startTick, startLiquidity, quantities, 0);
+    }
+
+    function MultiTickReward(
+        int24 startTick,
+        uint128 startLiquidity,
+        uint128[] memory quantities,
+        uint72 rewardChecksum
+    ) internal pure returns (RewardsUpdate memory) {
         return RewardsUpdate({
             onlyCurrent: false,
             onlyCurrentQuantity: 0,
             startTick: startTick,
             startLiquidity: startLiquidity,
-            quantities: quantities
+            quantities: quantities,
+            rewardChecksum: rewardChecksum
         });
     }
 

contracts/test/_reference/PoolUpdate.sol

@@ -19,6 +19,7 @@ struct RewardsUpdate {
     int24 startTick;
     uint128 startLiquidity;
     uint128[] quantities;
+    uint80 rewardChecksum;
 }
 
 using PoolUpdateLib for PoolUpdate global;
@@ -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,
+            bytes10(self.rewardChecksum)
         );
     }