node/vms/platformvm/state/stakers.go

// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package state

import (
	"errors"
	"fmt"
	"sync"

	"github.com/google/btree"

	"github.com/luxfi/container/iterator"
	"github.com/luxfi/database"
	"github.com/luxfi/ids"
)

var ErrAddingStakerAfterDeletion = errors.New("attempted to add a staker after deleting it")

type Stakers interface {
	CurrentStakers
	PendingStakers
}

type CurrentStakers interface {
	// GetCurrentValidator returns the [staker] describing the validator on
	// [netID] with [nodeID]. If the validator does not exist,
	// [database.ErrNotFound] is returned.
	GetCurrentValidator(netID ids.ID, nodeID ids.NodeID) (*Staker, error)

	// PutCurrentValidator adds the [staker] describing a validator to the
	// staker set.
	//
	// Invariant: [staker] is not currently a CurrentValidator
	PutCurrentValidator(staker *Staker) error

	// DeleteCurrentValidator removes the [staker] describing a validator from
	// the staker set.
	//
	// Invariant: [staker] is currently a CurrentValidator
	DeleteCurrentValidator(staker *Staker)

	// SetDelegateeReward sets the accrued delegation rewards for [nodeID] on
	// [netID] to [amount].
	SetDelegateeReward(netID ids.ID, nodeID ids.NodeID, amount uint64) error

	// GetDelegateeReward returns the accrued delegation rewards for [nodeID] on
	// [netID].
	GetDelegateeReward(netID ids.ID, nodeID ids.NodeID) (uint64, error)

	// GetCurrentDelegatorIterator returns the delegators associated with the
	// validator on [netID] with [nodeID]. Delegators are sorted by their
	// removal from current staker set.
	GetCurrentDelegatorIterator(chainID ids.ID, nodeID ids.NodeID) (iterator.Iterator[*Staker], error)

	// PutCurrentDelegator adds the [staker] describing a delegator to the
	// staker set.
	//
	// Invariant: [staker] is not currently a CurrentDelegator
	PutCurrentDelegator(staker *Staker)

	// DeleteCurrentDelegator removes the [staker] describing a delegator from
	// the staker set.
	//
	// Invariant: [staker] is currently a CurrentDelegator
	DeleteCurrentDelegator(staker *Staker)

	// GetCurrentStakerIterator returns stakers in order of their removal from
	// the current staker set.
	GetCurrentStakerIterator() (iterator.Iterator[*Staker], error)
}

type PendingStakers interface {
	// GetPendingValidator returns the Staker describing the validator on
	// [netID] with [nodeID]. If the validator does not exist,
	// [database.ErrNotFound] is returned.
	GetPendingValidator(netID ids.ID, nodeID ids.NodeID) (*Staker, error)

	// PutPendingValidator adds the [staker] describing a validator to the
	// staker set.
	PutPendingValidator(staker *Staker) error

	// DeletePendingValidator removes the [staker] describing a validator from
	// the staker set.
	DeletePendingValidator(staker *Staker)

	// GetPendingDelegatorIterator returns the delegators associated with the
	// validator on [netID] with [nodeID]. Delegators are sorted by their
	// removal from pending staker set.
	GetPendingDelegatorIterator(chainID ids.ID, nodeID ids.NodeID) (iterator.Iterator[*Staker], error)

	// PutPendingDelegator adds the [staker] describing a delegator to the
	// staker set.
	PutPendingDelegator(staker *Staker)

	// DeletePendingDelegator removes the [staker] describing a delegator from
	// the staker set.
	DeletePendingDelegator(staker *Staker)

	// GetPendingStakerIterator returns stakers in order of their removal from
	// the pending staker set.
	GetPendingStakerIterator() (iterator.Iterator[*Staker], error)
}

type baseStakers struct {
	// mu protects concurrent access to the btree and maps
	mu sync.RWMutex
	// netID --> nodeID --> current state for the validator of the chain
	validators map[ids.ID]map[ids.NodeID]*baseStaker
	stakers    *btree.BTreeG[*Staker]
	// netID --> nodeID --> diff for that validator since the last db write
	validatorDiffs map[ids.ID]map[ids.NodeID]*diffValidator
}

type baseStaker struct {
	validator  *Staker
	delegators *btree.BTreeG[*Staker]
}

func newBaseStakers() *baseStakers {
	return &baseStakers{
		validators:     make(map[ids.ID]map[ids.NodeID]*baseStaker),
		stakers:        btree.NewG(defaultTreeDegree, (*Staker).Less),
		validatorDiffs: make(map[ids.ID]map[ids.NodeID]*diffValidator),
	}
}

func (v *baseStakers) GetValidator(netID ids.ID, nodeID ids.NodeID) (*Staker, error) {
	v.mu.RLock()
	defer v.mu.RUnlock()
	chainValidators, ok := v.validators[netID]
	if !ok {
		return nil, database.ErrNotFound
	}
	validator, ok := chainValidators[nodeID]
	if !ok {
		return nil, database.ErrNotFound
	}
	if validator.validator == nil {
		return nil, database.ErrNotFound
	}
	return validator.validator, nil
}

func (v *baseStakers) PutValidator(staker *Staker) {
	v.mu.Lock()
	defer v.mu.Unlock()
	validator := v.getOrCreateValidatorLocked(staker.ChainID, staker.NodeID)
	validator.validator = staker

	validatorDiff := v.getOrCreateValidatorDiffLocked(staker.ChainID, staker.NodeID)
	validatorDiff.validatorStatus = added
	validatorDiff.validator = staker

	v.stakers.ReplaceOrInsert(staker)
}

func (v *baseStakers) DeleteValidator(staker *Staker) {
	v.mu.Lock()
	defer v.mu.Unlock()
	validator := v.getOrCreateValidatorLocked(staker.ChainID, staker.NodeID)
	validator.validator = nil
	v.pruneValidatorLocked(staker.ChainID, staker.NodeID)

	validatorDiff := v.getOrCreateValidatorDiffLocked(staker.ChainID, staker.NodeID)
	validatorDiff.validatorStatus = deleted
	validatorDiff.validator = staker

	v.stakers.Delete(staker)
}

func (v *baseStakers) GetDelegatorIterator(chainID ids.ID, nodeID ids.NodeID) iterator.Iterator[*Staker] {
	v.mu.RLock()
	defer v.mu.RUnlock()
	chainValidators, ok := v.validators[chainID]
	if !ok {
		return iterator.Empty[*Staker]{}
	}
	validator, ok := chainValidators[nodeID]
	if !ok {
		return iterator.Empty[*Staker]{}
	}
	// Collect items into a slice to avoid holding the lock during iteration
	var items []*Staker
	if validator.delegators != nil {
		validator.delegators.Ascend(func(item *Staker) bool {
			items = append(items, item)
			return true
		})
	}
	return iterator.FromSlice(items...)
}

func (v *baseStakers) PutDelegator(staker *Staker) {
	v.mu.Lock()
	defer v.mu.Unlock()
	validator := v.getOrCreateValidatorLocked(staker.ChainID, staker.NodeID)
	if validator.delegators == nil {
		validator.delegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
	}
	validator.delegators.ReplaceOrInsert(staker)

	validatorDiff := v.getOrCreateValidatorDiffLocked(staker.ChainID, staker.NodeID)
	if validatorDiff.addedDelegators == nil {
		validatorDiff.addedDelegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
	}
	validatorDiff.addedDelegators.ReplaceOrInsert(staker)

	v.stakers.ReplaceOrInsert(staker)
}

func (v *baseStakers) DeleteDelegator(staker *Staker) {
	v.mu.Lock()
	defer v.mu.Unlock()
	validator := v.getOrCreateValidatorLocked(staker.ChainID, staker.NodeID)
	if validator.delegators != nil {
		validator.delegators.Delete(staker)
	}
	v.pruneValidatorLocked(staker.ChainID, staker.NodeID)

	validatorDiff := v.getOrCreateValidatorDiffLocked(staker.ChainID, staker.NodeID)
	if validatorDiff.deletedDelegators == nil {
		validatorDiff.deletedDelegators = make(map[ids.ID]*Staker)
	}
	validatorDiff.deletedDelegators[staker.TxID] = staker

	v.stakers.Delete(staker)
}

func (v *baseStakers) GetStakerIterator() iterator.Iterator[*Staker] {
	v.mu.RLock()
	defer v.mu.RUnlock()
	// Collect items into a slice to avoid holding the lock during iteration
	var items []*Staker
	v.stakers.Ascend(func(item *Staker) bool {
		items = append(items, item)
		return true
	})
	return iterator.FromSlice(items...)
}

// LoadValidator adds a validator during state initialization.
// Unlike PutValidator, this does not track diffs since it's loading from database.
func (v *baseStakers) LoadValidator(staker *Staker) {
	v.mu.Lock()
	defer v.mu.Unlock()
	validator := v.getOrCreateValidatorLocked(staker.ChainID, staker.NodeID)
	validator.validator = staker
	v.stakers.ReplaceOrInsert(staker)
}

// LoadDelegator adds a delegator during state initialization.
// Unlike PutDelegator, this does not track diffs since it's loading from database.
func (v *baseStakers) LoadDelegator(staker *Staker) {
	v.mu.Lock()
	defer v.mu.Unlock()
	validator := v.getOrCreateValidatorLocked(staker.ChainID, staker.NodeID)
	if validator.delegators == nil {
		validator.delegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
	}
	validator.delegators.ReplaceOrInsert(staker)
	v.stakers.ReplaceOrInsert(staker)
}

// getOrCreateValidatorLocked requires the caller to hold v.mu (write lock)
func (v *baseStakers) getOrCreateValidatorLocked(netID ids.ID, nodeID ids.NodeID) *baseStaker {
	chainValidators, ok := v.validators[netID]
	if !ok {
		chainValidators = make(map[ids.NodeID]*baseStaker)
		v.validators[netID] = chainValidators
	}
	validator, ok := chainValidators[nodeID]
	if !ok {
		validator = &baseStaker{}
		chainValidators[nodeID] = validator
	}
	return validator
}

// pruneValidatorLocked assumes that the named validator is currently in the
// [validators] map. Requires the caller to hold v.mu (write lock).
func (v *baseStakers) pruneValidatorLocked(netID ids.ID, nodeID ids.NodeID) {
	chainValidators := v.validators[netID]
	validator := chainValidators[nodeID]
	if validator.validator != nil {
		return
	}
	if validator.delegators != nil && validator.delegators.Len() > 0 {
		return
	}
	delete(chainValidators, nodeID)
	if len(chainValidators) == 0 {
		delete(v.validators, netID)
	}
}

// getOrCreateValidatorDiffLocked requires the caller to hold v.mu (write lock)
func (v *baseStakers) getOrCreateValidatorDiffLocked(netID ids.ID, nodeID ids.NodeID) *diffValidator {
	chainValidatorDiffs, ok := v.validatorDiffs[netID]
	if !ok {
		chainValidatorDiffs = make(map[ids.NodeID]*diffValidator)
		v.validatorDiffs[netID] = chainValidatorDiffs
	}
	validatorDiff, ok := chainValidatorDiffs[nodeID]
	if !ok {
		validatorDiff = &diffValidator{
			validatorStatus: unmodified,
		}
		chainValidatorDiffs[nodeID] = validatorDiff
	}
	return validatorDiff
}

type diffStakers struct {
	// netID --> nodeID --> diff for that validator
	validatorDiffs map[ids.ID]map[ids.NodeID]*diffValidator
	addedStakers   *btree.BTreeG[*Staker]
	deletedStakers map[ids.ID]*Staker
}

type diffValidator struct {
	// validatorStatus describes whether a validator has been added or removed.
	//
	// validatorStatus is not affected by delegators ops so unmodified does not
	// mean that diffValidator hasn't change, since delegators may have changed.
	validatorStatus diffValidatorStatus
	validator       *Staker

	addedDelegators   *btree.BTreeG[*Staker]
	deletedDelegators map[ids.ID]*Staker
}

func (d *diffValidator) WeightDiff() (ValidatorWeightDiff, error) {
	weightDiff := ValidatorWeightDiff{
		Decrease: d.validatorStatus == deleted,
	}
	if d.validatorStatus != unmodified {
		weightDiff.Amount = d.validator.Weight
		// DO NOT set ValidationID here - it's set by L1 validator state management only
		// Setting it here causes TxID to change incorrectly for delegator operations
	}

	for _, staker := range d.deletedDelegators {
		if err := weightDiff.Sub(staker.Weight); err != nil {
			return ValidatorWeightDiff{}, fmt.Errorf("failed to decrease node weight diff: %w", err)
		}
	}

	addedDelegatorIterator := iterator.FromTree(d.addedDelegators)
	defer addedDelegatorIterator.Release()

	for addedDelegatorIterator.Next() {
		staker := addedDelegatorIterator.Value()

		if err := weightDiff.Add(staker.Weight); err != nil {
			return ValidatorWeightDiff{}, fmt.Errorf("failed to increase node weight diff: %w", err)
		}
	}

	return weightDiff, nil
}

// GetValidator attempts to fetch the validator with the given chainID and
// nodeID.
// Invariant: Assumes that the validator will never be removed and then added.
func (s *diffStakers) GetValidator(netID ids.ID, nodeID ids.NodeID) (*Staker, diffValidatorStatus) {
	chainValidatorDiffs, ok := s.validatorDiffs[netID]
	if !ok {
		return nil, unmodified
	}

	validatorDiff, ok := chainValidatorDiffs[nodeID]
	if !ok {
		return nil, unmodified
	}

	if validatorDiff.validatorStatus == added {
		return validatorDiff.validator, added
	}
	return nil, validatorDiff.validatorStatus
}

func (s *diffStakers) PutValidator(staker *Staker) error {
	validatorDiff := s.getOrCreateDiff(staker.ChainID, staker.NodeID)
	if validatorDiff.validatorStatus == deleted {
		// Enforce the invariant that a validator cannot be added after being
		// deleted.
		return ErrAddingStakerAfterDeletion
	}

	validatorDiff.validatorStatus = added
	validatorDiff.validator = staker

	if s.addedStakers == nil {
		s.addedStakers = btree.NewG(defaultTreeDegree, (*Staker).Less)
	}
	s.addedStakers.ReplaceOrInsert(staker)
	return nil
}

func (s *diffStakers) DeleteValidator(staker *Staker) {
	validatorDiff := s.getOrCreateDiff(staker.ChainID, staker.NodeID)
	if validatorDiff.validatorStatus == added {
		// This validator was added and immediately removed in this diff. We
		// treat it as if it was never added.
		validatorDiff.validatorStatus = unmodified
		s.addedStakers.Delete(validatorDiff.validator)
		validatorDiff.validator = nil
	} else {
		validatorDiff.validatorStatus = deleted
		validatorDiff.validator = staker
		if s.deletedStakers == nil {
			s.deletedStakers = make(map[ids.ID]*Staker)
		}
		s.deletedStakers[staker.TxID] = staker
	}
}

func (s *diffStakers) GetDelegatorIterator(
	parentIterator iterator.Iterator[*Staker],
	chainID ids.ID,
	nodeID ids.NodeID,
) iterator.Iterator[*Staker] {
	var (
		addedDelegatorIterator iterator.Iterator[*Staker] = iterator.Empty[*Staker]{}
		deletedDelegators      map[ids.ID]*Staker
	)
	if chainValidatorDiffs, ok := s.validatorDiffs[chainID]; ok {
		if validatorDiff, ok := chainValidatorDiffs[nodeID]; ok {
			addedDelegatorIterator = iterator.FromTree(validatorDiff.addedDelegators)
			deletedDelegators = validatorDiff.deletedDelegators
		}
	}

	return iterator.Filter(
		iterator.Merge(
			(*Staker).Less,
			parentIterator,
			addedDelegatorIterator,
		),
		func(staker *Staker) bool {
			_, ok := deletedDelegators[staker.TxID]
			return ok
		},
	)
}

func (s *diffStakers) PutDelegator(staker *Staker) {
	validatorDiff := s.getOrCreateDiff(staker.ChainID, staker.NodeID)
	if validatorDiff.addedDelegators == nil {
		validatorDiff.addedDelegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
	}
	validatorDiff.addedDelegators.ReplaceOrInsert(staker)

	if s.addedStakers == nil {
		s.addedStakers = btree.NewG(defaultTreeDegree, (*Staker).Less)
	}
	s.addedStakers.ReplaceOrInsert(staker)
}

func (s *diffStakers) DeleteDelegator(staker *Staker) {
	validatorDiff := s.getOrCreateDiff(staker.ChainID, staker.NodeID)
	if validatorDiff.deletedDelegators == nil {
		validatorDiff.deletedDelegators = make(map[ids.ID]*Staker)
	}
	validatorDiff.deletedDelegators[staker.TxID] = staker

	if s.deletedStakers == nil {
		s.deletedStakers = make(map[ids.ID]*Staker)
	}
	s.deletedStakers[staker.TxID] = staker
}

func (s *diffStakers) GetStakerIterator(parentIterator iterator.Iterator[*Staker]) iterator.Iterator[*Staker] {
	return iterator.Filter(
		iterator.Merge(
			(*Staker).Less,
			parentIterator,
			iterator.FromTree(s.addedStakers),
		),
		func(staker *Staker) bool {
			_, ok := s.deletedStakers[staker.TxID]
			return ok
		},
	)
}

func (s *diffStakers) getOrCreateDiff(netID ids.ID, nodeID ids.NodeID) *diffValidator {
	if s.validatorDiffs == nil {
		s.validatorDiffs = make(map[ids.ID]map[ids.NodeID]*diffValidator)
	}
	chainValidatorDiffs, ok := s.validatorDiffs[netID]
	if !ok {
		chainValidatorDiffs = make(map[ids.NodeID]*diffValidator)
		s.validatorDiffs[netID] = chainValidatorDiffs
	}
	validatorDiff, ok := chainValidatorDiffs[nodeID]
	if !ok {
		validatorDiff = &diffValidator{
			validatorStatus: unmodified,
		}
		chainValidatorDiffs[nodeID] = validatorDiff
	}
	return validatorDiff
}