Feat/pre ipo claim

contracts/dao/PreIPODistributor.sol

@@ -21,8 +21,10 @@ import "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol";
  * On the first deposit each address selects a delivery tranche, locked for the whole sale
  * (top-ups inherit it). The tranche does not affect accounting; it only records the choice.
  *
- * Allocation is computed off-chain; this contract only holds deposits. Deposits are locked and
- * can only be topped up, never withdrawn. Settlement and delivery are added via UUPS upgrade.
+ * Allocation is computed off-chain; the contract holds deposits and, after the review window,
+ * a finalized settlement merkle root drives claims. Deposits are locked and can only be topped
+ * up, never withdrawn by the user. Claiming pays the refund and, for the unlocked tranche,
+ * delivers the share token; the locked tranche only records the amount for off-chain delivery.
  */
 contract PreIPODistributor is
     Initializable,
@@ -33,11 +35,11 @@ contract PreIPODistributor is
     using SafeERC20 for IERC20;
 
     bytes32 public constant MANAGER = keccak256("MANAGER");
-    // BOT operates the settlement lifecycle in the settlement/claim upgrade; granted here at
-    // deploy so the upgrade needs no extra role setup.
     bytes32 public constant BOT = keccak256("BOT");
 
-    // Delivery tranche selected at first deposit; 0 = not selected.
+    // Delivery tranche, selected at an account's first deposit and locked for the sale; 0 = not selected.
+    // On claim, TRANCHE_UNLOCKED transfers the share token to the account immediately, while
+    // TRANCHE_LOCKED only records the share amount (delivered off-chain at maturity).
     uint8 public constant TRANCHE_UNLOCKED = 1;
     uint8 public constant TRANCHE_LOCKED = 2;
 
@@ -68,6 +70,24 @@ contract PreIPODistributor is
 
     uint64 public nextSaleId;
 
+    struct Settlement {
+        bytes32 root;               // finalized settlement root; claims verify against this
+        bytes32 pendingRoot;        // pending root awaiting finalize (0 = none pending)
+        uint256 pendingTotalRefund; // total refund committed by the pending root
+        uint256 totalRefund;        // total refund committed by the finalized root
+        uint256 lastSetTime;        // block time the pending root was set
+        uint256 refunded;           // cumulative refund already paid out via claim
+    }
+
+    // saleId => Settlement
+    mapping(uint64 => Settlement) public settlements;
+
+    // saleId => account => claimed
+    mapping(uint64 => mapping(address => bool)) public claimed;
+
+    // review window between setSettlementRoot (step 1) and finalizeSettlement (step 2); min 6h
+    uint256 public waitingPeriod;
+
     event CreateSale(
         uint64 indexed saleId,
         address depositToken,
@@ -108,6 +128,23 @@ contract PreIPODistributor is
         uint256 pubTotalDeposits
     );
 
+    event SetSettlementRoot(uint64 indexed saleId, bytes32 pendingRoot, uint256 totalRefund, uint256 setTime);
+
+    event FinalizeSettlement(uint64 indexed saleId, bytes32 root, uint256 finalizeTime);
+
+    event RevokeSettlementRoot(uint64 indexed saleId);
+
+    event WaitingPeriodUpdated(uint256 waitingPeriod);
+
+    event Claimed(
+        uint64 indexed saleId,
+        address indexed account,
+        uint256 refundAmount,
+        address shareToken,
+        uint256 tokenAmount,
+        uint8 tranche
+    );
+
     event EmergencyWithdraw(address indexed token, address indexed to, uint256 amount);
 
     /// @custom:oz-upgrades-unsafe-allow constructor
@@ -132,6 +169,13 @@ contract PreIPODistributor is
         _setRoleAdmin(BOT, MANAGER);
     }
 
+    /// @dev Initializes settlement/claim state introduced by this version. Runs once, after the
+    ///      upgrade — call it atomically via upgradeToAndCall(newImpl, abi.encodeCall(initializeV2, ())).
+    function initializeV2() external reinitializer(2) {
+        waitingPeriod = 6 hours;
+        emit WaitingPeriodUpdated(waitingPeriod);
+    }
+
     /// @dev Create a sale (whitelist round). The public round is opened later via setPublicRound.
     function createSale(
         address _depositToken,
@@ -193,6 +237,11 @@ contract PreIPODistributor is
     {
         Sale storage sale = sales[_saleId];
         require(sale.whitelistRoot != bytes32(0), "Invalid saleId");
+        // cannot open/reschedule a public round once settlement is pending or finalized
+        require(
+            settlements[_saleId].pendingRoot == bytes32(0) && settlements[_saleId].root == bytes32(0),
+            "Settlement started"
+        );
         require(sale.pubStartTime == 0 || sale.pubStartTime > block.timestamp, "Public round started");
         require(_pubStartTime > sale.endTime, "Public must follow WL");
         require(_pubStartTime > block.timestamp, "Invalid pub start");
@@ -225,14 +274,15 @@ contract PreIPODistributor is
         require(_amount >= sale.minDeposit, "Below min deposit");
 
         // A non-zero existing WL deposit implies the caller already passed the proof check.
-        if (deposits[_saleId][msg.sender] == 0) {
+        uint256 prior = deposits[_saleId][msg.sender];
+        if (prior == 0) {
             bytes32 leaf = keccak256(abi.encode(block.chainid, msg.sender));
             require(MerkleProof.verifyCalldata(_proof, sale.whitelistRoot, leaf), "Invalid proof");
         }
 
         uint8 tranche = _applyTranche(_saleId, _tranche);
 
-        uint256 userTotal = deposits[_saleId][msg.sender] + _amount;
+        uint256 userTotal = prior + _amount;
         deposits[_saleId][msg.sender] = userTotal;
         sale.totalDeposits += _amount;
 
@@ -276,7 +326,146 @@ contract PreIPODistributor is
         }
     }
 
+    /// @dev Set the pending settlement root (step 1 of 2); starts the review window.
+    ///      A new pending root cannot be set while one is already in flight.
+    function setSettlementRoot(uint64 _saleId, bytes32 _root, uint256 _totalRefund)
+        external
+        onlyRole(BOT)
+    {
+        Sale storage sale = sales[_saleId];
+        require(sale.whitelistRoot != bytes32(0), "Invalid saleId");
+        require(!sale.paused, "Sale paused");
+        // deposit windows must be closed so no deposit can land after the off-chain snapshot
+        require(block.timestamp > sale.endTime, "WL round not ended");
+        require(sale.pubStartTime == 0 || block.timestamp > sale.pubEndTime, "Public round not ended");
+
+        Settlement storage s = settlements[_saleId];
+        // settlement is single-shot: once a root is finalized it cannot be replaced
+        require(s.root == bytes32(0), "Already finalized");
+        require(_root != bytes32(0), "Invalid root");
+        require(s.pendingRoot == bytes32(0), "Pending root in flight");
+        require(_totalRefund <= sale.totalDeposits + sale.pubTotalDeposits, "Refund exceeds deposits");
+
+        s.pendingRoot = _root;
+        s.pendingTotalRefund = _totalRefund;
+        s.lastSetTime = block.timestamp;
+
+        emit SetSettlementRoot(_saleId, _root, _totalRefund, block.timestamp);
+    }
+
+    /// @dev Finalize the pending settlement root (step 2 of 2); only after the review window.
+    function finalizeSettlement(uint64 _saleId) external onlyRole(BOT) {
+        require(!sales[_saleId].paused, "Sale paused");
+        Settlement storage s = settlements[_saleId];
+        require(s.pendingRoot != bytes32(0), "No pending root");
+        require(block.timestamp >= s.lastSetTime + waitingPeriod, "Review window not passed");
+
+        s.root = s.pendingRoot;
+        s.totalRefund = s.pendingTotalRefund;
+        s.pendingRoot = bytes32(0);
+        s.pendingTotalRefund = 0;
+
+        emit FinalizeSettlement(_saleId, s.root, block.timestamp);
+    }
+
+    /// @dev Revoke the pending settlement root before it is finalized.
+    function revokeSettlementRoot(uint64 _saleId) external onlyRole(MANAGER) {
+        Settlement storage s = settlements[_saleId];
+        require(s.pendingRoot != bytes32(0), "No pending root");
+        s.pendingRoot = bytes32(0);
+        s.pendingTotalRefund = 0;
+        emit RevokeSettlementRoot(_saleId);
+    }
+
+    /// @dev Claim a finalized allocation (whitelist + public rounds combined), once per account.
+    ///      Pays the refund; for the unlocked tranche also transfers the share token, while the
+    ///      locked tranche only records the amount (delivered off-chain at maturity). Permanent.
+    ///      Permissionless: anyone may call it on behalf of any account; the refund and any share
+    ///      delivery always go to `_account` (the address bound in the leaf), never to msg.sender.
+    /// @param _saleId Sale id
+    /// @param _account Account the allocation belongs to and that receives refund/shares
+    /// @param _refundAmount Refund amount in the deposit token (part of the leaf)
+    /// @param _shareToken Share token to deliver for the unlocked tranche (part of the leaf)
+    /// @param _tokenAmount Share amount (delivered for unlocked, recorded for locked; part of the leaf)
+    /// @param _proof Merkle proof of the leaf against the finalized settlement root
+    function claim(
+        uint64 _saleId,
+        address _account,
+        uint256 _refundAmount,
+        address _shareToken,
+        uint256 _tokenAmount,
+        bytes32[] calldata _proof
+    ) external nonReentrant {
+        require(!sales[_saleId].paused, "Sale paused");
+        Settlement storage s = settlements[_saleId];
+        require(s.root != bytes32(0), "Not finalized");
+        require(!claimed[_saleId][_account], "Already claimed");
+
+        bytes32 leaf =
+            keccak256(abi.encode(block.chainid, _saleId, _account, _refundAmount, _shareToken, _tokenAmount));
+        require(MerkleProof.verifyCalldata(_proof, s.root, leaf), "Invalid proof");
+
+        claimed[_saleId][_account] = true;
+
+        uint8 tranche = userTranche[_saleId][_account];
+
+        if (_refundAmount > 0) {
+            // aggregate refunds can never exceed the committed total
+            require(s.refunded + _refundAmount <= s.totalRefund, "Refund over total");
+            s.refunded += _refundAmount;
+            IERC20(sales[_saleId].depositToken).safeTransfer(_account, _refundAmount);
+        }
+        // Unlocked tranche receives the share token now; locked tranche is only recorded
+        // (delivered off-chain at maturity) via the Claimed event. A zero share token is only
+        // valid for the locked tranche.
+        if (tranche == TRANCHE_UNLOCKED) {
+            require(_shareToken != address(0), "Share token required");
+            if (_tokenAmount > 0) {
+                IERC20(_shareToken).safeTransfer(_account, _tokenAmount);
+            }
+        }
+
+        emit Claimed(_saleId, _account, _refundAmount, _shareToken, _tokenAmount, tranche);
+    }
+
+    /// @dev Read-only preview of what claim() would do for the given leaf inputs, without changing
+    ///      state. Mirrors claim()'s validation and delivery routing.
+    /// @return valid True if the sale is finalized and the proof matches the leaf
+    /// @return alreadyClaimed True if the account has already claimed
+    /// @return tranche The account's locked tranche (0 = never deposited)
+    /// @return sendShares True if a successful claim would transfer the share token now
+    ///         (unlocked tranche, non-zero amount and share token)
+    function previewClaim(
+        uint64 _saleId,
+        address _account,
+        uint256 _refundAmount,
+        address _shareToken,
+        uint256 _tokenAmount,
+        bytes32[] calldata _proof
+    ) external view returns (bool valid, bool alreadyClaimed, uint8 tranche, bool sendShares) {
+        bytes32 leaf =
+            keccak256(abi.encode(block.chainid, _saleId, _account, _refundAmount, _shareToken, _tokenAmount));
+        bytes32 root = settlements[_saleId].root;
+        valid = root != bytes32(0) && MerkleProof.verifyCalldata(_proof, root, leaf);
+        alreadyClaimed = claimed[_saleId][_account];
+        tranche = userTranche[_saleId][_account];
+        sendShares =
+            valid && !alreadyClaimed && tranche == TRANCHE_UNLOCKED && _shareToken != address(0) && _tokenAmount > 0;
+    }
+
+    /// @dev Update the review window between set and finalize (min 6h).
+    function setWaitingPeriod(uint256 _waitingPeriod) external onlyRole(MANAGER) {
+        require(_waitingPeriod >= 6 hours, "Waiting period too short");
+        waitingPeriod = _waitingPeriod;
+        emit WaitingPeriodUpdated(_waitingPeriod);
+    }
+
     /// @dev Manager (multisig) safety valve; there is no normal withdrawal path.
+    ///      WARNING: this does NOT adjust any internal accounting (totalDeposits, pubTotalDeposits,
+    ///      settlement totals). Withdrawing the deposit token can leave the contract without enough
+    ///      balance to satisfy outstanding refunds, causing finalized claims to revert and stranding
+    ///      those refunds. After using it, the tracked totals must be restaged (or the contract
+    ///      upgraded) before normal claims can resume. Use only in emergencies.
     function emergencyWithdraw(address _token, address _to, uint256 _amount)
         external
         onlyRole(MANAGER)