Use generic btree implementation (ava-labs#2522)

utils/hashing/consistent/ring.go

@@ -12,8 +12,8 @@ import (
 )
 
 var (
-	_ Ring       = (*hashRing)(nil)
-	_ btree.Item = (*ringItem)(nil)
+	_ Ring                     = (*hashRing)(nil)
+	_ btree.LessFunc[ringItem] = ringItem.Less
 
 	errEmptyRing = errors.New("ring doesn't have any members")
 )
@@ -154,7 +154,7 @@ type hashRing struct {
 	virtualNodes int
 
 	lock sync.RWMutex
-	ring *btree.BTree
+	ring *btree.BTreeG[ringItem]
 }
 
 // RingConfig configures settings for a Ring.
@@ -172,7 +172,7 @@ func NewHashRing(config RingConfig) Ring {
 	return &hashRing{
 		hasher:       config.Hasher,
 		virtualNodes: config.VirtualNodes,
-		ring:         btree.New(config.Degree),
+		ring:         btree.NewG(config.Degree, ringItem.Less),
 	}
 }
 
@@ -200,8 +200,7 @@ func (h *hashRing) get(key Hashable) (Hashable, error) {
 			hash:  hash,
 			value: key,
 		},
-		func(itemIntf btree.Item) bool {
-			item := itemIntf.(ringItem)
+		func(item ringItem) bool {
 			if hash < item.hash {
 				result = item.value
 				return false
@@ -213,7 +212,8 @@ func (h *hashRing) get(key Hashable) (Hashable, error) {
 	// If found nothing ascending the tree, we need to wrap around the ring to
 	// the left-most (min) node.
 	if result == nil {
-		result = h.ring.Min().(ringItem).value
+		min, _ := h.ring.Min()
+		result = min.value
 	}
 	return result, nil
 }
@@ -260,9 +260,7 @@ func (h *hashRing) remove(key Hashable) bool {
 		item := ringItem{
 			hash: virtualNodeHash,
 		}
-		if h.ring.Delete(item) != nil {
-			removed = true
-		}
+		_, removed = h.ring.Delete(item)
 	}
 	return removed
 }
@@ -278,6 +276,6 @@ type ringItem struct {
 	value Hashable
 }
 
-func (r ringItem) Less(than btree.Item) bool {
-	return r.hash < than.(ringItem).hash
+func (r ringItem) Less(than ringItem) bool {
+	return r.hash < than.hash
 }

vms/platformvm/state/staker.go

@@ -14,7 +14,7 @@ import (
 	"github.com/ava-labs/avalanchego/vms/platformvm/txs"
 )
 
-var _ btree.Item = (*Staker)(nil)
+var _ btree.LessFunc[*Staker] = (*Staker).Less
 
 // StakerIterator defines an interface for iterating over a set of stakers.
 type StakerIterator interface {
@@ -63,11 +63,7 @@ type Staker struct {
 //     lesser one.
 //  3. If the priorities are also the same, the one with the lesser txID is
 //     lesser.
-//
-// Invariant: [thanIntf] is a *Staker.
-func (s *Staker) Less(thanIntf btree.Item) bool {
-	than := thanIntf.(*Staker)
-
+func (s *Staker) Less(than *Staker) bool {
 	if s.NextTime.Before(than.NextTime) {
 		return true
 	}

vms/platformvm/state/stakers.go

@@ -90,20 +90,20 @@ type PendingStakers interface {
 type baseStakers struct {
 	// subnetID --> nodeID --> current state for the validator of the subnet
 	validators map[ids.ID]map[ids.NodeID]*baseStaker
-	stakers    *btree.BTree
+	stakers    *btree.BTreeG[*Staker]
 	// subnetID --> 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.BTree
+	delegators *btree.BTreeG[*Staker]
 }
 
 func newBaseStakers() *baseStakers {
 	return &baseStakers{
 		validators:     make(map[ids.ID]map[ids.NodeID]*baseStaker),
-		stakers:        btree.New(defaultTreeDegree),
+		stakers:        btree.NewG(defaultTreeDegree, (*Staker).Less),
 		validatorDiffs: make(map[ids.ID]map[ids.NodeID]*diffValidator),
 	}
 }
@@ -161,13 +161,13 @@ func (v *baseStakers) GetDelegatorIterator(subnetID ids.ID, nodeID ids.NodeID) S
 func (v *baseStakers) PutDelegator(staker *Staker) {
 	validator := v.getOrCreateValidator(staker.SubnetID, staker.NodeID)
 	if validator.delegators == nil {
-		validator.delegators = btree.New(defaultTreeDegree)
+		validator.delegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
 	}
 	validator.delegators.ReplaceOrInsert(staker)
 
 	validatorDiff := v.getOrCreateValidatorDiff(staker.SubnetID, staker.NodeID)
 	if validatorDiff.addedDelegators == nil {
-		validatorDiff.addedDelegators = btree.New(defaultTreeDegree)
+		validatorDiff.addedDelegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
 	}
 	validatorDiff.addedDelegators.ReplaceOrInsert(staker)
 
@@ -242,7 +242,7 @@ func (v *baseStakers) getOrCreateValidatorDiff(subnetID ids.ID, nodeID ids.NodeI
 type diffStakers struct {
 	// subnetID --> nodeID --> diff for that validator
 	validatorDiffs map[ids.ID]map[ids.NodeID]*diffValidator
-	addedStakers   *btree.BTree
+	addedStakers   *btree.BTreeG[*Staker]
 	deletedStakers map[ids.ID]*Staker
 }
 
@@ -253,7 +253,7 @@ type diffValidator struct {
 	validatorDeleted bool
 	validator        *Staker
 
-	addedDelegators   *btree.BTree
+	addedDelegators   *btree.BTreeG[*Staker]
 	deletedDelegators map[ids.ID]*Staker
 }
 
@@ -294,7 +294,7 @@ func (s *diffStakers) PutValidator(staker *Staker) {
 	validatorDiff.validator = staker
 
 	if s.addedStakers == nil {
-		s.addedStakers = btree.New(defaultTreeDegree)
+		s.addedStakers = btree.NewG(defaultTreeDegree, (*Staker).Less)
 	}
 	s.addedStakers.ReplaceOrInsert(staker)
 }
@@ -343,12 +343,12 @@ func (s *diffStakers) GetDelegatorIterator(
 func (s *diffStakers) PutDelegator(staker *Staker) {
 	validatorDiff := s.getOrCreateDiff(staker.SubnetID, staker.NodeID)
 	if validatorDiff.addedDelegators == nil {
-		validatorDiff.addedDelegators = btree.New(defaultTreeDegree)
+		validatorDiff.addedDelegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
 	}
 	validatorDiff.addedDelegators.ReplaceOrInsert(staker)
 
 	if s.addedStakers == nil {
-		s.addedStakers = btree.New(defaultTreeDegree)
+		s.addedStakers = btree.NewG(defaultTreeDegree, (*Staker).Less)
 	}
 	s.addedStakers.ReplaceOrInsert(staker)
 }

vms/platformvm/state/state.go

@@ -1228,7 +1228,7 @@ func (s *state) loadCurrentValidators() error {
 
 			validator := s.currentStakers.getOrCreateValidator(staker.SubnetID, staker.NodeID)
 			if validator.delegators == nil {
-				validator.delegators = btree.New(defaultTreeDegree)
+				validator.delegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
 			}
 			validator.delegators.ReplaceOrInsert(staker)
 
@@ -1314,7 +1314,7 @@ func (s *state) loadPendingValidators() error {
 
 			validator := s.pendingStakers.getOrCreateValidator(staker.SubnetID, staker.NodeID)
 			if validator.delegators == nil {
-				validator.delegators = btree.New(defaultTreeDegree)
+				validator.delegators = btree.NewG(defaultTreeDegree, (*Staker).Less)
 			}
 			validator.delegators.ReplaceOrInsert(staker)
 

vms/platformvm/state/tree_iterator.go

@@ -23,7 +23,7 @@ type treeIterator struct {
 
 // NewTreeIterator returns a new iterator of the stakers in [tree] in ascending
 // order. Note that it isn't safe to modify [tree] while iterating over it.
-func NewTreeIterator(tree *btree.BTree) StakerIterator {
+func NewTreeIterator(tree *btree.BTreeG[*Staker]) StakerIterator {
 	if tree == nil {
 		return EmptyIterator
 	}
@@ -34,9 +34,9 @@ func NewTreeIterator(tree *btree.BTree) StakerIterator {
 	it.wg.Add(1)
 	go func() {
 		defer it.wg.Done()
-		tree.Ascend(func(i btree.Item) bool {
+		tree.Ascend(func(i *Staker) bool {
 			select {
-			case it.next <- i.(*Staker):
+			case it.next <- i:
 				return true
 			case <-it.release:
 				return false

vms/platformvm/state/tree_iterator_test.go

@@ -31,7 +31,7 @@ func TestTreeIterator(t *testing.T) {
 		},
 	}
 
-	tree := btree.New(defaultTreeDegree)
+	tree := btree.NewG(defaultTreeDegree, (*Staker).Less)
 	for _, staker := range stakers {
 		require.Nil(tree.ReplaceOrInsert(staker))
 	}
@@ -68,7 +68,7 @@ func TestTreeIteratorEarlyRelease(t *testing.T) {
 		},
 	}
 
-	tree := btree.New(defaultTreeDegree)
+	tree := btree.NewG(defaultTreeDegree, (*Staker).Less)
 	for _, staker := range stakers {
 		require.Nil(tree.ReplaceOrInsert(staker))
 	}