fix(gov/staker): removeStake not persisting updated reward state (#827)

* fix: add missing tree update logic

Author: notJoon

contract/r/gnoswap/v1/gov/staker/emission_reward_mananger.gno renamed to contract/r/gnoswap/v1/gov/staker/emission_reward_manager.gno

@@ -7,6 +7,8 @@ import (
 	u256 "gno.land/p/gnoswap/uint256"
 )
 
+var errFailedToCastRewardState = "failed to cast rewardStates's element to *EmissionRewardState: %T"
+
 // EmissionRewardManager manages the distribution of emission rewards to stakers.
 type EmissionRewardManager struct {
 	// rewardStates maps address to EmissionRewardState for tracking individual staker rewards
@@ -78,7 +80,7 @@ func (e *EmissionRewardManager) calculateAccumulatedRewardX128PerStake(
 	currentTimestamp int64,
 ) (*u256.Uint, error) {
 	// If we're looking at a past timestamp, return current state
-	if e.accumulatedTimestamp > currentTimestamp {
+	if currentTimestamp < e.accumulatedTimestamp {
 		return e.accumulatedRewardX128PerStake, nil
 	}
 
@@ -87,14 +89,16 @@ func (e *EmissionRewardManager) calculateAccumulatedRewardX128PerStake(
 		return e.accumulatedRewardX128PerStake, nil
 	}
 
-	// Calculate the newly distributed rewards since last update
+	// Newly distributed rewards since last update
 	distributedAmountDelta := currentDistributedAmount - e.distributedAmount
-	distributedAmountDeltaX128 := u256.NewUintFromInt64(distributedAmountDelta)
-	distributedAmountDeltaX128 = u256.Zero().Lsh(distributedAmountDeltaX128, 128)
+	if distributedAmountDelta <= 0 {
+		// Non-positive delta. nothing to do more.
+		return e.accumulatedRewardX128PerStake, nil
+	}
 
-	// Calculate reward per stake for the new distribution
+	// Reward per stake for the new distribution
 	distributedAmountDeltaX128PerStake := u256.Zero().Div(
-		distributedAmountDeltaX128,
+		u256.Zero().Lsh(u256.NewUintFromInt64(distributedAmountDelta), 128),
 		u256.NewUintFromInt64(e.totalStakedAmount),
 	)
 
@@ -117,12 +121,12 @@ func (e *EmissionRewardManager) updateAccumulatedRewardX128PerStake(
 	currentDistributedAmount int64,
 	currentTimestamp int64,
 ) error {
-	// Don't update if we're looking at a past timestamp
-	if e.accumulatedTimestamp > currentTimestamp {
+	// DO NOT apply out-of-order timestamps
+	if currentTimestamp < e.accumulatedTimestamp {
 		return nil
 	}
 
-	// Don't update if no tokens are staked
+	// to avoid accumulating a large delta later.
 	if e.totalStakedAmount == 0 {
 		return nil
 	}
@@ -136,7 +140,7 @@ func (e *EmissionRewardManager) updateAccumulatedRewardX128PerStake(
 		return err
 	}
 
-	e.accumulatedRewardX128PerStake = accumulatedRewardX128PerStake
+	e.accumulatedRewardX128PerStake = accumulatedRewardX128PerStake.Clone()
 	e.distributedAmount = currentDistributedAmount
 	e.accumulatedTimestamp = currentTimestamp
 
@@ -147,53 +151,49 @@ func (e *EmissionRewardManager) updateAccumulatedRewardX128PerStake(
 // This method ensures rewards are properly calculated before the stake change.
 // Adds stake for specified address and updates reward calculations.
 func (e *EmissionRewardManager) addStake(address string, amount int64, currentTimestamp int64) error {
-	rewardStateI, ok := e.rewardStates.Get(address)
-	if !ok {
-		rewardStateI = NewEmissionRewardState(e.accumulatedRewardX128PerStake)
+	rewardState, ok, err := e.getRewardState(address)
+	if err != nil {
+		return err
 	}
-
-	rewardState, ok := rewardStateI.(*EmissionRewardState)
 	if !ok {
-		return ufmt.Errorf(
-			"failed to cast rewardStates's element to *EmissionRewardState: %T",
-			rewardStateI,
-		)
+		// if the address is unseen, initialize a snapshot to avoid nil deref
+		rewardState = NewEmissionRewardState(e.accumulatedRewardX128PerStake.Clone())
 	}
 
-	err := rewardState.addStakeWithUpdateRewardDebtX128(amount, e.accumulatedRewardX128PerStake, currentTimestamp)
+	err = rewardState.addStakeWithUpdateRewardDebtX128(amount, e.accumulatedRewardX128PerStake, currentTimestamp)
 	if err != nil {
 		return err
 	}
 
 	e.rewardStates.Set(address, rewardState)
 	e.totalStakedAmount = e.totalStakedAmount + amount
-
 	return nil
 }
 
 // removeStake removes a stake for an address and updates their reward state.
 // This method ensures rewards are properly calculated before the stake change.
 // Removes stake for specified address and updates reward calculations.
 func (e *EmissionRewardManager) removeStake(address string, amount int64, currentTimestamp int64) error {
-	rewardStateI, ok := e.rewardStates.Get(address)
-	if !ok {
-		rewardStateI = NewEmissionRewardState(e.accumulatedRewardX128PerStake.Clone())
+	rewardState, ok, err := e.getRewardState(address)
+	if err != nil {
+		return err
 	}
-
-	rewardState, ok := rewardStateI.(*EmissionRewardState)
 	if !ok {
-		return ufmt.Errorf(
-			"failed to cast rewardStates's element to *EmissionRewardState: %T",
-			rewardStateI,
-		)
+		// if the address is unseen, initialize a snapshot to avoid nil deref
+		rewardState = NewEmissionRewardState(e.accumulatedRewardX128PerStake.Clone())
 	}
 
-	err := rewardState.removeStakeWithUpdateRewardDebtX128(amount, e.accumulatedRewardX128PerStake, currentTimestamp)
+	err = rewardState.removeStakeWithUpdateRewardDebtX128(amount, e.accumulatedRewardX128PerStake, currentTimestamp)
 	if err != nil {
 		return err
 	}
 
-	e.totalStakedAmount = e.totalStakedAmount - amount
+	// persist updated state
+	e.rewardStates.Set(address, rewardState)
+	e.totalStakedAmount -= amount
+	if e.totalStakedAmount < 0 {
+		e.totalStakedAmount = 0 // defensive clamp
+	}
 
 	return nil
 }
@@ -202,17 +202,9 @@ func (e *EmissionRewardManager) removeStake(address string, amount int64, curren
 // This method calculates and returns the amount of rewards claimed.
 // Claims available rewards for specified address.
 func (e *EmissionRewardManager) claimRewards(address string, currentTimestamp int64) (claimedRewardAmount int64, err error) {
-	rewardStateI, ok := e.rewardStates.Get(address)
-	if !ok {
-		return 0, nil
-	}
-
-	rewardState, ok := rewardStateI.(*EmissionRewardState)
-	if !ok {
-		return 0, ufmt.Errorf(
-			"failed to cast rewardStates's element to *EmissionRewardState: %T",
-			rewardStateI,
-		)
+	rewardState, ok, err := e.getRewardState(address)
+	if err != nil || !ok {
+		return 0, err
 	}
 
 	claimedRewardAmount, cErr := rewardState.claimRewardsWithUpdateRewardDebtX128(e.accumulatedRewardX128PerStake, currentTimestamp)
@@ -221,7 +213,6 @@ func (e *EmissionRewardManager) claimRewards(address string, currentTimestamp in
 	}
 
 	e.rewardStates.Set(address, rewardState)
-
 	return claimedRewardAmount, nil
 }
 
@@ -237,3 +228,15 @@ func NewEmissionRewardManager() *EmissionRewardManager {
 		rewardStates:                  avl.NewTree(),
 	}
 }
+
+func (e *EmissionRewardManager) getRewardState(addr string) (*EmissionRewardState, bool, error) {
+	ri, ok := e.rewardStates.Get(addr)
+	if !ok {
+		return nil, false, nil
+	}
+	rs, castOk := ri.(*EmissionRewardState)
+	if !castOk {
+		return nil, false, ufmt.Errorf(errFailedToCastRewardState, ri)
+	}
+	return rs, true, nil
+}

contract/r/gnoswap/v1/gov/staker/emission_reward_manager_test.gno

@@ -392,3 +392,72 @@ func TestEmissionRewardManager_updateAccumulatedRewardX128PerStake(t *testing.T)
 		})
 	}
 }
+
+// removeStake bug reproduction 1:
+// When removeStake is called for a new address (not present in the tree),
+// internally a NewEmissionRewardState(...) is created. If rewardStates.Set(...) is missing,
+// the new state is not persisted in the tree, and this was a bug.
+// With the fixed code, the state should be saved correctly.
+func TestEmissionRewardManager_removeStake_persistsStateForNewAddress(t *testing.T) {
+	manager := NewEmissionRewardManager()
+
+	// Prepare: advance global accumulator once for safety
+	manager.accumulatedRewardX128PerStake = u256.NewUint(0)
+	manager.updateAccumulatedRewardX128PerStake(0, 10)
+
+	// When: calling removeStake on an address that does not exist yet
+	addr := "user-new"
+	uassert.Equal(t, manager.totalStakedAmount, int64(0))
+	uassert.Equal(t, manager.GetTotalStakedAmount(), int64(0))
+
+	err := manager.removeStake(addr, 100, 20)
+	uassert.NoError(t, err)
+
+	// Then: state should be persisted in the tree
+	// (this fails in the buggy version where Set was missing)
+	if ri, ok := manager.rewardStates.Get(addr); !ok {
+		t.Fatalf("expected reward state to be persisted for new address after removeStake, but not found (bug)")
+	} else {
+		// Type and field checks
+		rs, castOK := ri.(*EmissionRewardState)
+		uassert.True(t, castOK)
+		uassert.NotEqual(t, rs, nil)
+		uassert.Equal(t, rs.stakedAmount, int64(0)) // clamped to 0, no negative values
+		// rewardDebtX128 should be a cloned snapshot
+		uassert.NotEqual(t, rs.rewardDebtX128, nil)
+	}
+}
+
+// removeStake bug reproduction 2 (defensive check):
+// After adding stake for an existing address and then removing part of it,
+// the updated staked amount must be persisted in the tree.
+// In the buggy version, this may rely on pointer aliasing instead of Set persistence.
+// This test ensures round-trip read/write correctness.
+func TestEmissionRewardManager_removeStake_updatesPersistedState(t *testing.T) {
+	manager := NewEmissionRewardManager()
+
+	addr := "user1"
+	// Given: add stake → creates and saves state
+	err := manager.addStake(addr, 1000, 50)
+	uassert.NoError(t, err)
+	uassert.Equal(t, manager.GetTotalStakedAmount(), int64(1000))
+
+	// When: remove part of the stake
+	err = manager.removeStake(addr, 300, 100)
+	uassert.NoError(t, err)
+	uassert.Equal(t, manager.GetTotalStakedAmount(), int64(700))
+
+	// Then: re-read from tree should reflect the reduced stake
+	ri, ok := manager.rewardStates.Get(addr)
+	uassert.True(t, ok)
+	rs, castOK := ri.(*EmissionRewardState)
+	uassert.True(t, castOK)
+	uassert.NotEqual(t, rs, nil)
+	uassert.Equal(t, rs.stakedAmount, int64(700))
+
+	// Extra check: subsequent operations should remain consistent
+	claimed, err := manager.claimRewards(addr, 120)
+	uassert.NoError(t, err)
+	// Since no rewards were distributed, claimed amount should be zero
+	uassert.Equal(t, claimed, int64(0))
+}