Update merkle airdrop to have conditions as stages

contracts/scripts/deployments/facets/DeployDropFacet.s.sol

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.23;
+
+//interfaces
+
+//libraries
+
+//contracts
+import {Deployer} from "contracts/scripts/common/Deployer.s.sol";
+import {FacetHelper} from "contracts/test/diamond/Facet.t.sol";
+import {DropFacet} from "contracts/src/tokens/drop/DropFacet.sol";
+
+contract DeployDropFacet is Deployer, FacetHelper {
+  // FacetHelper
+  constructor() {
+    addSelector(DropFacet.claimWithPenalty.selector);
+    addSelector(DropFacet.setClaimConditions.selector);
+    addSelector(DropFacet.getActiveClaimConditionId.selector);
+    addSelector(DropFacet.getClaimConditionById.selector);
+    addSelector(DropFacet.getSupplyClaimedByWallet.selector);
+  }
+
+  // Deploying
+  function versionName() public pure override returns (string memory) {
+    return "dropFacet";
+  }
+
+  function initializer() public pure override returns (bytes4) {
+    return DropFacet.__DropFacet_init.selector;
+  }
+
+  // function makeInitData(address claimToken) public pure returns (bytes memory) {
+  //   return abi.encodeWithSelector(initializer(), claimToken);
+  // }
+
+  function __deploy(address deployer) public override returns (address) {
+    vm.startBroadcast(deployer);
+    DropFacet dropFacet = new DropFacet();
+    vm.stopBroadcast();
+    return address(dropFacet);
+  }
+}

contracts/src/tokens/drop/DropFacet.sol

@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.23;
+
+// interfaces
+import {IDropFacet} from "contracts/src/tokens/drop/IDropFacet.sol";
+// libraries
+import {DropStorage} from "contracts/src/tokens/drop/DropStorage.sol";
+import {CurrencyTransfer} from "contracts/src/utils/libraries/CurrencyTransfer.sol";
+import {BasisPoints} from "contracts/src/utils/libraries/BasisPoints.sol";
+
+// contracts
+import {Facet} from "contracts/src/diamond/facets/Facet.sol";
+import {DropFacetBase} from "contracts/src/tokens/drop/DropFacetBase.sol";
+import {OwnableBase} from "contracts/src/diamond/facets/ownable/OwnableBase.sol";
+
+contract DropFacet is IDropFacet, DropFacetBase, OwnableBase, Facet {
+  using DropStorage for DropStorage.Layout;
+
+  function __DropFacet_init() external onlyInitializing {
+    _addInterface(type(IDropFacet).interfaceId);
+  }
+
+  function claimWithPenalty(
+    uint256 conditionId,
+    address account,
+    uint256 quantity,
+    bytes32[] calldata allowlistProof
+  ) external {
+    DropStorage.Layout storage ds = DropStorage.layout();
+
+    _verifyClaim(ds, conditionId, account, quantity, allowlistProof);
+
+    ClaimCondition storage condition = ds.getClaimConditionById(conditionId);
+
+    uint256 amount = quantity;
+    uint256 penaltyBps = condition.penaltyBps;
+    if (penaltyBps > 0) {
+      uint256 penaltyAmount = BasisPoints.calculate(quantity, penaltyBps);
+      amount = quantity - penaltyAmount;
+    }
+
+    _updateClaim(ds, conditionId, account, amount);
+
+    CurrencyTransfer.safeTransferERC20(
+      condition.currency,
+      address(this),
+      account,
+      amount
+    );
+
+    emit DropFacet_Claimed_WithPenalty(
+      conditionId,
+      msg.sender,
+      account,
+      amount
+    );
+  }
+
+  function claimAndStake(
+    address account,
+    uint256 quantity,
+    bytes32[] calldata allowlistProof
+  ) external {}
+
+  ///@inheritdoc IDropFacet
+  function setClaimConditions(
+    ClaimCondition[] calldata conditions,
+    bool resetEligibility
+  ) external onlyOwner {
+    DropStorage.Layout storage ds = DropStorage.layout();
+    _setClaimConditions(ds, conditions, resetEligibility);
+  }
+
+  ///@inheritdoc IDropFacet
+  function getActiveClaimConditionId() external view returns (uint256) {
+    return _getActiveConditionId(DropStorage.layout());
+  }
+
+  ///@inheritdoc IDropFacet
+  function getClaimConditionById(
+    uint256 conditionId
+  ) external view returns (ClaimCondition memory) {
+    return DropStorage.layout().getClaimConditionById(conditionId);
+  }
+
+  ///@inheritdoc IDropFacet
+  function getSupplyClaimedByWallet(
+    address account,
+    uint256 conditionId
+  ) external view returns (uint256) {
+    return DropStorage.layout().getSupplyClaimedByWallet(conditionId, account);
+  }
+}

contracts/src/tokens/drop/DropFacetBase.sol

@@ -0,0 +1,174 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.23;
+
+// interfaces
+import {IDropFacetBase} from "./IDropFacet.sol";
+
+// libraries
+import {DropStorage} from "./DropStorage.sol";
+import {CustomRevert} from "contracts/src/utils/libraries/CustomRevert.sol";
+import {MerkleProofLib} from "solady/utils/MerkleProofLib.sol";
+
+// contracts
+
+abstract contract DropFacetBase is IDropFacetBase {
+  using DropStorage for DropStorage.Layout;
+  using MerkleProofLib for bytes32[];
+
+  function _getActiveConditionId(
+    DropStorage.Layout storage ds
+  ) internal view returns (uint256) {
+    uint256 conditionStartId = ds.conditionStartId;
+    uint256 conditionCount = ds.conditionCount;
+
+    for (
+      uint256 i = conditionStartId + conditionCount;
+      i > conditionStartId;
+      i--
+    ) {
+      ClaimCondition memory condition = ds.conditionById[i - 1];
+      if (
+        block.timestamp >= condition.startTimestamp &&
+        (condition.endTimestamp == 0 ||
+          block.timestamp < condition.endTimestamp)
+      ) {
+        return i - 1;
+      }
+    }
+
+    CustomRevert.revertWith(DropFacet__NoActiveClaimCondition.selector);
+  }
+
+  function _verifyClaim(
+    DropStorage.Layout storage ds,
+    uint256 conditionId,
+    address account,
+    uint256 quantity,
+    bytes32[] calldata proof
+  ) internal view {
+    ClaimCondition memory condition = ds.getClaimConditionById(conditionId);
+
+    if (condition.merkleRoot == bytes32(0)) {
+      CustomRevert.revertWith(DropFacet__MerkleRootNotSet.selector);
+    }
+
+    if (quantity == 0) {
+      CustomRevert.revertWith(
+        DropFacet__QuantityMustBeGreaterThanZero.selector
+      );
+    }
+
+    // Check if the total claimed supply (including the current claim) exceeds the maximum claimable supply
+    if (condition.supplyClaimed + quantity > condition.maxClaimableSupply) {
+      CustomRevert.revertWith(DropFacet__ExceedsMaxClaimableSupply.selector);
+    }
+
+    if (block.timestamp < condition.startTimestamp) {
+      CustomRevert.revertWith(DropFacet__ClaimHasNotStarted.selector);
+    }
+
+    if (
+      condition.endTimestamp > 0 && block.timestamp >= condition.endTimestamp
+    ) {
+      CustomRevert.revertWith(DropFacet__ClaimHasEnded.selector);
+    }
+
+    // check if already claimed
+    if (ds.supplyClaimedByWallet[conditionId][account] > 0) {
+      CustomRevert.revertWith(DropFacet__AlreadyClaimed.selector);
+    }
+
+    bytes32 leaf = _createLeaf(account, quantity);
+    if (!proof.verifyCalldata(condition.merkleRoot, leaf)) {
+      CustomRevert.revertWith(DropFacet__InvalidProof.selector);
+    }
+  }
+
+  function _setClaimConditions(
+    DropStorage.Layout storage ds,
+    ClaimCondition[] calldata conditions,
+    bool resetEligibility
+  ) internal {
+    // get the existing claim condition count and start id
+    uint256 existingStartId = ds.conditionStartId;
+    uint256 existingConditionCount = ds.conditionCount;
+
+    /// @dev If the claim conditions are being reset, we assign a new uid to the claim conditions.
+    /// which ends up resetting the eligibility of the claim conditions in `supplyClaimedByWallet`.
+    uint256 newConditionCount = conditions.length;
+    uint256 newStartId = existingStartId;
+    if (resetEligibility) {
+      newStartId = existingStartId + existingConditionCount;
+    }
+
+    ds.conditionCount = newConditionCount;
+    ds.conditionStartId = newStartId;
+
+    uint256 lastConditionTimestamp;
+    for (uint256 i = 0; i < newConditionCount; i++) {
+      if (lastConditionTimestamp >= conditions[i].startTimestamp) {
+        CustomRevert.revertWith(
+          DropFacet__ClaimConditionsNotInAscendingOrder.selector
+        );
+      }
+
+      // check that amount already claimed is less than or equal to the max claimable supply
+      uint256 amountAlreadyClaimed = ds
+        .conditionById[newStartId + i]
+        .supplyClaimed;
+
+      if (amountAlreadyClaimed > conditions[i].maxClaimableSupply) {
+        CustomRevert.revertWith(DropFacet__CannotSetClaimConditions.selector);
+      }
+
+      ds.conditionById[newStartId + i] = conditions[i];
+      ds.conditionById[newStartId + i].supplyClaimed = amountAlreadyClaimed;
+      lastConditionTimestamp = conditions[i].startTimestamp;
+    }
+
+    // if _resetEligibility is true, we assign new uids to the claim conditions
+    // so we delete claim conditions with UID < newStartId
+    if (resetEligibility) {
+      for (uint256 i = existingStartId; i < newStartId; i++) {
+        delete ds.conditionById[i];
+      }
+    } else {
+      if (existingConditionCount > newConditionCount) {
+        for (uint256 i = newConditionCount; i < existingConditionCount; i++) {
+          delete ds.conditionById[newStartId + i];
+        }
+      }
+    }
+
+    emit DropFacet_ClaimConditionsUpdated(conditions, resetEligibility);
+  }
+
+  function _updateClaim(
+    DropStorage.Layout storage ds,
+    uint256 conditionId,
+    address account,
+    uint256 amount
+  ) internal {
+    ds.conditionById[conditionId].supplyClaimed += amount;
+    ds.supplyClaimedByWallet[conditionId][account] += amount;
+  }
+
+  // =============================================================
+  //                        Utilities
+  // =============================================================
+  function _createLeaf(
+    address account,
+    uint256 amount
+  ) internal pure returns (bytes32 leaf) {
+    assembly ("memory-safe") {
+      // Store the account address at memory location 0
+      mstore(0, account)
+      // Store the amount at memory location 0x20 (32 bytes after the account address)
+      mstore(0x20, amount)
+      // Compute the keccak256 hash of the account and amount, and store it at memory location 0
+      mstore(0, keccak256(0, 0x40))
+      // Compute the keccak256 hash of the previous hash (stored at memory location 0) and store it in the leaf variable
+      leaf := keccak256(0, 0x20)
+    }
+  }
+}

contracts/src/tokens/drop/DropStorage.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.23;
+
+// interfaces
+import {IDropFacetBase} from "./IDropFacet.sol";
+
+// libraries
+
+// contracts
+
+library DropStorage {
+  // keccak256(abi.encode(uint256(keccak256("diamond.facets.drop.storage")) - 1)) & ~bytes32(uint256(0xff))
+  bytes32 constant STORAGE_SLOT =
+    0xeda6a1e2ce6f1639b6d3066254ca87a2daf51c4f0ad5038d408bbab6cc2cab00;
+
+  struct Layout {
+    address claimToken;
+    uint256 conditionStartId;
+    uint256 conditionCount;
+    mapping(uint256 conditionId => mapping(address => uint256)) supplyClaimedByWallet;
+    mapping(uint256 conditionId => IDropFacetBase.ClaimCondition) conditionById;
+  }
+
+  function layout() internal pure returns (Layout storage l) {
+    assembly {
+      l.slot := STORAGE_SLOT
+    }
+  }
+
+  function getClaimToken(Layout storage ds) internal view returns (address) {
+    return ds.claimToken;
+  }
+
+  function getClaimConditionById(
+    Layout storage ds,
+    uint256 conditionId
+  ) internal view returns (IDropFacetBase.ClaimCondition storage) {
+    return ds.conditionById[conditionId];
+  }
+
+  function getSupplyClaimedByWallet(
+    Layout storage ds,
+    uint256 conditionId,
+    address account
+  ) internal view returns (uint256) {
+    return ds.supplyClaimedByWallet[conditionId][account];
+  }
+}