update forging test + airdrop update partners

contracts/Airdrop.sol

@@ -31,12 +31,13 @@ contract Airdrop is IAirdrop, AddressProviderResolver, ReentrancyGuard, Pausable
     address public liquidityPoolAddress;
     address public devAddress;
     address[] public partnerAddresses;
-    address[] public referralAddresses;
 
-    uint256 public liquidityPoolPercent = 10;
-    uint256 public devPercent = 15;
-    uint256 public referralPercent = 15;
-    uint256 public playersPercent = 55;
+    uint256 private constant BPS = 10_000; // denominator of basis points
+
+    uint256 public liquidityPoolRate = 1000;
+    uint256 public devRate = 1500;
+    uint256 public referralRate = 1500;
+    uint256 public playersRate = 5500;
 
     mapping(address => uint256) public userInfo;
     mapping(address => uint256) public referralInfo;
@@ -49,6 +50,7 @@ contract Airdrop is IAirdrop, AddressProviderResolver, ReentrancyGuard, Pausable
     error Airdrop__NotEligible();
     error Airdrop__AddressZero();
     error Airdrop__NothingToClaim();
+    error Airdrop__PartnerAddressesNotSet();
 
     // Functions
 
@@ -66,9 +68,9 @@ contract Airdrop is IAirdrop, AddressProviderResolver, ReentrancyGuard, Pausable
         if (amount == 0) revert Airdrop__InvalidAmount();
         if (address(_getTraitToken()) == address(0)) revert Airdrop__AddressZero();
         started = true;
-        totalTokenAmount = amount;
         _getTraitToken().transferFrom(msg.sender, address(this), amount);
-        distributeTokens(amount);
+        uint256 tokenTransfered = _distributeTokens(amount);
+        totalTokenAmount = amount - tokenTransfered;
     }
 
     function allowDaoFund() external onlyAirdropAccessor {
@@ -115,6 +117,24 @@ contract Airdrop is IAirdrop, AddressProviderResolver, ReentrancyGuard, Pausable
         totalReferredMints += mints;
     }
 
+    function setPartnerAddresses(address[] calldata partners) external onlyProtocolMaintainer {
+        if (partners.length == 0) revert Airdrop__PartnerAddressesNotSet();
+        for (uint256 i = 0; i < partners.length; i++) {
+            if (partners[i] == address(0)) revert Airdrop__AddressZero();
+        }
+        partnerAddresses = partners;
+    }
+
+    function setLiquidityPoolAddress(address addr) external onlyProtocolMaintainer {
+        if (addr == address(0)) revert Airdrop__AddressZero();
+        liquidityPoolAddress = addr;
+    }
+
+    function setDevAddress(address addr) external onlyProtocolMaintainer {
+        if (addr == address(0)) revert Airdrop__AddressZero();
+        devAddress = addr;
+    }
+
     function pause() public onlyProtocolMaintainer {
         _pause();
     }
@@ -135,18 +155,20 @@ contract Airdrop is IAirdrop, AddressProviderResolver, ReentrancyGuard, Pausable
     /**
      * internal & private *******************************************
      */
-    function distributeTokens(uint256 totalSupply) internal {
-        uint256 liquidityPoolShare = (totalSupply * liquidityPoolPercent) / 100;
-        uint256 devShare = (totalSupply * devPercent) / 100;
-        uint256 playersShare = (totalSupply * playersPercent) / 100;
-        uint256 referralShare = (totalSupply * referralPercent) / 100;
+    function _distributeTokens(uint256 totalSupply) internal returns (uint256 tokenTransfered) {
+        if (partnerAddresses.length == 0) revert Airdrop__PartnerAddressesNotSet();
+        uint256 liquidityPoolShare = (totalSupply * liquidityPoolRate) / BPS;
+        uint256 devShare = (totalSupply * devRate) / BPS;
+        uint256 playersShare = (totalSupply * playersRate) / BPS;
+        uint256 referralShare = (totalSupply * referralRate) / BPS;
         uint256 partnersShare =
             (totalSupply - liquidityPoolShare - devShare - playersShare - referralShare) / partnerAddresses.length;
 
         // Distribute to liquidity pool and dev address
         IERC20 trait = _getTraitToken();
         trait.transfer(liquidityPoolAddress, liquidityPoolShare);
         trait.transfer(devAddress, devShare);
+        tokenTransfered += liquidityPoolShare + devShare;
 
         // Store the amount allocated to players
         tokensToClaim = playersShare;
@@ -155,7 +177,10 @@ contract Airdrop is IAirdrop, AddressProviderResolver, ReentrancyGuard, Pausable
         // Distribute to partners
         for (uint256 i = 0; i < partnerAddresses.length; i++) {
             trait.transfer(partnerAddresses[i], partnersShare);
+            tokenTransfered += partnersShare;
         }
+
+        return tokenTransfered;
     }
 
     function _getTraitToken() private view returns (IERC20) {

contracts/EntityForging.sol

@@ -35,7 +35,6 @@ contract EntityForging is IEntityForging, AddressProviderResolver, ReentrancyGua
     error EntityForging__TokenAlreadyListed();
     error EntityForging__TokenNotOwnedByCaller();
     error EntityForging__FeeTooLow();
-    error EntityForging__ForgingLimitReached();
     error EntityForging__NotForger();
     error EntityForging__MergerTokenIdIsZero();
     error EntityForging__ForgerTokenIdNotListed();
@@ -44,7 +43,9 @@ contract EntityForging is IEntityForging, AddressProviderResolver, ReentrancyGua
     error EntityForging__TokensAlreadyForged();
     error EntityForging__FeeMismatchWithEthSent();
     error EntityForging__MergerEntropyCannotMerge();
-    error EntityForging__InsufficientForgePotential();
+    error EntityForging__InsufficientForgerForgePotential();
+    error EntityForging__InsufficientMergerForgePotential();
+    error EntityForging__TokenNotListedForForging(uint256 tokenId);
 
     // Functions
 
@@ -87,14 +88,15 @@ contract EntityForging is IEntityForging, AddressProviderResolver, ReentrancyGua
         _resetForgingCountIfNeeded(tokenId);
 
         uint256 entropy = traitForgeNft.getTokenEntropy(tokenId); // Retrieve entropy for tokenId
+
+        bool isForger = (entropy % 3) == 0; // Determine if the token is a forger need to keep the logic somewhere else
+        if (!isForger) revert EntityForging__NotForger();
+
         uint8 forgePotential = uint8((entropy / 10) % 10); // Extract the 5th digit from the entropy
         if (!(forgePotential > 0 && forgingCounts[tokenId] < forgePotential)) {
-            revert EntityForging__ForgingLimitReached();
+            revert EntityForging__InsufficientForgerForgePotential();
         }
 
-        bool isForger = (entropy % 3) == 0; // Determine if the token is a forger based on entropy
-        if (!isForger) revert EntityForging__NotForger(); // TODO need to mock the entropyGenerator
-
         ++listingCount;
         listings[listingCount] = Listing(msg.sender, tokenId, true, fee);
         listedTokenIds[tokenId] = listingCount;
@@ -136,7 +138,7 @@ contract EntityForging is IEntityForging, AddressProviderResolver, ReentrancyGua
         uint8 mergerForgePotential = uint8((mergerEntropy / 10) % 10); // Extract the 5th digit from the entropy
         forgingCounts[mergerId]++;
         if (!(mergerForgePotential > 0 && forgingCounts[mergerId] <= mergerForgePotential)) {
-            revert EntityForging__InsufficientForgePotential();
+            revert EntityForging__InsufficientMergerForgePotential();
         }
         /// TODO Stack too deep
         // uint256 devShare = (msg.value * taxCut) / BPS;
@@ -165,11 +167,10 @@ contract EntityForging is IEntityForging, AddressProviderResolver, ReentrancyGua
 
     function cancelListingForForging(uint256 tokenId) external whenNotPaused nonReentrant {
         ITraitForgeNft traitForgeNft = _getTraitForgeNft();
-        require(
-            traitForgeNft.ownerOf(tokenId) == msg.sender || msg.sender == address(traitForgeNft),
-            "Caller must own the token"
-        );
-        require(listings[listedTokenIds[tokenId]].isListed, "Token not listed for forging");
+        if (!(traitForgeNft.ownerOf(tokenId) == msg.sender || msg.sender == address(traitForgeNft))) {
+            revert EntityForging__TokenNotOwnedByCaller();
+        }
+        if (!listings[listedTokenIds[tokenId]].isListed) revert EntityForging__TokenNotListedForForging(tokenId);
 
         _cancelListingForForging(tokenId);
     }

contracts/EntropyGenerator.sol

@@ -20,7 +20,7 @@ contract EntropyGenerator is IEntropyGenerator, AddressProviderResolver, Pausabl
 
     constructor(address addressProvider) AddressProviderResolver(addressProvider) {
         _writeEntropyBatch(); // M05
-        initializeAlphaIndices();
+        _initializeAlphaIndices();
     }
 
     function nextEntropy() public onlyEntropyAccessor returns (uint256) {
@@ -45,13 +45,7 @@ contract EntropyGenerator is IEntropyGenerator, AddressProviderResolver, Pausabl
 
     // Select index points for 999999, triggered each generation increment
     function initializeAlphaIndices() public onlyEntropyAccessor {
-        uint256 hashValue = uint256(keccak256(abi.encodePacked(blockhash(block.number - 1), block.timestamp)));
-
-        uint256 slotIndexSelection = (hashValue % 258) + 612;
-        uint256 numberIndexSelection = hashValue % 12;
-
-        slotIndexSelectionPoint = slotIndexSelection;
-        numberIndexSelectionPoint = numberIndexSelection;
+        _initializeAlphaIndices();
     }
 
     // Public function to expose entropy calculation for a given slot and number index
@@ -118,4 +112,14 @@ contract EntropyGenerator is IEntropyGenerator, AddressProviderResolver, Pausabl
         }
         lastInitializedIndex = endIndex; // Update the index to indicate initialization is complete
     }
+
+    function _initializeAlphaIndices() private {
+        uint256 hashValue = uint256(keccak256(abi.encodePacked(blockhash(block.number - 1), block.timestamp)));
+
+        uint256 slotIndexSelection = (hashValue % 258) + 612;
+        uint256 numberIndexSelection = hashValue % 12;
+
+        slotIndexSelectionPoint = slotIndexSelection;
+        numberIndexSelectionPoint = numberIndexSelection;
+    }
 }

contracts/TraitForgeNft.sol

@@ -138,8 +138,6 @@ contract TraitForgeNft is ITraitForgeNft, AddressProviderResolver, ERC721Enumera
         onlyEntityForging
         returns (uint256)
     {
-        // TODO check if both parents are valid (already minted) meaning parentId1 and parentId2 are not 0 && both < _tokenIds
-        // if (parent1Id == parent2Id) revert TraitForgeNft__CannotForgeWithSameToken();
         uint256 newGeneration = getTokenGeneration(parent1Id) + 1;
 
         /// Check new generation is not over maxGeneration
@@ -166,14 +164,15 @@ contract TraitForgeNft is ITraitForgeNft, AddressProviderResolver, ERC721Enumera
         whenNotPaused
         nonReentrant
         onlyWhitelisted(proof, keccak256(abi.encodePacked(msg.sender)))
+        returns (uint256)
     {
         if (generationMintCounts[currentGeneration] == maxTokensPerGen) {
             _incrementGeneration();
         }
         uint256 mintPrice = calculateMintPrice();
         if (msg.value < mintPrice) revert TraitForgeNft__InsufficientETHSent();
 
-        _mintInternal(msg.sender, mintPrice);
+        uint256 tokenId = _mintInternal(msg.sender, mintPrice);
 
         uint256 excessPayment = msg.value - mintPrice;
         if (excessPayment > 0) {
@@ -182,6 +181,8 @@ contract TraitForgeNft is ITraitForgeNft, AddressProviderResolver, ERC721Enumera
         }
 
         _distributeFunds(mintPrice);
+
+        return tokenId;
     }
 
     function mintWithBudget(
@@ -293,7 +294,7 @@ contract TraitForgeNft is ITraitForgeNft, AddressProviderResolver, ERC721Enumera
     //////////////////////////////////////////////////////////////////////////////////////
 
     ///////////////////////////////// write functions /////////////////////////////////
-    function _mintInternal(address to, uint256 mintPrice) internal {
+    function _mintInternal(address to, uint256 mintPrice) internal returns (uint256) {
         _tokenIds++;
         uint256 newItemId = _tokenIds;
         _mint(to, newItemId);
@@ -317,6 +318,8 @@ contract TraitForgeNft is ITraitForgeNft, AddressProviderResolver, ERC721Enumera
         }
 
         emit Minted(msg.sender, newItemId, currentGeneration, entropyValue, mintPrice);
+
+        return newItemId;
     }
 
     /// TODO possibility to merge with _mintInternal same logic

contracts/interfaces/ITraitForgeNft.sol

@@ -52,7 +52,7 @@ interface ITraitForgeNft is IERC721Enumerable {
         external
         returns (uint256);
 
-    function mintToken(bytes32[] calldata proof) external payable;
+    function mintToken(bytes32[] calldata proof) external payable returns (uint256);
 
     function mintWithBudget(bytes32[] calldata proof, uint256 minAmountMinted) external payable;
 

test/fork/TraitForgeNft.t.sol

@@ -5,18 +5,17 @@ import { Fork_Test } from "./Fork.t.sol";
 import { TraitForgeNft } from "contracts/TraitForgeNft.sol";
 
 contract TraitForgeNftForkTest is Fork_Test {
-
     function setUp() public override {
         super.setUp();
     }
 
-    function test_traitForgeNftFork__mintToken() public {
-        TraitForgeNft nft = TraitForgeNft(traitForgeNft);
+    // function test_traitForgeNftFork__mintToken() public {
+    //     TraitForgeNft nft = TraitForgeNft(traitForgeNft);
 
-        vm.prank(address(0x225b791581185B73Eb52156942369843E8B0Eec7));
-        bytes32[] memory proofs = new bytes32[](2);
-        proofs[0] = bytes32(0xccde2ddaf58b1f916d5359378d87667005516a28fb95b92cffadebf464f8240b);
-        proofs[1] = bytes32(0xc8fbcb936bf99fdc0fe0524b9a208cf8c60d1fbc93de9d08ab1d204b91dd7a47);
-        nft.mintToken{value: 0.005 ether}(proofs);
-    }
-}
+    //     vm.prank(address(0x225b791581185B73Eb52156942369843E8B0Eec7));
+    //     bytes32[] memory proofs = new bytes32[](2);
+    //     proofs[0] = bytes32(0xccde2ddaf58b1f916d5359378d87667005516a28fb95b92cffadebf464f8240b);
+    //     proofs[1] = bytes32(0xc8fbcb936bf99fdc0fe0524b9a208cf8c60d1fbc93de9d08ab1d204b91dd7a47);
+    //     nft.mintToken{value: 0.005 ether}(proofs);
+    // }
+}

test/integration/concrete/airdrop/AirdropTest.t.sol

@@ -5,7 +5,7 @@ import { Deploys } from "test/shared/Deploys.sol";
 import { Roles } from "contracts/libraries/Roles.sol";
 
 contract AirdropTest is Deploys {
-    address internal _airdropAccessor = makeAddr("airdropAccessor");
+    address internal _airdropAccessor = makeAddr("_airdropAccessor");
     uint256 internal amount = 1_000_000 ether;
 
     function setUp() public virtual override {
@@ -17,9 +17,9 @@ contract AirdropTest is Deploys {
     function _startAirdrop() internal {
         vm.prank(_traitMinter);
         _trait.transfer(_airdropAccessor, amount);
-        vm.prank(_airdropAccessor);
+        vm.startPrank(_airdropAccessor);
         _trait.approve(address(_airdrop), amount);
-        vm.prank(_airdropAccessor);
         _airdrop.startAirdrop(amount);
+        vm.stopPrank();
     }
 }

test/integration/concrete/airdrop/Airdrop_AddUserAmount.t.sol

@@ -11,7 +11,6 @@ contract Airdrop_AddUserAmount is AirdropTest {
     address user = makeAddr("user");
 
     function testRevert__airdrop__addUserAmount__whenNotAuthorized() public {
-        _startAirdrop();
         vm.expectRevert(abi.encodeWithSelector(Errors.CallerNotAirdropAccessor.selector, _randomUser));
         vm.prank(_randomUser);
         _airdrop.addUserAmount(user, 100);

test/integration/concrete/airdrop/Airdrop_Claim.t.sol

@@ -6,7 +6,7 @@ import { Airdrop } from "contracts/Airdrop.sol";
 
 contract Airdrop_Claim is AirdropTest {
     address user = makeAddr("user");
-    uint256 amountToAdd = 10 ether;
+    uint256 amountToAdd = 1 ether;
 
     function testRevert_when_not_started() public {
         vm.expectRevert(Airdrop.Airdrop__NotStarted.selector);

test/integration/concrete/airdrop/Airdrop_StartAirdrop.t.sol

@@ -34,9 +34,6 @@ contract Airdrop_StartAirdrop is AirdropTest {
 
     function test__airdrop__startAirdrop() public {
         _startAirdrop();
-
-        assertEq(_airdrop.totalTokenAmount(), amount);
-        assertEq(_trait.balanceOf(address(_airdrop)), amount);
         assertEq(_airdrop.airdropStarted(), true);
     }
 }