Add extra data to one time distribution (#9)

* Add extra data to one time distribution

* Fix forge snapshots

* Replace nightly with stable

* Replace nightly with stable 2

* Remove forge gas snapshot

* Update forge std

* Update gas snapshots

* Fix snapshot

.forge-snapshots/EthAaveLeverageStrategyTest_test_rescueLendingAssets2.snap

@@ -1 +1 @@
-500064
+500051

.forge-snapshots/MerkleDistributorTest_test_constructor.snap

@@ -1 +1 @@
-1399588
+1410652

.forge-snapshots/MerkleDistributorTest_test_distributeOneTime.snap

@@ -1 +1 @@
-72196
+74933

.forge-snapshots/MerkleDistributorTest_test_setRewardsRoot.snap

@@ -1 +1 @@
-52114
+52092

.gas-snapshot

@@ -0,0 +1,52 @@
+EthAaveLeverageStrategyTest:test_claimExitedAssets() (gas: 4140025)
+EthAaveLeverageStrategyTest:test_claimExitedAssets_InvalidExitQueueTicket() (gas: 1110826)
+EthAaveLeverageStrategyTest:test_claimExitedAssets_NoExitPosition() (gas: 43521)
+EthAaveLeverageStrategyTest:test_deposit_ExitingProxy() (gas: 1125441)
+EthAaveLeverageStrategyTest:test_deposit_HasPosition() (gas: 2044855)
+EthAaveLeverageStrategyTest:test_deposit_NoPosition() (gas: 889041)
+EthAaveLeverageStrategyTest:test_deposit_WithoutApproval() (gas: 344057)
+EthAaveLeverageStrategyTest:test_deposit_ZeroShares() (gas: 34281)
+EthAaveLeverageStrategyTest:test_enterExitQueue() (gas: 1359245)
+EthAaveLeverageStrategyTest:test_enterExitQueue_ExitingProxy() (gas: 1121545)
+EthAaveLeverageStrategyTest:test_enterExitQueue_InvalidPositionPercent() (gas: 908584)
+EthAaveLeverageStrategyTest:test_enterExitQueue_NoPosition() (gas: 66966)
+EthAaveLeverageStrategyTest:test_forceEnterExitQueue_ForceExitConfigChecksNotPassed() (gas: 1146334)
+EthAaveLeverageStrategyTest:test_forceEnterExitQueue_NoForceExitConfig() (gas: 901003)
+EthAaveLeverageStrategyTest:test_forceEnterExitQueue_borrowForceExitLtvPercent() (gas: 1209757)
+EthAaveLeverageStrategyTest:test_forceEnterExitQueue_vaultForceExitLtvPercent() (gas: 1183792)
+EthAaveLeverageStrategyTest:test_permit() (gas: 1093572)
+EthAaveLeverageStrategyTest:test_receiveFlashLoan_InvalidCaller() (gas: 31251)
+EthAaveLeverageStrategyTest:test_rescueLendingAssets() (gas: 1901155)
+EthAaveLeverageStrategyTest:test_rescueLendingAssets_InvalidPosition() (gas: 61474)
+EthAaveLeverageStrategyTest:test_rescueLendingAssets_InvalidSlippage() (gas: 33988)
+EthAaveLeverageStrategyTest:test_rescueLendingAssets_NoBalancerPoolIdConfig() (gas: 1039251)
+EthAaveLeverageStrategyTest:test_rescueLendingAssets_ZeroAssets() (gas: 60434)
+EthAaveLeverageStrategyTest:test_rescueVaultAssets() (gas: 2740245)
+EthAaveLeverageStrategyTest:test_rescueVaultAssets_ExitPositionNotProcessed() (gas: 1131761)
+EthAaveLeverageStrategyTest:test_rescueVaultAssets_ExitQueueNotEntered() (gas: 43073)
+EthAaveLeverageStrategyTest:test_rescueVaultAssets_InvalidExitQueueTicket() (gas: 1111494)
+EthAaveLeverageStrategyTest:test_rescueVaultAssets_NoRescueVaultConfig() (gas: 1621794)
+EthAaveLeverageStrategyTest:test_upgradeProxy() (gas: 999571)
+EthAaveLeverageStrategyTest:test_upgradeProxy_NoVaultUpgradeConfig() (gas: 900157)
+EthAaveLeverageStrategyTest:test_upgradeProxy_NotRegisteredProxy() (gas: 45298)
+EthAaveLeverageStrategyTest:test_upgradeProxy_VaultUpgradeConfigSameAddress() (gas: 980507)
+EthAaveLeverageStrategyTest:test_upgradeProxy_VaultUpgradeConfigZeroAddress() (gas: 959850)
+EthAaveLeverageStrategyTest:test_upgradeProxy_WithExitingPosition() (gas: 1101835)
+MerkleDistributorTest:test_claim() (gas: 731451)
+MerkleDistributorTest:test_constructor() (gas: 1420037)
+MerkleDistributorTest:test_distributeOneTime() (gas: 334784)
+MerkleDistributorTest:test_distributePeriodically() (gas: 358162)
+MerkleDistributorTest:test_setRewardsDelay() (gas: 71336)
+MerkleDistributorTest:test_setRewardsMinOracles() (gas: 142651)
+MerkleDistributorTest:test_setRewardsRoot() (gas: 401435)
+StakeHelpersTest:test_calculateStake() (gas: 338892)
+StakeHelpersTest:test_calculateUnstake() (gas: 524501)
+StrategiesRegistryTest:test_addStrategyProxy() (gas: 240040)
+StrategiesRegistryTest:test_initialize() (gas: 62604)
+StrategiesRegistryTest:test_setStrategy() (gas: 168598)
+StrategiesRegistryTest:test_setStrategyConfig() (gas: 116083)
+VaultUserLtvTrackerTest:test_multipleStakeMint() (gas: 724889)
+VaultUserLtvTrackerTest:test_multipleVaults() (gas: 1035108)
+VaultUserLtvTrackerTest:test_normalUserNormalUser() (gas: 839774)
+VaultUserLtvTrackerTest:test_normalUserZeroUser() (gas: 595217)
+VaultUserLtvTrackerTest:test_zeroLtv() (gas: 36451)

.github/workflows/coverage.yml

@@ -21,7 +21,7 @@ jobs:
       - name: Install foundry
         uses: foundry-rs/foundry-toolchain@v1
         with:
-          version: nightly
+          version: stable
       - name: Run Forge build
         run: |
           forge --version

.github/workflows/lint.yml

@@ -15,7 +15,7 @@ jobs:
       - name: Install Foundry
         uses: foundry-rs/foundry-toolchain@v1
         with:
-          version: nightly
+          version: stable
 
       - name: Lint
         run: forge fmt --check

.github/workflows/tests.yml

@@ -16,13 +16,13 @@ jobs:
       - name: Install Foundry
         uses: foundry-rs/foundry-toolchain@v1
         with:
-          version: nightly
+          version: stable
 
       - name: Build
         run: forge build
 
       - name: Run tests
-        run: forge test --isolate -vvv
+        run: forge snapshot --isolate -vvv
         env:
           FORGE_SNAPSHOT_CHECK: true
           MAINNET_RPC_URL: ${{ secrets.MAINNET_RPC_URL }}

.gitmodules

@@ -7,9 +7,6 @@
 [submodule "lib/openzeppelin-contracts-upgradeable"]
 	path = lib/openzeppelin-contracts-upgradeable
 	url = https://github.com/OpenZeppelin/openzeppelin-contracts-upgradeable
-[submodule "lib/forge-gas-snapshot"]
-	path = lib/forge-gas-snapshot
-	url = https://github.com/marktoda/forge-gas-snapshot
 [submodule "lib/aave-v3-origin"]
 	path = lib/aave-v3-origin
 	url = https://github.com/aave-dao/aave-v3-origin

snapshots/EthAaveLeverageStrategyTest.json

@@ -0,0 +1,15 @@
+{
+  "EthAaveLeverageStrategyTest_test_claimExitedAssets1": "566164",
+  "EthAaveLeverageStrategyTest_test_claimExitedAssets2": "1008200",
+  "EthAaveLeverageStrategyTest_test_claimExitedAssets3": "1003269",
+  "EthAaveLeverageStrategyTest_test_deposit_HasPosition": "778215",
+  "EthAaveLeverageStrategyTest_test_deposit_NoPosition": "905013",
+  "EthAaveLeverageStrategyTest_test_enterExitQueue": "229194",
+  "EthAaveLeverageStrategyTest_test_forceEnterExitQueue_borrowForceExitLtvPercent": "275313",
+  "EthAaveLeverageStrategyTest_test_forceEnterExitQueue_vaultForceExitLtvPercent": "248645",
+  "EthAaveLeverageStrategyTest_test_permit": "370326",
+  "EthAaveLeverageStrategyTest_test_rescueLendingAssets1": "611767",
+  "EthAaveLeverageStrategyTest_test_rescueLendingAssets2": "503767",
+  "EthAaveLeverageStrategyTest_test_rescueVaultAssets": "650101",
+  "EthAaveLeverageStrategyTest_test_upgradeProxy": "55780"
+}

snapshots/MerkleDistributorTest.json

@@ -0,0 +1,9 @@
+{
+  "MerkleDistributorTest_test_claim": "146428",
+  "MerkleDistributorTest_test_constructor": "1410695",
+  "MerkleDistributorTest_test_distributeOneTime": "74966",
+  "MerkleDistributorTest_test_distributePeriodically": "73090",
+  "MerkleDistributorTest_test_setRewardsDelay": "30905",
+  "MerkleDistributorTest_test_setRewardsMinOracles": "36907",
+  "MerkleDistributorTest_test_setRewardsRoot": "52139"
+}

snapshots/StrategiesRegistryTest.json

@@ -0,0 +1,5 @@
+{
+  "StrategiesRegistryTest_test_addStrategyProxy": "71090",
+  "StrategiesRegistryTest_test_setStrategy": "48814",
+  "StrategiesRegistryTest_test_setStrategyConfig": "73969"
+}

src/MerkleDistributor.sol

@@ -135,11 +135,16 @@ contract MerkleDistributor is Ownable2Step, EIP712, IMerkleDistributor {
     }
 
     /// @inheritdoc IMerkleDistributor
-    function distributeOneTime(address token, uint256 amount, string calldata rewardsIpfsHash) external onlyOwner {
+    function distributeOneTime(
+        address token,
+        uint256 amount,
+        string calldata rewardsIpfsHash,
+        bytes calldata extraData
+    ) external onlyOwner {
         if (amount == 0) revert InvalidAmount();
 
         SafeERC20.safeTransferFrom(IERC20(token), msg.sender, address(this), amount);
-        emit OneTimeDistributionAdded(msg.sender, token, amount, rewardsIpfsHash);
+        emit OneTimeDistributionAdded(msg.sender, token, amount, rewardsIpfsHash, extraData);
     }
 
     /// @inheritdoc IMerkleDistributor

src/interfaces/IMerkleDistributor.sol

@@ -58,9 +58,10 @@ interface IMerkleDistributor {
      * @param token The address of the token
      * @param amount The amount of tokens to distribute
      * @param rewardsIpfsHash The IPFS hash of the rewards
+     * @param extraData The extra data for the distribution
      */
     event OneTimeDistributionAdded(
-        address indexed caller, address indexed token, uint256 amount, string rewardsIpfsHash
+        address indexed caller, address indexed token, uint256 amount, string rewardsIpfsHash, bytes extraData
     );
 
     /**
@@ -166,8 +167,14 @@ interface IMerkleDistributor {
      * @param token The address of the token
      * @param amount The amount of tokens to distribute
      * @param rewardsIpfsHash The IPFS hash of the rewards
+     * @param extraData The extra data for the distribution
      */
-    function distributeOneTime(address token, uint256 amount, string calldata rewardsIpfsHash) external;
+    function distributeOneTime(
+        address token,
+        uint256 amount,
+        string calldata rewardsIpfsHash,
+        bytes calldata extraData
+    ) external;
 
     /**
      * @notice Claim the tokens for a user

test/MerkleDistributor.t.sol

@@ -3,7 +3,6 @@
 pragma solidity ^0.8.26;
 
 import {Test} from 'forge-std/Test.sol';
-import {GasSnapshot} from 'forge-gas-snapshot/GasSnapshot.sol';
 import {Ownable} from '@openzeppelin/contracts/access/Ownable.sol';
 import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol';
 import {MessageHashUtils} from '@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol';
@@ -13,7 +12,7 @@ import {Errors} from '@stakewise-core/libraries/Errors.sol';
 import {MerkleDistributor} from '../src/MerkleDistributor.sol';
 import {IMerkleDistributor} from '../src/interfaces/IMerkleDistributor.sol';
 
-contract MerkleDistributorTest is Test, GasSnapshot {
+contract MerkleDistributorTest is Test {
     uint256 public constant forkBlockNumber = 21_264_254;
 
     MerkleDistributor public distributor;
@@ -32,9 +31,9 @@ contract MerkleDistributorTest is Test, GasSnapshot {
     }
 
     function test_constructor() public {
-        snapStart('MerkleDistributorTest_test_constructor');
+        vm.startSnapshotGas('MerkleDistributorTest_test_constructor');
         MerkleDistributor distributor2 = new MerkleDistributor(address(keeper), owner, 1 days, 2);
-        snapEnd();
+        vm.stopSnapshotGas();
         assertEq(distributor2.rewardsDelay(), 1 days, 'Should correctly set rewardsDelay');
         assertEq(distributor2.rewardsMinOracles(), 2, 'Should correctly set rewardsMinOracles');
         assertEq(distributor2.owner(), owner, 'Should correctly set owner');
@@ -60,9 +59,9 @@ contract MerkleDistributorTest is Test, GasSnapshot {
         vm.startPrank(owner);
         vm.expectEmit(true, false, false, true);
         emit IMerkleDistributor.RewardsMinOraclesUpdated(owner, 1);
-        snapStart('MerkleDistributorTest_test_setRewardsMinOracles');
+        vm.startSnapshotGas('MerkleDistributorTest_test_setRewardsMinOracles');
         distributor.setRewardsMinOracles(1);
-        snapEnd();
+        vm.stopSnapshotGas();
         vm.stopPrank();
 
         assertEq(distributor.rewardsMinOracles(), 1, 'Should correctly update rewardsMinOracles');
@@ -78,9 +77,9 @@ contract MerkleDistributorTest is Test, GasSnapshot {
         vm.startPrank(owner);
         vm.expectEmit(true, false, false, true);
         emit IMerkleDistributor.RewardsDelayUpdated(owner, newDelay);
-        snapStart('MerkleDistributorTest_test_setRewardsDelay');
+        vm.startSnapshotGas('MerkleDistributorTest_test_setRewardsDelay');
         distributor.setRewardsDelay(newDelay);
-        snapEnd();
+        vm.stopSnapshotGas();
         vm.stopPrank();
 
         assertEq(distributor.rewardsDelay(), newDelay, 'Should correctly update rewardsDelay');
@@ -102,11 +101,11 @@ contract MerkleDistributorTest is Test, GasSnapshot {
         vm.expectRevert(abi.encodeWithSelector(IMerkleDistributor.InvalidDuration.selector));
         distributor.distributePeriodically(address(swiseToken), amount, 3600, 0, '');
 
-        snapStart('MerkleDistributorTest_test_distributePeriodically');
+        vm.startSnapshotGas('MerkleDistributorTest_test_distributePeriodically');
         vm.expectEmit(true, true, false, true);
         emit IMerkleDistributor.PeriodicDistributionAdded(owner, address(swiseToken), amount, 3600, 86_400, '');
         distributor.distributePeriodically(address(swiseToken), amount, 3600, 86_400, '');
-        snapEnd();
+        vm.stopSnapshotGas();
         vm.stopPrank();
 
         assertEq(swiseToken.balanceOf(address(distributor)), amount, 'Tokens should be transferred to distributor');
@@ -117,7 +116,7 @@ contract MerkleDistributorTest is Test, GasSnapshot {
 
         // Ensure unauthorized accounts cannot call the function
         vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(this)));
-        distributor.distributeOneTime(address(swiseToken), amount, 'ipfsHash');
+        distributor.distributeOneTime(address(swiseToken), amount, 'ipfsHash', 'extraData');
 
         // Impersonate the owner to approve SWISE and distribute tokens
         vm.startPrank(owner);
@@ -126,18 +125,18 @@ contract MerkleDistributorTest is Test, GasSnapshot {
 
         // Test invalid amount (zero)
         vm.expectRevert(abi.encodeWithSelector(IMerkleDistributor.InvalidAmount.selector));
-        distributor.distributeOneTime(address(swiseToken), 0, 'ipfsHash');
+        distributor.distributeOneTime(address(swiseToken), 0, 'ipfsHash', 'extraData');
 
-        snapStart('MerkleDistributorTest_test_distributeOneTime');
+        vm.startSnapshotGas('MerkleDistributorTest_test_distributeOneTime');
 
         // Expect the correct event to be emitted
         vm.expectEmit(true, true, false, true);
-        emit IMerkleDistributor.OneTimeDistributionAdded(owner, address(swiseToken), amount, 'ipfsHash');
+        emit IMerkleDistributor.OneTimeDistributionAdded(owner, address(swiseToken), amount, 'ipfsHash', 'extraData');
 
         // Perform the one-time distribution
-        distributor.distributeOneTime(address(swiseToken), amount, 'ipfsHash');
+        distributor.distributeOneTime(address(swiseToken), amount, 'ipfsHash', 'extraData');
 
-        snapEnd();
+        vm.stopSnapshotGas();
 
         vm.stopPrank();
 
@@ -193,7 +192,7 @@ contract MerkleDistributorTest is Test, GasSnapshot {
         signatures = _sign(oraclePrivateKey, newRewardsRoot, newIpfsHash, nonceBefore);
 
         // Start snapshot for state and events
-        snapStart('MerkleDistributorTest_test_setRewardsRoot');
+        vm.startSnapshotGas('MerkleDistributorTest_test_setRewardsRoot');
 
         // Expect the correct event to be emitted
         vm.expectEmit(true, true, false, true);
@@ -203,7 +202,7 @@ contract MerkleDistributorTest is Test, GasSnapshot {
         distributor.setRewardsRoot(newRewardsRoot, newIpfsHash, signatures);
 
         // End snapshot
-        snapEnd();
+        vm.stopSnapshotGas();
 
         // Stop impersonation
         vm.stopPrank();
@@ -268,13 +267,13 @@ contract MerkleDistributorTest is Test, GasSnapshot {
         emit IMerkleDistributor.RewardsClaimed(address(this), user, tokens, cumulativeAmounts);
 
         // Start snapshot for state and events
-        snapStart('MerkleDistributorTest_test_claim');
+        vm.startSnapshotGas('MerkleDistributorTest_test_claim');
 
         // Successful claim
         distributor.claim(user, tokens, cumulativeAmounts, merkleProof);
 
         // End snapshot
-        snapEnd();
+        vm.stopSnapshotGas();
 
         assertEq(swiseToken.balanceOf(user), swiseAmount);
         assertEq(daiToken.balanceOf(user), daiAmount);

test/StrategiesRegistry.t.sol

@@ -3,12 +3,11 @@
 pragma solidity ^0.8.26;
 
 import {Test} from 'forge-std/Test.sol';
-import {GasSnapshot} from 'forge-gas-snapshot/GasSnapshot.sol';
 import {Ownable} from '@openzeppelin/contracts/access/Ownable.sol';
 import {Errors} from '@stakewise-core/libraries/Errors.sol';
 import {IStrategiesRegistry, StrategiesRegistry} from '../src/StrategiesRegistry.sol';
 
-contract StrategiesRegistryTest is Test, GasSnapshot {
+contract StrategiesRegistryTest is Test {
     StrategiesRegistry public registry;
     address public owner = address(0x123);
     address public strategy = address(0x456);
@@ -47,9 +46,9 @@ contract StrategiesRegistryTest is Test, GasSnapshot {
         // Set a new strategy and verify it's enabled
         vm.expectEmit(true, false, false, true);
         emit IStrategiesRegistry.StrategyUpdated(owner, strategy, true);
-        snapStart('StrategiesRegistryTest_test_setStrategy');
+        vm.startSnapshotGas('StrategiesRegistryTest_test_setStrategy');
         registry.setStrategy(strategy, true);
-        snapEnd();
+        vm.stopSnapshotGas();
         assertEq(registry.strategies(strategy), true);
 
         // Try setting it again with the same status, expect revert
@@ -87,9 +86,9 @@ contract StrategiesRegistryTest is Test, GasSnapshot {
         // add proxy
         vm.expectEmit(true, true, true, false);
         emit IStrategiesRegistry.StrategyProxyAdded(strategy, proxyId, proxy);
-        snapStart('StrategiesRegistryTest_test_addStrategyProxy');
+        vm.startSnapshotGas('StrategiesRegistryTest_test_addStrategyProxy');
         registry.addStrategyProxy(proxyId, proxy);
-        snapEnd();
+        vm.stopSnapshotGas();
         assertEq(registry.strategyProxies(proxy), true);
         assertEq(registry.strategyProxyIdToProxy(proxyId), proxy);
 
@@ -111,9 +110,9 @@ contract StrategiesRegistryTest is Test, GasSnapshot {
         vm.startPrank(owner);
         vm.expectEmit(true, false, false, true);
         emit IStrategiesRegistry.StrategyConfigUpdated(strategyId, configName, configValue);
-        snapStart('StrategiesRegistryTest_test_setStrategyConfig');
+        vm.startSnapshotGas('StrategiesRegistryTest_test_setStrategyConfig');
         registry.setStrategyConfig(strategyId, configName, configValue);
-        snapEnd();
+        vm.stopSnapshotGas();
         vm.stopPrank();
     }
 }