add nft chance booster hook

src/prize-hooks/examples/nft-chance-booster/NftChanceBoosterHook.sol

@@ -0,0 +1,154 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import { IPrizeHooks } from "pt-v5-vault/interfaces/IPrizeHooks.sol";
+import { IERC721 } from "openzeppelin-v5/token/ERC721/IERC721.sol";
+import { PrizePool, TwabController } from "pt-v5-prize-pool/PrizePool.sol";
+import { UniformRandomNumber } from "uniform-random-number/UniformRandomNumber.sol";
+
+uint256 constant PICK_GAS_ESTIMATE = 60_000; // probably lower, but we set it higher to avoid a reversion
+
+/// @notice Thrown if the boosted vault address is the zero address.
+error BoostedVaultAddressZero();
+
+/// @notice Thrown if the prize pool address is the zero address.
+error PrizePoolAddressZero();
+
+/// @notice Thrown if the nft collection address is the zero address.
+error NftCollectionAddressZero();
+
+/// @notice Thrown if the invalid token ID bounds are provided.
+/// @param tokenIdLowerBound The provided lower bound
+/// @param tokenIdUpperBound The provided upper bound
+error InvalidTokenIdBounds(uint256 tokenIdLowerBound, uint256 tokenIdUpperBound);
+
+/// @title PoolTogether V5 - Hook that boosts prizes on a vault for holders of an NFT collection
+/// @notice This prize hook awards the winner's prize to a random holder of a NFT in a specific collection.
+/// The winning NFT holder is limited between a specified ID range. The NFT holder must also have some amount
+/// of the specified prize vault tokens to be eligible to win. If a winner is picked who does not meet the
+/// requirements, the prize will be contributed on behalf of the specified prize vault instead.
+/// @dev This contract works best with NFTs that have iterating IDs ex. IDs: (1,2,3,4,5,...)
+/// @author G9 Software Inc.
+contract NftChanceBoosterHook is IPrizeHooks {
+
+    /// @notice Emitted when a vault is boosted with a prize re-contribution.
+    /// @param prizePool The prize pool the vault was boosted on
+    /// @param boostedVault The boosted vault
+    /// @param prizeAmount The amount of prize tokens contributed
+    event BoostedVaultWithPrize(address indexed prizePool, address indexed boostedVault, uint256 prizeAmount);
+
+    /// @notice The ERC721 token whose holders will have a chance to win prizes
+    IERC721 public immutable nftCollection;
+
+    /// @notice The prize pool that is awarding prizes
+    PrizePool public immutable prizePool;
+
+    /// @notice The twab controller associated with the prize pool
+    TwabController public immutable twabController;
+
+    /// @notice The vault that is being boosted
+    address public immutable boostedVault;
+
+    /// @notice The minimum TWAB that the selected winner must have to win a prize
+    uint256 public immutable minTwabOverPrizePeriod;
+
+    /// @notice The lower bound of eligible NFT IDs (inclusive)
+    uint256 public immutable tokenIdLowerBound;
+
+    /// @notice The upper bound of eligible NFT IDs (inclusive)
+    uint256 public immutable tokenIdUpperBound;
+
+    /// @notice Constructor to deploy the hook contract
+    /// @param nftCollection_ The ERC721 token whose holders will have a chance to win prizes
+    /// @param prizePool_ The prize pool that is awarding prizes
+    /// @param boostedVault_ The The vault that is being boosted
+    /// @param minTwabOverPrizePeriod_ The minimum TWAB that the selected winner must have over the prize
+    /// period to win the prize; if set to zero, no balance is needed.
+    /// @param tokenIdLowerBound_ The lower bound of eligible NFT IDs (inclusive)
+    /// @param tokenIdUpperBound_ The upper bound of eligible NFT IDs (inclusive)
+    constructor(
+        IERC721 nftCollection_,
+        PrizePool prizePool_,
+        address boostedVault_,
+        uint256 minTwabOverPrizePeriod_,
+        uint256 tokenIdLowerBound_,
+        uint256 tokenIdUpperBound_
+    ) {
+        if (address(0) == address(nftCollection_)) revert NftCollectionAddressZero();
+        if (address(0) == address(prizePool_)) revert PrizePoolAddressZero();
+        if (address(0) == boostedVault_) revert BoostedVaultAddressZero();
+        if (tokenIdUpperBound_ < tokenIdLowerBound_) revert InvalidTokenIdBounds(tokenIdLowerBound_, tokenIdUpperBound_);
+
+        nftCollection = nftCollection_;
+        prizePool = prizePool_;
+        twabController = prizePool_.twabController();
+        boostedVault = boostedVault_;
+        minTwabOverPrizePeriod = minTwabOverPrizePeriod_;
+        tokenIdLowerBound = tokenIdLowerBound_;
+        tokenIdUpperBound = tokenIdUpperBound_;
+    }
+
+    /// @inheritdoc IPrizeHooks
+    /// @dev This prize hook uses the random number from the last awarded prize pool draw to randomly select
+    /// the receiver of the prize from a list of current NFT holders. The prize tier and prize index are also
+    /// used to provide variance in the entropy for each prize so there can be multiple winners per draw.
+    /// @dev Tries to select a winner until the call runs out of gas before reverting to the backup action of 
+    /// contributing the prize on behalf of the boosted vault.
+    function beforeClaimPrize(address, uint8 _tier, uint32 _prizeIndex, uint96, address) external view returns (address, bytes memory) {
+        uint256 _tierStartTime;
+        uint256 _tierEndTime;
+        uint256 _winningRandomNumber = prizePool.getWinningRandomNumber();
+        {
+            uint24 _tierEndDrawId = prizePool.getLastAwardedDrawId();
+            uint24 _tierStartDrawId = prizePool.computeRangeStartDrawIdInclusive(
+                _tierEndDrawId,
+                prizePool.getTierAccrualDurationInDraws(_tier)
+            );
+            _tierStartTime = prizePool.drawOpensAt(_tierStartDrawId);
+            _tierEndTime = prizePool.drawClosesAt(_tierEndDrawId);
+        }
+        uint256 _pickAttempt;
+        for (; gasleft() >= PICK_GAS_ESTIMATE; _pickAttempt++) {
+            address _ownerOfToken;
+            {
+                uint256 _randomTokenId;
+                uint256 _numTokens = 1 + tokenIdUpperBound - tokenIdLowerBound;
+                if (_numTokens == 1) {
+                    _randomTokenId = tokenIdLowerBound;
+                } else {
+                    _randomTokenId = tokenIdLowerBound + UniformRandomNumber.uniform(
+                        uint256(keccak256(abi.encode(_winningRandomNumber, _tier, _prizeIndex, _pickAttempt))),
+                        _numTokens
+                    );
+                }
+                try nftCollection.ownerOf(_randomTokenId) returns (address _ownerOfResult) {
+                    _ownerOfToken = _ownerOfResult;
+                } catch { }
+            }
+            if (_ownerOfToken != address(0)) {
+                uint256 _recipientTwab;
+                if (minTwabOverPrizePeriod > 0) {
+                    _recipientTwab = twabController.getTwabBetween(boostedVault, _ownerOfToken, _tierStartTime, _tierEndTime);
+                }
+                if (_recipientTwab >= minTwabOverPrizePeriod) {
+                    // The owner of the selected NFT wins the prize!
+                    return (_ownerOfToken, abi.encode(_pickAttempt));
+                }
+            }
+        }
+        // By default, if no NFT winner can be determined, the prize will be sent to the prize pool and
+        // contributed on behalf of the boosted prize vault.
+        return (address(prizePool), abi.encode(_pickAttempt));
+    }
+
+    /// @inheritdoc IPrizeHooks
+    /// @dev If the recipient is set to the prize pool, the prize will be contributed on behalf of the vault
+    /// that is being boosted. Otherwise, it will do nothing (the prize will have already been sent to the
+    /// randomly selected NFT winner).
+    function afterClaimPrize(address, uint8, uint32, uint256 _prizeAmount, address _prizeRecipient, bytes memory) external {
+        if (_prizeRecipient == address(prizePool) && _prizeAmount > 0) {
+            prizePool.contributePrizeTokens(boostedVault, _prizeAmount);
+            emit BoostedVaultWithPrize(address(prizePool), boostedVault, _prizeAmount);
+        }
+    }
+}

test/prize-hooks/examples/nft-chance-booster/NftChanceBoosterHook.t.sol

@@ -0,0 +1,132 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.19;
+
+import { Test } from "forge-std/Test.sol";
+import { console2 } from "forge-std/console2.sol";
+
+import { NftChanceBoosterHook, PrizePool } from "src/prize-hooks/examples/nft-chance-booster/NftChanceBoosterHook.sol";
+import { PrizeHooks, IPrizeHooks } from "pt-v5-vault/interfaces/IPrizeHooks.sol";
+import { PrizePool } from "pt-v5-prize-pool/PrizePool.sol";
+import { ERC721ConsecutiveMock } from "openzeppelin-v5/mocks/token/ERC721ConsecutiveMock.sol";
+
+contract NftChanceBoosterHookTest is Test {
+
+    uint256 fork;
+    uint256 forkBlock = 17582112;
+    uint256 forkTimestamp = 1721996771;
+    uint256 randomNumber = 282830497779024192640724388550852704286534307968011569641355386343626319848;
+    uint96 nftIdOffset = 5;
+    uint256 minTwab = 1e18;
+    address[] holders;
+    uint96[] holderNumTokens;
+    address[] diverseHolders;
+    uint96[] diverseHoldersNumTokens;
+
+    NftChanceBoosterHook public nftBooster;
+    ERC721ConsecutiveMock public nft;
+    PrizePool public prizePool = PrizePool(address(0x45b2010d8A4f08b53c9fa7544C51dFd9733732cb));
+
+    function setUp() public {
+        fork = vm.createFork('base', forkBlock);
+        vm.selectFork(fork);
+        vm.warp(forkTimestamp);
+
+        holders.push(makeAddr('bob'));
+        holders.push(makeAddr('alice'));
+        holderNumTokens.push(9);
+        holderNumTokens.push(1);
+
+        for (uint256 i = 0; i < 100; i++) {
+            diverseHolders.push(makeAddr(string(abi.encode(i))));
+            diverseHoldersNumTokens.push(1);
+        }
+
+        nft = new ERC721ConsecutiveMock("Test NFT", "TNFT", nftIdOffset, holders, holders, holderNumTokens);
+
+        nftBooster = new NftChanceBoosterHook(nft, prizePool, address(this), minTwab, nftIdOffset, nftIdOffset + 9);
+    }
+
+    function testRecipientIsPrizePoolIfNotEligible() public {
+        address bob = holders[0];
+        address alice = holders[1];
+        assertEq(prizePool.twabController().delegateBalanceOf(address(this), bob), 0);
+        assertEq(prizePool.twabController().delegateBalanceOf(address(this), alice), 0);
+        // check a bunch of numbers and ensure there are no valid winners selected
+        for (uint256 i = 0; i < 1000; i++) {
+            vm.mockCall(address(prizePool), abi.encodeWithSelector(PrizePool.getWinningRandomNumber.selector), abi.encode(randomNumber + i));
+            address recipient = callBeforeClaimPrize(0, 0);
+            assertEq(recipient, address(prizePool));
+        }
+    }
+
+    function testRecipientIsNotPrizePoolIfEligible() public {
+        address bob = holders[0];
+        address alice = holders[1];
+        vm.warp(prizePool.drawOpensAt(1));
+        prizePool.twabController().mint(bob, 1e18);
+        prizePool.twabController().mint(alice, 1e18);
+        vm.warp(forkTimestamp);
+        assertEq(prizePool.twabController().delegateBalanceOf(address(this), bob), 1e18);
+        assertEq(prizePool.twabController().delegateBalanceOf(address(this), alice), 1e18);
+        // check a bunch of numbers and ensure there are no valid winners selected
+        uint256 numBobWins;
+        uint256 numAliceWins;
+        for (uint256 i = 0; i < 1000; i++) {
+            vm.mockCall(address(prizePool), abi.encodeWithSelector(PrizePool.getWinningRandomNumber.selector), abi.encode(randomNumber + i));
+            address recipient = callBeforeClaimPrize(0, 0);
+            if (recipient == bob) numBobWins++;
+            else if (recipient == alice) numAliceWins++;
+        }
+        assertGt(numBobWins, 800);
+        assertGt(numAliceWins, 50);
+    }
+
+    function testRecipientRetries() public {
+        address bob = holders[0];
+        address alice = holders[1];
+        vm.warp(prizePool.drawOpensAt(1));
+        prizePool.twabController().mint(bob, 1e18);
+        vm.warp(forkTimestamp);
+        assertEq(prizePool.twabController().delegateBalanceOf(address(this), bob), 1e18);
+        // Alice was not minted any balance, so she will not be eligible, but since bob has 5 tokens and alice has 1,
+        // it should still be very likely that bob is selected every time given 3 picks.
+        assertEq(prizePool.twabController().delegateBalanceOf(address(this), alice), 0);
+        // check a bunch of numbers and ensure there are no valid winners selected
+        uint256 numBobWins;
+        uint256 numAliceWins;
+        for (uint256 i = 0; i < 1000; i++) {
+            vm.mockCall(address(prizePool), abi.encodeWithSelector(PrizePool.getWinningRandomNumber.selector), abi.encode(randomNumber + i));
+            address recipient = callBeforeClaimPrize(0, 0);
+            if (recipient == bob) numBobWins++;
+            else if (recipient == alice) numAliceWins++;
+        }
+        assertGt(numBobWins, 900);
+        assertEq(numAliceWins, 0);
+    }
+
+    function testRetriesDoesNotRevert() public {
+        nft = new ERC721ConsecutiveMock("Test NFT", "TNFT", nftIdOffset, diverseHolders, diverseHolders, diverseHoldersNumTokens);
+        nftBooster = new NftChanceBoosterHook(nft, prizePool, address(this), minTwab, nftIdOffset, nftIdOffset + diverseHolders.length - 1);
+
+        // check a bunch of numbers and ensure there are no valid winners selected
+        for (uint256 i = 0; i < 1000; i++) {
+            vm.mockCall(address(prizePool), abi.encodeWithSelector(PrizePool.getWinningRandomNumber.selector), abi.encode(randomNumber + i));
+            address recipient = callBeforeClaimPrize(0, 0);
+            assertEq(recipient, address(prizePool));
+        }
+    }
+
+    function callBeforeClaimPrize(uint8 tier, uint32 prizeIndex) internal returns (address) {
+        (bool success, bytes memory data) = address(nftBooster).call{ gas: 150_000 }(abi.encodeWithSelector(IPrizeHooks.beforeClaimPrize.selector, address(0), tier, prizeIndex, 0, address(0)));
+        require(success, "beforeClaimPrize failed");
+        (address recipient, bytes memory hookData) = abi.decode(data, (address,bytes));
+        // if (hookData.length > 0) {
+        //     uint256 pickAttempt = abi.decode(hookData, (uint256));
+        //     if (pickAttempt > 0) {
+        //         console2.log("pick attempt", pickAttempt);
+        //     }
+        // }
+        return recipient;
+    }
+
+}