Account for validation ending before delegation

contracts/staking/PoSValidatorManager.sol

@@ -337,52 +337,57 @@ abstract contract PoSValidatorManager is IPoSValidatorManager, ValidatorManager
 
         // Ensure the delegator is active
         Delegator memory delegator = $._delegatorStakes[delegationID];
+        Validator memory validator = getValidator(validationID);
         require(
             delegator.status == DelegatorStatus.Active, "PoSValidatorManager: delegation not active"
         );
         // Only the delegation owner can end the delegation.
         require(
             delegator.owner == _msgSender(), "PoSValidatorManager: delegation not owned by sender"
         );
-        uint64 nonce = _incrementAndGetNonce(validationID);
 
         // Set the delegator status to pending removed, so that it can be properly removed in
         // the complete step, even if the delivered nonce is greater than the nonce used to
         // initialize the removal.
         delegator.status = DelegatorStatus.PendingRemoved;
-        delegator.endedAt = uint64(block.timestamp);
-        delegator.endingNonce = nonce;
-
-        $._delegatorStakes[delegationID] = delegator;
-
-        Validator memory validator = getValidator(validationID);
-        require(validator.weight > delegator.weight, "PoSValidatorManager: Invalid weight");
 
-        // Check that removing this delegator would not exceed the maximum churn rate.
-        // We only need to check this is the validator is still active. If the validator ends its validation
-        // period, the weight of all its delegators will be added to the churn tracker at that time. Ending
-        // a delegation whose validator has ended validating has no impact on the stake weight of the chain.
         if (validator.status == ValidatorStatus.Active) {
+            // Check that removing this delegator would not exceed the maximum churn rate.
+            // We only need to check this is the validator is still active. If the validator ends its validation
+            // period, the weight of all its delegators will be added to the churn tracker at that time. Ending
+            // a delegation whose validator has ended validating has no impact on the stake weight of the chain.
             _checkAndUpdateChurnTrackerRemoval(delegator.weight);
-        }
 
-        uint64 newValidatorWeight = validator.weight - delegator.weight;
-        _setValidatorWeight(validationID, newValidatorWeight);
+            delegator.endingNonce = _incrementAndGetNonce(validationID);
+            delegator.endedAt = uint64(block.timestamp);
+
+            uint64 newValidatorWeight = validator.weight - delegator.weight;
+            _setValidatorWeight(validationID, newValidatorWeight);
+
+            // Submit the message to the Warp precompile.
+            bytes32 messageID = WARP_MESSENGER.sendWarpMessage(
+                ValidatorMessages.packSetSubnetValidatorWeightMessage(
+                    validationID, delegator.endingNonce, newValidatorWeight
+                )
+            );
+
+            emit ValidatorWeightUpdate({
+                validationID: validationID,
+                nonce: delegator.endingNonce,
+                validatorWeight: newValidatorWeight,
+                setWeightMessageID: messageID
+            });
+        } else {
+            delegator.endingNonce = validator.messageNonce;
+            delegator.endedAt = validator.endedAt;
+        }
 
-        // Submit the message to the Warp precompile.
-        bytes32 messageID = WARP_MESSENGER.sendWarpMessage(
-            ValidatorMessages.packSetSubnetValidatorWeightMessage(
-                validationID, nonce, newValidatorWeight
-            )
-        );
+        $._delegatorStakes[delegationID] = delegator;
 
         emit DelegatorRemovalInitialized({
             delegationID: delegationID,
             validationID: validationID,
-            nonce: nonce,
-            validatorWeight: newValidatorWeight,
-            endTime: block.timestamp,
-            setWeightMessageID: messageID
+            endTime: delegator.endedAt
         });
     }
 

contracts/staking/interfaces/IPoSValidatorManager.sol

@@ -67,6 +67,7 @@ interface IPoSValidatorManager is IValidatorManager {
     /**
      * @notice Event emitted when a delegator registration is completed
      * @param delegationID The ID of the delegation
+     * @param validationID The ID of the validation period
      * @param nonce The message nonce used to update the validator weight, as returned by the P-Chain
      * @param startTime The time at which the registration was completed
      */
@@ -80,17 +81,24 @@ interface IPoSValidatorManager is IValidatorManager {
     /**
      * @notice Event emitted when delegator removal is initiated
      * @param delegationID The ID of the delegation
+     * @param validationID The ID of the validation period
+     * @param endTime The time at which the removal was initiated
+     */
+    event DelegatorRemovalInitialized(
+        bytes32 indexed delegationID, bytes32 indexed validationID, uint256 endTime
+    );
+
+    /**
+     * @notice Event emitted when delegator removal is initiated
+     * @param validationID The ID of the validation period
      * @param nonce The message nonce used to update the validator weight
      * @param validatorWeight The updated validator weight that is sent to the P-Chain
-     * @param endTime The time at which the removal was initiated
      * @param setWeightMessageID The ID of the Warp message that updates the validator's weight on the P-Chain
      */
-    event DelegatorRemovalInitialized(
-        bytes32 indexed delegationID,
+    event ValidatorWeightUpdate(
         bytes32 indexed validationID,
         uint64 indexed nonce,
         uint64 validatorWeight,
-        uint256 endTime,
         bytes32 setWeightMessageID
     );
 

contracts/staking/tests/PoSValidatorManagerTests.t.sol

@@ -52,11 +52,13 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
     );
 
     event DelegatorRemovalInitialized(
-        bytes32 indexed delegationID,
+        bytes32 indexed delegationID, bytes32 indexed validationID, uint256 endTime
+    );
+
+    event ValidatorWeightUpdate(
         bytes32 indexed validationID,
         uint64 indexed nonce,
         uint64 validatorWeight,
-        uint256 endTime,
         bytes32 setWeightMessageID
     );
 
@@ -791,14 +793,19 @@ abstract contract PoSValidatorManagerTest is ValidatorManagerTest {
         _mockSendWarpMessage(setValidatorWeightPayload, bytes32(0));
 
         vm.expectEmit(true, true, true, true, address(posValidatorManager));
-        emit DelegatorRemovalInitialized({
-            delegationID: delegationID,
+        emit ValidatorWeightUpdate({
             validationID: validationID,
             nonce: expectedNonce,
             validatorWeight: expectedValidatorWeight,
-            endTime: endDelegationTimestamp,
             setWeightMessageID: bytes32(0)
         });
+        vm.expectEmit(true, true, true, true, address(posValidatorManager));
+        emit DelegatorRemovalInitialized({
+            delegationID: delegationID,
+            validationID: validationID,
+            endTime: endDelegationTimestamp
+        });
+
         vm.prank(delegatorAddress);
         posValidatorManager.initializeEndDelegation(delegationID, false, 0);
         return delegationID;