Fixed withdrawals

script/utils/interfaces/ICreateX.sol

@@ -77,12 +77,10 @@ interface ICreateX {
         payable
         returns (address newContract);
 
-    function deployCreate2AndInit(
-        bytes memory initCode,
-        bytes memory data,
-        Values memory values,
-        address refundAddress
-    ) external payable returns (address newContract);
+    function deployCreate2AndInit(bytes memory initCode, bytes memory data, Values memory values, address refundAddress)
+        external
+        payable
+        returns (address newContract);
 
     function deployCreate2AndInit(bytes memory initCode, bytes memory data, Values memory values)
         external
@@ -124,12 +122,10 @@ interface ICreateX {
         payable
         returns (address newContract);
 
-    function deployCreate3AndInit(
-        bytes memory initCode,
-        bytes memory data,
-        Values memory values,
-        address refundAddress
-    ) external payable returns (address newContract);
+    function deployCreate3AndInit(bytes memory initCode, bytes memory data, Values memory values, address refundAddress)
+        external
+        payable
+        returns (address newContract);
 
     function deployCreate3AndInit(bytes memory initCode, bytes memory data, Values memory values)
         external

src/contracts/delegator/NetworkRestakeDelegator.sol

@@ -130,8 +130,9 @@ contract NetworkRestakeDelegator is BaseDelegator, INetworkRestakeDelegator {
             revert AlreadySet();
         }
 
-        _totalOperatorNetworkShares[subnetwork]
-        .push(Time.timestamp(), totalOperatorNetworkShares(subnetwork) - operatorNetworkShares_ + shares);
+        _totalOperatorNetworkShares[subnetwork].push(
+            Time.timestamp(), totalOperatorNetworkShares(subnetwork) - operatorNetworkShares_ + shares
+        );
         _operatorNetworkShares[subnetwork][operator].push(Time.timestamp(), shares);
 
         emit SetOperatorNetworkShares(subnetwork, operator, shares);

src/contracts/slasher/VetoSlasher.sol

@@ -252,8 +252,7 @@ contract VetoSlasher is BaseSlasher, IVetoSlasher {
             }
 
             if (resolver_ != address(uint160(_resolver[subnetwork].latest()))) {
-                _resolver[subnetwork]
-                .push(
+                _resolver[subnetwork].push(
                     (IVault(vault_).currentEpochStart() + resolverSetEpochsDelay * IVault(vault_).epochDuration())
                     .toUint48(),
                     uint160(resolver_)

src/contracts/vault/Vault.sol

@@ -20,6 +20,7 @@ import {Time} from "@openzeppelin/contracts/utils/types/Time.sol";
 
 contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVault {
     using Checkpoints for Checkpoints.Trace256;
+    using Checkpoints for Checkpoints.Trace208;
     using Math for uint256;
     using SafeCast for uint256;
     using SafeERC20 for IERC20;
@@ -40,8 +41,14 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
      * @inheritdoc IVault
      */
     function totalStake() public view returns (uint256) {
-        uint256 epoch = currentEpoch();
-        return activeStake() + withdrawals[epoch] + withdrawals[epoch + 1];
+        uint208 lastBucket = _timeToBucket.latest();
+        uint256 lastWithdrawalShares = withdrawalShares[lastBucket];
+        return activeStake()
+            + (lastWithdrawalShares > 0
+                    ? (_withdrawalSharesPrefixes.latest()
+                        - _withdrawalSharesPrefixes.upperLookupRecent(uint48(block.timestamp)))
+                    .mulDiv(withdrawals[lastBucket], lastWithdrawalShares)
+                    : 0);
     }
 
     /**
@@ -69,17 +76,11 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
     /**
      * @inheritdoc IVault
      */
-    function withdrawalsOf(uint256 epoch, address account) public view returns (uint256) {
-        return
-            ERC4626Math.previewRedeem(withdrawalSharesOf[epoch][account], withdrawals[epoch], withdrawalShares[epoch]);
-    }
-
-    /**
-     * @inheritdoc IVault
-     */
-    function slashableBalanceOf(address account) external view returns (uint256) {
-        uint256 epoch = currentEpoch();
-        return activeBalanceOf(account) + withdrawalsOf(epoch, account) + withdrawalsOf(epoch + 1, account);
+    function withdrawalsOf(uint256 index, address account) public view returns (uint256) {
+        uint256 bucketIndex = _timeToBucket.upperLookupRecent(withdrawalUnlockAt(index, account));
+        return ERC4626Math.previewRedeem(
+            withdrawalSharesOf(index, account), withdrawals[bucketIndex], withdrawalShares[bucketIndex]
+        );
     }
 
     /**
@@ -116,9 +117,9 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
 
         mintedShares = ERC4626Math.previewDeposit(depositedAmount, activeShares_, activeStake_);
 
-        _activeStake.push(Time.timestamp(), activeStake_ + depositedAmount);
-        _activeShares.push(Time.timestamp(), activeShares_ + mintedShares);
-        _activeSharesOf[onBehalfOf].push(Time.timestamp(), activeSharesOf(onBehalfOf) + mintedShares);
+        _activeStake.push(uint48(block.timestamp), activeStake_ + depositedAmount);
+        _activeShares.push(uint48(block.timestamp), activeShares_ + mintedShares);
+        _activeSharesOf[onBehalfOf].push(uint48(block.timestamp), activeSharesOf(onBehalfOf) + mintedShares);
 
         emit Deposit(msg.sender, onBehalfOf, depositedAmount, mintedShares);
     }
@@ -176,38 +177,38 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
     /**
      * @inheritdoc IVault
      */
-    function claim(address recipient, uint256 epoch) external nonReentrant returns (uint256 amount) {
+    function claim(address recipient, uint256 index) external nonReentrant returns (uint256 amount) {
         if (recipient == address(0)) {
             revert InvalidRecipient();
         }
 
-        amount = _claim(epoch);
+        amount = _claim(index);
 
         IERC20(collateral).safeTransfer(recipient, amount);
 
-        emit Claim(msg.sender, recipient, epoch, amount);
+        emit Claim(msg.sender, recipient, index, amount);
     }
 
     /**
      * @inheritdoc IVault
      */
-    function claimBatch(address recipient, uint256[] calldata epochs) external nonReentrant returns (uint256 amount) {
+    function claimBatch(address recipient, uint256[] calldata indexes) external nonReentrant returns (uint256 amount) {
         if (recipient == address(0)) {
             revert InvalidRecipient();
         }
 
-        uint256 length = epochs.length;
+        uint256 length = indexes.length;
         if (length == 0) {
             revert InvalidLengthEpochs();
         }
 
         for (uint256 i; i < length; ++i) {
-            amount += _claim(epochs[i]);
+            amount += _claim(indexes[i]);
         }
 
         IERC20(collateral).safeTransfer(recipient, amount);
 
-        emit ClaimBatch(msg.sender, recipient, epochs, amount);
+        emit ClaimBatch(msg.sender, recipient, indexes, amount);
     }
 
     /**
@@ -217,46 +218,30 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
         if (msg.sender != slasher) {
             revert NotSlasher();
         }
-
-        uint256 currentEpoch_ = currentEpoch();
-        uint256 captureEpoch = epochAt(captureTimestamp);
-        if ((currentEpoch_ > 0 && captureEpoch < currentEpoch_ - 1) || captureEpoch > currentEpoch_) {
+        if (captureTimestamp + epochDuration < uint48(block.timestamp) || captureTimestamp >= uint48(block.timestamp)) {
             revert InvalidCaptureEpoch();
         }
+        uint208 lastBucket = _timeToBucket.latest();
+        uint256 lastWithdrawals = withdrawals[lastBucket];
+        uint256 lastWithdrawalShares = withdrawalShares[lastBucket];
+        uint256 unmaturedWithdrawalShares =
+            _withdrawalSharesPrefixes.latest() - _withdrawalSharesPrefixes.upperLookupRecent(uint48(block.timestamp));
+        uint256 unmaturedWithdrawals =
+            lastWithdrawalShares > 0 ? unmaturedWithdrawalShares.mulDiv(lastWithdrawals, lastWithdrawalShares) : 0;
 
         uint256 activeStake_ = activeStake();
-        uint256 nextWithdrawals = withdrawals[currentEpoch_ + 1];
-        if (captureEpoch == currentEpoch_) {
-            uint256 slashableStake = activeStake_ + nextWithdrawals;
-            slashedAmount = Math.min(amount, slashableStake);
-            if (slashedAmount > 0) {
-                uint256 activeSlashed = slashedAmount.mulDiv(activeStake_, slashableStake);
-                uint256 nextWithdrawalsSlashed = slashedAmount - activeSlashed;
-
-                _activeStake.push(Time.timestamp(), activeStake_ - activeSlashed);
-                withdrawals[captureEpoch + 1] = nextWithdrawals - nextWithdrawalsSlashed;
-            }
-        } else {
-            uint256 withdrawals_ = withdrawals[currentEpoch_];
-            uint256 slashableStake = activeStake_ + withdrawals_ + nextWithdrawals;
-            slashedAmount = Math.min(amount, slashableStake);
-            if (slashedAmount > 0) {
-                uint256 activeSlashed = slashedAmount.mulDiv(activeStake_, slashableStake);
-                uint256 nextWithdrawalsSlashed = slashedAmount.mulDiv(nextWithdrawals, slashableStake);
-                uint256 withdrawalsSlashed = slashedAmount - activeSlashed - nextWithdrawalsSlashed;
-
-                if (withdrawals_ < withdrawalsSlashed) {
-                    nextWithdrawalsSlashed += withdrawalsSlashed - withdrawals_;
-                    withdrawalsSlashed = withdrawals_;
-                }
+        uint256 slashableStake = activeStake_ + unmaturedWithdrawals;
+        slashedAmount = Math.min(amount, slashableStake);
+        if (slashedAmount > 0) {
+            _timeToBucket.push(uint48(block.timestamp), lastBucket + 1);
+            withdrawals[lastBucket] = lastWithdrawals - unmaturedWithdrawals;
+            withdrawalShares[lastBucket] = lastWithdrawalShares - unmaturedWithdrawalShares;
+            withdrawalShares[lastBucket + 1] = unmaturedWithdrawalShares;
 
-                _activeStake.push(Time.timestamp(), activeStake_ - activeSlashed);
-                withdrawals[currentEpoch_ + 1] = nextWithdrawals - nextWithdrawalsSlashed;
-                withdrawals[currentEpoch_] = withdrawals_ - withdrawalsSlashed;
-            }
-        }
+            uint256 activeSlashed = slashedAmount.mulDiv(activeStake_, slashableStake);
+            _activeStake.push(uint48(block.timestamp), activeStake_ - activeSlashed);
+            withdrawals[lastBucket + 1] = unmaturedWithdrawals - (slashedAmount - activeSlashed);
 
-        if (slashedAmount > 0) {
             IERC20(collateral).safeTransfer(burner, slashedAmount);
         }
 
@@ -370,39 +355,36 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
         virtual
         returns (uint256 mintedShares)
     {
-        _activeSharesOf[msg.sender].push(Time.timestamp(), activeSharesOf(msg.sender) - burnedShares);
-        _activeShares.push(Time.timestamp(), activeShares() - burnedShares);
-        _activeStake.push(Time.timestamp(), activeStake() - withdrawnAssets);
-
-        uint256 epoch = currentEpoch() + 1;
-        uint256 withdrawals_ = withdrawals[epoch];
-        uint256 withdrawalsShares_ = withdrawalShares[epoch];
+        _activeSharesOf[msg.sender].push(uint48(block.timestamp), activeSharesOf(msg.sender) - burnedShares);
+        _activeShares.push(uint48(block.timestamp), activeShares() - burnedShares);
+        _activeStake.push(uint48(block.timestamp), activeStake() - withdrawnAssets);
 
-        mintedShares = ERC4626Math.previewDeposit(withdrawnAssets, withdrawalsShares_, withdrawals_);
+        uint256 lastBucket = _timeToBucket.latest();
+        mintedShares =
+            ERC4626Math.previewDeposit(withdrawnAssets, withdrawalShares[lastBucket], withdrawals[lastBucket]);
+        withdrawals[lastBucket] += withdrawnAssets;
+        withdrawalShares[lastBucket] += mintedShares;
 
-        withdrawals[epoch] = withdrawals_ + withdrawnAssets;
-        withdrawalShares[epoch] = withdrawalsShares_ + mintedShares;
-        withdrawalSharesOf[epoch][claimer] += mintedShares;
+        uint48 unlockAt = uint48(block.timestamp) + epochDuration;
+        _withdrawalsOf[claimer].push(Withdrawal(false, unlockAt, mintedShares));
+        _withdrawalSharesPrefixes.push(unlockAt, _withdrawalSharesPrefixes.latest() + mintedShares);
 
         emit Withdraw(msg.sender, claimer, withdrawnAssets, burnedShares, mintedShares);
     }
 
-    function _claim(uint256 epoch) internal returns (uint256 amount) {
-        if (epoch >= currentEpoch()) {
-            revert InvalidEpoch();
-        }
-
-        if (isWithdrawalsClaimed[epoch][msg.sender]) {
+    function _claim(uint256 index) internal returns (uint256 amount) {
+        Withdrawal storage withdrawal = _withdrawalsOf[msg.sender][index];
+        if (withdrawal.claimed) {
             revert AlreadyClaimed();
         }
-
-        amount = withdrawalsOf(epoch, msg.sender);
-
+        if (withdrawal.unlockAt >= block.timestamp) {
+            revert WithdrawalNotMatured();
+        }
+        amount = withdrawalsOf(index, msg.sender);
         if (amount == 0) {
             revert InsufficientClaim();
         }
-
-        isWithdrawalsClaimed[epoch][msg.sender] = true;
+        withdrawal.claimed = true;
     }
 
     function _initialize(uint64, address, bytes memory data) internal virtual override {
@@ -442,7 +424,6 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
 
         burner = params.burner;
 
-        epochDurationInit = Time.timestamp();
         epochDuration = params.epochDuration;
 
         depositWhitelist = params.depositWhitelist;
@@ -467,6 +448,27 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
         }
     }
 
+    /**
+     * @inheritdoc IVault
+     */
+    function migrateWithdrawalsOf(address account, uint48 epoch) public {
+        if (_isEpochWithdrawalsClaimed[epoch][account]) {
+            revert();
+        }
+        uint256 shares = _epochWithdrawalSharesOf[epoch][account];
+        uint48 unlockAt = _epochDurationInit + (epoch + 1) * epochDuration;
+        if (unlockAt >= _withdrawalSharesPrefixes.at(0)._key) {
+            shares = ERC4626Math.previewRedeem(shares, _epochWithdrawals[epoch], _epochWithdrawalShares[epoch]);
+        } else if (withdrawalShares[epoch] == 0) {
+            withdrawals[epoch] = _epochWithdrawals[epoch];
+            withdrawalShares[epoch] = _epochWithdrawalShares[epoch];
+            _timeToBucket._trace._checkpoints[epoch]._key = unlockAt;
+            _timeToBucket._trace._checkpoints[epoch]._value = epoch;
+        }
+        _withdrawalsOf[account].push(Withdrawal(false, unlockAt, shares));
+        _isEpochWithdrawalsClaimed[epoch][account] = true;
+    }
+
     function _migrate(
         uint64,
         /* oldVersion */
@@ -477,6 +479,17 @@ contract Vault is VaultStorage, MigratableEntity, AccessControlUpgradeable, IVau
         internal
         override
     {
-        revert();
+        uint48 epoch = (block.timestamp - _epochDurationInit).toUint48() / epochDuration;
+        uint256 epochWithdrawals = _epochWithdrawals[epoch];
+        uint48 nextEpochStart = _epochDurationInit + (epoch + 1) * epochDuration;
+        _withdrawalSharesPrefixes.push(nextEpochStart, epochWithdrawals);
+        epochWithdrawals += _epochWithdrawals[epoch + 1];
+        _withdrawalSharesPrefixes.push(nextEpochStart + epochDuration, epochWithdrawals);
+        assembly ("memory-safe") {
+            sstore(_timeToBucket.slot, epoch)
+        }
+        _timeToBucket.push(nextEpochStart, epoch);
+        withdrawals[epoch] = epochWithdrawals;
+        withdrawalShares[epoch] = epochWithdrawals;
     }
 }