Add bucket bankruptcy scenario for rewards manager (#930)

* Add bucket bankruptcy scenario for rewards manager

* Fix evm reverts

* PR feedback

* Update prepare test methods to add position in NFT if there is no position in it

tests/forge/invariants/PositionsAndRewards/BucketBankruptcyERC20PoolRewardsInvariants.t.sol

@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: UNLICENSED
+
+pragma solidity 0.8.18;
+
+import { Strings } from '@openzeppelin/contracts/utils/Strings.sol';
+
+import { Maths }             from 'src/libraries/internal/Maths.sol';
+import { Pool }              from 'src/base/Pool.sol';
+import { ERC20Pool }         from 'src/ERC20Pool.sol';
+import { ERC721Pool }        from 'src/ERC721Pool.sol';
+import { ERC20PoolFactory }  from 'src/ERC20PoolFactory.sol';
+import { ERC721PoolFactory } from 'src/ERC721PoolFactory.sol';
+import { PositionManager }   from 'src/PositionManager.sol';
+import { RewardsManager }    from 'src/RewardsManager.sol';
+
+import { TokenWithNDecimals } from '../../utils/Tokens.sol';
+
+import { BucketBankruptcyERC20PoolRewardsHandler } from './handlers/BucketBankruptcyERC20PoolRewardsHandler.sol';
+import { RewardsInvariants }                   from './RewardsInvariants.t.sol';
+
+contract BucketBankruptcyERC20PoolRewardsInvariants is RewardsInvariants {
+
+    BucketBankruptcyERC20PoolRewardsHandler internal _bucketBankruptcyerc20poolrewardsHandler;
+
+    function setUp() public override virtual {
+
+        super.setUp();
+
+        _erc20poolFactory  = new ERC20PoolFactory(address(_ajna));
+        _erc20impl         = _erc20poolFactory.implementation();
+        _erc721poolFactory = new ERC721PoolFactory(address(_ajna));
+        _erc721impl        = _erc721poolFactory.implementation();
+        _positionManager   = new PositionManager(_erc20poolFactory, _erc721poolFactory);
+        _rewardsManager    = new RewardsManager(address(_ajna), _positionManager);
+
+        uint256 noOfPools = vm.envOr("NO_OF_POOLS", uint256(10));
+
+        for (uint256 i = 0; i < noOfPools; ++i) {
+            address collateral = address(new TokenWithNDecimals(string(abi.encodePacked("Collateral", Strings.toString(i + 1))), "C", uint8(vm.envOr("COLLATERAL_PRECISION", uint256(18)))));
+            address quote      = address(new TokenWithNDecimals(string(abi.encodePacked("Quote", Strings.toString(i + 1))), "Q", uint8(vm.envOr("QUOTE_PRECISION", uint256(18)))));
+            address pool       = address(_erc20poolFactory.deployPool(collateral, quote, 0.05 * 10**18));
+
+            excludeContract(collateral);
+            excludeContract(quote);
+            excludeContract(pool);
+            _pools.push(pool);
+        }
+
+        // fund the rewards manager with 100M ajna
+        _ajna.mint(address(_rewardsManager), 100_000_000 * 1e18);
+
+        excludeContract(address(_ajna));
+        excludeContract(address(_quote));
+        excludeContract(address(_erc20poolFactory));
+        excludeContract(address(_erc721poolFactory));
+        excludeContract(address(_poolInfo));
+        excludeContract(address(_erc20impl));
+        excludeContract(address(_erc721impl));
+        excludeContract(address(_positionManager));
+        excludeContract(address(_rewardsManager));
+
+        _bucketBankruptcyerc20poolrewardsHandler = new BucketBankruptcyERC20PoolRewardsHandler(
+            address(_rewardsManager),
+            address(_positionManager),
+            _pools,
+            address(_ajna),
+            address(_poolInfo),
+            NUM_ACTORS,
+            address(this)
+        );
+
+        _handler = address(_bucketBankruptcyerc20poolrewardsHandler);
+    }
+}

tests/forge/invariants/PositionsAndRewards/handlers/BasePositionPoolHandler.sol

@@ -0,0 +1,155 @@
+// SPDX-License-Identifier: UNLICENSED
+
+pragma solidity 0.8.18;
+
+import '@openzeppelin/contracts/utils/structs/EnumerableSet.sol';
+
+import { Maths } from 'src/libraries/internal/Maths.sol';
+
+import { UnboundedPositionPoolHandler } from './unbounded/UnboundedPositionPoolHandler.sol';
+
+abstract contract BasePositionPoolHandler is UnboundedPositionPoolHandler {
+
+    using EnumerableSet for EnumerableSet.UintSet;
+
+    /********************************/
+    /***  Prepare Tests Functions ***/
+    /********************************/
+
+    function _preMemorializePositions(
+        uint256 noOfBuckets_,
+        uint256 amountToAdd_
+    ) internal returns (uint256 tokenId_, uint256[] memory indexes_) {
+        noOfBuckets_ = constrictToRange(noOfBuckets_, 1, buckets.length());
+        indexes_     = getRandomElements(noOfBuckets_, buckets.values());
+        uint256[] memory lpBalances = new uint256[](noOfBuckets_);
+
+        for (uint256 i = 0; i < noOfBuckets_; i++) {
+
+            uint256 bucketIndex = indexes_[i];
+
+            // ensure actor has a position
+            (uint256 lpBalanceBefore,) = _pool.lenderInfo(bucketIndex, _actor);
+
+            // add quote token if they don't have a position
+            if (lpBalanceBefore == 0) {
+                // bound amount
+                uint256 boundedAmount = constrictToRange(amountToAdd_, Maths.max(_pool.quoteTokenScale(), MIN_QUOTE_AMOUNT), MAX_QUOTE_AMOUNT);
+                _ensureQuoteAmount(_actor, boundedAmount);
+                try _pool.addQuoteToken(boundedAmount, bucketIndex, block.timestamp + 1 minutes, false) {
+                } catch (bytes memory err) {
+                    _ensurePoolError(err);
+                }
+            }
+
+            (lpBalances[i], ) = _pool.lenderInfo(bucketIndex, _actor);
+        }
+
+        _pool.increaseLPAllowance(address(_positionManager), indexes_, lpBalances);
+
+        // mint position NFT
+        tokenId_ = _mint();
+    }
+
+    function _preRedeemPositions(
+        uint256 noOfBuckets_,
+        uint256 amountToAdd_
+    ) internal returns (uint256 tokenId_, uint256[] memory indexes_) {
+
+        (tokenId_, indexes_) = _getNFTPosition(noOfBuckets_, amountToAdd_);
+
+        // approve positionManager to transfer LP tokens
+        address[] memory transferors = new address[](1);
+        transferors[0] = address(_positionManager);
+
+        _pool.approveLPTransferors(transferors);
+    }
+
+    function _preBurn(
+        uint256 bucketIndex_,
+        uint256 amountToAdd_
+    ) internal returns (uint256 tokenId_) { 
+        uint256[] memory indexes;
+
+        // check and create the position
+        (tokenId_, indexes) = _preRedeemPositions(bucketIndex_, amountToAdd_);
+
+        _redeemPositions(tokenId_, indexes);
+    }
+
+    function _preMoveLiquidity(
+        uint256 amountToMove_,
+        uint256 toIndex_
+    ) internal returns (uint256 tokenId_, uint256 boundedFromIndex_, uint256 boundedToIndex_) {
+        boundedToIndex_   = constrictToRange(toIndex_,   LENDER_MIN_BUCKET_INDEX, LENDER_MAX_BUCKET_INDEX);
+
+        uint256[] memory indexes;
+        (tokenId_, indexes) = _getNFTPosition(1, amountToMove_);
+        boundedFromIndex_   = indexes.length != 0 ? indexes[0]: 0;
+    }
+
+    function _getNFTPosition(
+        uint256 noOfBuckets_,
+        uint256 amountToAdd_
+    ) internal returns (uint256 tokenId_, uint256[] memory indexes_) {
+
+        // Check for exisiting nft positions in PositionManager
+        uint256[] memory tokenIds = getTokenIdsByActor(address(_actor));
+
+        if (tokenIds.length != 0 ) {
+            // use existing position NFT
+            tokenId_ = tokenIds[constrictToRange(randomSeed(), 0, tokenIds.length - 1)];
+            indexes_ = getBucketIndexesByTokenId(tokenId_);
+            if (indexes_.length != 0) {
+                noOfBuckets_ = constrictToRange(noOfBuckets_, 1, indexes_.length);
+
+                // pick random indexes from all positions
+                indexes_ = getRandomElements(noOfBuckets_, indexes_);
+            }
+        } else {
+            // create a position for the actor
+            (tokenId_, indexes_) = _preMemorializePositions(noOfBuckets_, amountToAdd_); 
+            _memorializePositions(tokenId_, indexes_);
+        }
+    }
+
+    function _getPosition(
+        uint256 bucketIndex_,
+        uint256 amountToAdd_
+    ) internal returns (uint256[] memory indexes_) {
+        // ensure actor has a position
+        (uint256 lpBalanceBefore,) = _pool.lenderInfo(bucketIndex_, _actor);
+
+        // add quote token if they don't have a position
+        if (lpBalanceBefore == 0) {
+            // bound amount
+            uint256 boundedAmount = constrictToRange(amountToAdd_, Maths.max(_pool.quoteTokenScale(), MIN_QUOTE_AMOUNT), MAX_QUOTE_AMOUNT);
+            _ensureQuoteAmount(_actor, boundedAmount);
+            try _pool.addQuoteToken(boundedAmount, bucketIndex_, block.timestamp + 1 minutes, false) {
+            } catch (bytes memory err) {
+                _ensurePoolError(err);
+            }
+        }
+
+        indexes_ = new uint256[](1);
+        indexes_[0] = bucketIndex_;
+
+        uint256[] memory lpBalances = new uint256[](1);
+
+        (lpBalances[0], ) = _pool.lenderInfo(bucketIndex_, _actor);
+        _pool.increaseLPAllowance(address(_positionManager), indexes_, lpBalances);
+    }
+
+    function getRandomElements(uint256 noOfElements, uint256[] memory arr) internal returns (uint256[] memory randomBuckets_) {
+        randomBuckets_ = new uint256[](noOfElements);
+
+        for (uint256 i = 0; i < noOfElements; i++) {
+            uint256 bucketIndex = constrictToRange(randomSeed(), 0, arr.length - 1 - i);
+            uint256 bucket = arr[bucketIndex];
+            randomBuckets_[i] = bucket;
+
+            // put last element from array to choosen array and next time pick new random element from first to last second element.
+            arr[bucketIndex] = arr[arr.length - 1 - i];
+        }
+    }
+}

tests/forge/invariants/PositionsAndRewards/handlers/BaseRewardsPoolHandler.sol

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: UNLICENSED
+
+pragma solidity 0.8.18;
+
+import { Strings } from '@openzeppelin/contracts/utils/Strings.sol';
+
+import { BasePositionPoolHandler }     from './BasePositionPoolHandler.sol';
+import { UnboundedRewardsPoolHandler } from './unbounded/UnboundedRewardsPoolHandler.sol';
+
+abstract contract BaseRewardsPoolHandler is UnboundedRewardsPoolHandler, BasePositionPoolHandler {
+
+    /*******************************/
+    /*** Prepare Tests Functions ***/
+    /*******************************/
+
+    function _preStake(
+        uint256 bucketIndex_,
+        uint256 amountToAdd_
+    ) internal returns (uint256 tokenId_, uint256[] memory indexes_) {
+
+        // retreive or create a NFT position
+        (tokenId_, indexes_) = _getNFTPosition(bucketIndex_, amountToAdd_);
+
+        // Approve rewards contract to transfer token
+        _positionManager.approve(address(_rewardsManager), tokenId_); 
+    }
+
+    function _preUnstake(
+        uint256 bucketIndex_,
+        uint256 amountToAdd_,
+        uint256 numberOfEpochs_
+    ) internal returns (uint256 tokenId_, uint256[] memory indexes_) {
+        (tokenId_, indexes_) = _getStakedPosition(bucketIndex_, amountToAdd_);
+
+        if (indexes_.length != 0) {
+            _advanceEpochRewardStakers(
+                amountToAdd_,
+                indexes_,
+                numberOfEpochs_
+            );
+        }
+    }
+
+    function _getStakedPosition(
+        uint256 bucketIndex_,
+        uint256 amountToAdd_
+    ) internal returns (uint256 tokenId_, uint256[] memory indexes_) {
+
+        // Check for exisiting staked positions in RewardsManager
+        uint256[] memory tokenIds = getStakedTokenIdsByActor(_actor);
+
+        if (tokenIds.length != 0 ) {
+            // use existing position NFT
+            tokenId_ = tokenIds[0];
+            indexes_ = getBucketIndexesByTokenId(tokenId_);
+            updateTokenAndPoolAddress(_positionManager.poolKey(tokenId_));
+
+            // create position in NFT if not already there
+            if (indexes_.length == 0) {
+                indexes_ = _getPosition(bucketIndex_, amountToAdd_);
+                _memorializePositions(tokenId_, indexes_);
+            }
+
+        } else {
+            // retreive or create a NFT position
+            (tokenId_, indexes_) = _getNFTPosition(bucketIndex_, amountToAdd_);
+            updateTokenAndPoolAddress(_positionManager.poolKey(tokenId_));
+
+            // approve rewards contract to transfer token
+            _positionManager.approve(address(_rewardsManager), tokenId_);
+
+            // stake the position
+            _stake(tokenId_);
+        }
+    }
+}