Added function to estimate prizes for both canary tiers

GenerationSoftware · Apr 1, 2024 · 8e34849 · 8e34849
1 parent 434e36c
commit 8e34849
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/src/abstract/TieredLiquidityDistributor.sol b/src/abstract/TieredLiquidityDistributor.sol
@@ -533,6 +533,19 @@ contract TieredLiquidityDistributor {
     return 0;
   }
 
+  /// @notice Estimates the prize count for the given tier, including BOTH canary tiers
+  /// @param numTiers The number of tiers
+  /// @return The estimated prize count across all tiers, including both canary tiers.
+  function estimatedPrizeCountWithBothCanaries(
+    uint8 numTiers
+  ) public view returns (uint32) {
+    if (numTiers >= MINIMUM_NUMBER_OF_TIERS && numTiers <= MAXIMUM_NUMBER_OF_TIERS) {
+      return estimatedPrizeCount(numTiers) + uint32(TierCalculationLib.prizeCount(numTiers - 1));
+    } else {
+      return 0;
+    }
+  }
+
   /// @notice Estimates the number of tiers for the given prize count.
   /// @dev Can return lower than the minimum, so that minimum can be detected
   /// @param _prizeCount The number of prizes that were claimed

diff --git a/test/abstract/TieredLiquidityDistributor.t.sol b/test/abstract/TieredLiquidityDistributor.t.sol
@@ -411,6 +411,15 @@ contract TieredLiquidityDistributorTest is Test {
     assertEq(distributor.estimatedPrizeCount(12), 0, "num tiers 12");
   }
 
+  function testEstimatedPrizeCountWithBothCanaries() public {
+    assertEq(distributor.estimatedPrizeCountWithBothCanaries(3), 0, "num tiers 3");
+    assertEq(distributor.estimatedPrizeCountWithBothCanaries(4), 20 + 4**3, "num tiers 4");
+    assertEq(distributor.estimatedPrizeCountWithBothCanaries(5), 80 + 4**4, "num tiers 5");
+    assertEq(distributor.estimatedPrizeCountWithBothCanaries(6), 320 + 4**5, "num tiers 6");
+    assertEq(distributor.estimatedPrizeCountWithBothCanaries(7), 1283 + 4**6, "num tiers 7");
+    assertEq(distributor.estimatedPrizeCountWithBothCanaries(12), 0, "num tiers 12");
+  }
+
   function testSumTierPrizeCounts() public {
     // 16 canary 1 daily + 64 canary 2 daily = 80
     assertEq(distributor.sumTierPrizeCounts(5), 80, "num tiers 5");