chore: simplify IsLocked logic (#1177)

darrenkel · web-flow · commit f8678d760ff5 · 2025-07-22T13:04:25.000+01:00
diff --git a/builtin/staker/delegations.go b/builtin/staker/delegations.go
@@ -115,14 +115,14 @@ func (d *delegations) DisableAutoRenew(delegationID thor.Bytes32) error {
 	if delegation.Stake.Sign() == 0 {
 		return errors.New("delegation is not active")
 	}
-	if validation.Status == StatusExit {
-		return errors.New("delegation is not active, withdraw is available")
+	if delegation.Ended(validation) {
+		return errors.New("delegation is not active")
 	}
 
 	weight := delegation.Weight()
 
 	// the delegation's funds have already been locked, so we need to move them to non-recurring, but still locked
-	if delegation.IsLocked(validation) {
+	if delegation.Started(validation) {
 		// move the delegation's portion of locked to non-recurring.
 		// this will make the funds available at the end of the current iteration
 		aggregation.CurrentRecurringVET = big.NewInt(0).Sub(aggregation.CurrentRecurringVET, delegation.Stake)
@@ -141,8 +141,15 @@ func (d *delegations) DisableAutoRenew(delegationID thor.Bytes32) error {
 		aggregation.PendingRecurringWeight = big.NewInt(0).Sub(aggregation.PendingRecurringWeight, weight)
 	}
 
+	// TODO: In a future PR this won't be possible, so it will be removed. This is backwards compatible according to the unit tests.
+	// - In future: delegations auto added as auto-renew, then will have to signal an exit to withdraw in the next staking period.
 	// set the delegation's exit iteration
-	lastIteration := validation.CurrentIteration() + 1
+	var lastIteration uint32
+	if delegation.Started(validation) {
+		lastIteration = validation.CurrentIteration()
+	} else {
+		lastIteration = delegation.FirstIteration
+	}
 	delegation.LastIteration = &lastIteration
 	delegation.AutoRenew = false
 
@@ -161,8 +168,8 @@ func (d *delegations) EnableAutoRenew(delegationID thor.Bytes32) error {
 	if delegation.AutoRenew {
 		return errors.New("delegation is already autoRenew")
 	}
-	if validation.Status == StatusExit {
-		return errors.New("delegation is not active, withdraw is available")
+	if delegation.Ended(validation) {
+		return errors.New("delegation is not active")
 	}
 	weight := delegation.Weight()
 
@@ -173,7 +180,7 @@ func (d *delegations) EnableAutoRenew(delegationID thor.Bytes32) error {
 		return errors.New("validation's next period stake exceeds max stake")
 	}
 
-	if delegation.IsLocked(validation) {
+	if delegation.Started(validation) {
 		// move the delegation's portion of non-recurring to locked.
 		// this means the funds will not be available until the validator is inactive, or the delegation signals an exit
 		// and completes the current staking period
@@ -205,19 +212,14 @@ func (d *delegations) Withdraw(delegationID thor.Bytes32) (*big.Int, error) {
 	if err != nil {
 		return nil, err
 	}
-	if delegation.IsLocked(validation) {
+	started := delegation.Started(validation)
+	finished := delegation.Ended(validation)
+	if started && !finished {
 		return nil, errors.New("delegation is not eligible for withdraw")
 	}
 	weight := delegation.Weight()
 
-	delegationStarted := delegation.FirstIteration <= validation.CompleteIterations
-	if delegationStarted && validation.Status != StatusQueued {
-		// the stake has moved to withdrawable since we checked if the validation is locked above
-		if aggregation.WithdrawableVET.Cmp(delegation.Stake) < 0 {
-			return nil, errors.New("not enough withdraw VET")
-		}
-		aggregation.WithdrawableVET = big.NewInt(0).Sub(aggregation.WithdrawableVET, delegation.Stake)
-	} else {
+	if !started {
 		if delegation.AutoRenew { // delegation's stake is pending locked
 			if aggregation.PendingRecurringVET.Cmp(delegation.Stake) < 0 {
 				return nil, errors.New("not enough pending locked VET")
@@ -239,6 +241,14 @@ func (d *delegations) Withdraw(delegationID thor.Bytes32) (*big.Int, error) {
 		}
 	}
 
+	if finished {
+		// the stake has moved to withdrawable since we checked if the validation is locked above
+		if aggregation.WithdrawableVET.Cmp(delegation.Stake) < 0 {
+			return nil, errors.New("not enough withdraw VET")
+		}
+		aggregation.WithdrawableVET = big.NewInt(0).Sub(aggregation.WithdrawableVET, delegation.Stake)
+	}
+
 	stake := delegation.Stake
 	delegation.Stake = big.NewInt(0)
 	// remove the delegation from the mapping after the withdraw
diff --git a/builtin/staker/delegations_test.go b/builtin/staker/delegations_test.go
@@ -52,7 +52,8 @@ func Test_IsLocked(t *testing.T) {
 			CompleteIterations: 2,
 		}
 
-		assert.False(t, d.IsLocked(v), "should not be locked when complete iterations is equal to last iteration")
+		assert.True(t, d.Started(v), "should not be locked when complete iterations is equal to last iteration")
+		assert.True(t, d.Ended(v), "should be locked when first is less than current and last is equal to current")
 	})
 
 	t.Run("Incomplete Staking Periods", func(t *testing.T) {
@@ -69,7 +70,8 @@ func Test_IsLocked(t *testing.T) {
 			CompleteIterations: 3,
 		}
 
-		assert.True(t, d.IsLocked(v), "should be locked when first is less than current and last is greater")
+		assert.True(t, d.Started(v), "should be started when complete iterations is greater than first iteration")
+		assert.False(t, d.Ended(v), "should not be locked when first is less than current and last is greater than current")
 	})
 
 	t.Run("Delegation Not Started", func(t *testing.T) {
@@ -86,7 +88,8 @@ func Test_IsLocked(t *testing.T) {
 			CompleteIterations: 3,
 		}
 
-		assert.False(t, d.IsLocked(v), "should not be locked if delegation has not started yet")
+		assert.False(t, d.Started(v), "should not be started when complete iterations is less than first iteration")
+		assert.False(t, d.Ended(v), "should not be locked when first is greater than current and last is greater than current")
 	})
 	t.Run("Staker is Queued", func(t *testing.T) {
 		d := &Delegation{
@@ -101,7 +104,8 @@ func Test_IsLocked(t *testing.T) {
 			CompleteIterations: 0,
 		}
 
-		assert.False(t, d.IsLocked(v), "should not be locked when validation status is queued")
+		assert.False(t, d.Started(v), "should not be started when validation status is queued")
+		assert.False(t, d.Ended(v), "should not be locked when validation status is queued")
 	})
 
 	t.Run("Exit block not defined", func(t *testing.T) {
@@ -117,7 +121,8 @@ func Test_IsLocked(t *testing.T) {
 			CompleteIterations: 0,
 		}
 
-		assert.True(t, d.IsLocked(v), "should be locked when last iteration is nil and first equals current")
+		assert.True(t, d.Started(v), "should be started when first iteration is less than current")
+		assert.False(t, d.Ended(v), "should not be locked when last iteration is nil and first equals current")
 	})
 }
 
@@ -662,7 +667,7 @@ func Test_Delegations_EnableAutoRenew_MatchStakeReached(t *testing.T) {
 	assert.NoError(t, err)
 	validation, err := staker.Get(validator.ID)
 	assert.NoError(t, err)
-	assert.False(t, delegation1.IsLocked(validation))
+	assert.False(t, delegation1.Started(validation))
 
 	// Delegation should become active
 	_, _, err = staker.Housekeep(validator.Period)
@@ -671,7 +676,7 @@ func Test_Delegations_EnableAutoRenew_MatchStakeReached(t *testing.T) {
 	assert.NoError(t, err)
 	validation, err = staker.Get(validator.ID)
 	assert.NoError(t, err)
-	assert.True(t, delegation1.IsLocked(validation))
+	assert.True(t, delegation1.Started(validation))
 
 	// Enable auto renew for delegation. Should be possible since the max stake won't be reached.
 	assert.NoError(t, staker.UpdateDelegationAutoRenew(delegationID, true))
diff --git a/builtin/staker/types.go b/builtin/staker/types.go
@@ -131,23 +131,37 @@ func (d *Delegation) Weight() *big.Int {
 	return weight
 }
 
-// IsLocked returns whether the delegator is locked for the current staking period.
-func (d *Delegation) IsLocked(validation *Validation) bool {
+// Started returns whether the delegation became locked
+func (d *Delegation) Started(validation *Validation) bool {
 	if d.IsEmpty() {
 		return false
 	}
-	// validation is not active, so the delegator is not locked
-	if validation.Status != StatusActive {
+	if validation.Status == StatusQueued {
+		return false // Delegation cannot start if the validation is not active
+	}
+	currentStakingPeriod := validation.CurrentIteration()
+	return currentStakingPeriod >= d.FirstIteration
+}
+
+// Ended returns whether the delegation has ended
+// It returns true if:
+// - the delegation's exit iteration is less than the current staking period
+// - OR if the validation is in exit status and the delegation has started
+func (d *Delegation) Ended(validation *Validation) bool {
+	if d.IsEmpty() {
 		return false
 	}
-	current := validation.CurrentIteration()
-	if d.LastIteration == nil {
-		return current >= d.FirstIteration
+	if validation.Status == StatusQueued {
+		return false // Delegation cannot end if the validation is not active
+	}
+	if validation.Status == StatusExit && d.Started(validation) {
+		return true // Delegation is ended if the validation is in exit status
 	}
-	if current < d.FirstIteration {
+	currentStakingPeriod := validation.CurrentIteration()
+	if d.LastIteration == nil {
 		return false
 	}
-	return current == d.FirstIteration || current < *d.LastIteration
+	return *d.LastIteration < currentStakingPeriod
 }
 
 // Aggregation represents the total amount of VET locked for a given validation's delegations.
diff --git a/builtin/staker_native.go b/builtin/staker_native.go
@@ -311,7 +311,8 @@ func init() {
 			if delegation.LastIteration != nil {
 				lastPeriod = *delegation.LastIteration
 			}
-			return []any{delegation.ValidationID, delegation.Stake, delegation.FirstIteration, lastPeriod, delegation.Multiplier, delegation.AutoRenew, delegation.IsLocked(validator), ""}
+			locked := delegation.Started(validator) && !delegation.Ended(validator)
+			return []any{delegation.ValidationID, delegation.Stake, delegation.FirstIteration, lastPeriod, delegation.Multiplier, delegation.AutoRenew, locked, ""}
 		}},
 		{"native_getRewards", func(env *xenv.Environment) []any {
 			var args struct {
