Drop support for go 1.20

[btrtrc.git] / request-strategy / piece-request-order.go

diff --git a/request-strategy/piece-request-order.go b/request-strategy/piece-request-order.go
index efe75604ecd94ae275618bd0ddc3865564ba7878..3056741db389a453de42e87899125ae563a5d658 100644 (file)
--- a/request-strategy/piece-request-order.go
+++ b/request-strategy/piece-request-order.go
@@ -1,19 +1,22 @@
-package request_strategy
+package requestStrategy
 
-import (
-       "github.com/anacrolix/torrent/metainfo"
-       "github.com/google/btree"
-)
+import "github.com/anacrolix/torrent/metainfo"
 
-func NewPieceOrder() *PieceRequestOrder {
+type Btree interface {
+       Delete(pieceRequestOrderItem)
+       Add(pieceRequestOrderItem)
+       Scan(func(pieceRequestOrderItem) bool)
+}
+
+func NewPieceOrder(btree Btree, cap int) *PieceRequestOrder {
        return &PieceRequestOrder{
-               tree: btree.New(32),
-               keys: make(map[PieceRequestOrderKey]PieceRequestOrderState),
+               tree: btree,
+               keys: make(map[PieceRequestOrderKey]PieceRequestOrderState, cap),
        }
 }
 
 type PieceRequestOrder struct {
-       tree *btree.BTree
+       tree Btree
        keys map[PieceRequestOrderKey]PieceRequestOrderState
 }
 
@@ -25,7 +28,7 @@ type PieceRequestOrderKey struct {
 type PieceRequestOrderState struct {
        Priority     piecePriority
        Partial      bool
-       Availability int64
+       Availability int
 }
 
 type pieceRequestOrderItem struct {
@@ -33,43 +36,31 @@ type pieceRequestOrderItem struct {
        state PieceRequestOrderState
 }
 
-func (me pieceRequestOrderItem) Less(other btree.Item) bool {
-       otherConcrete := other.(pieceRequestOrderItem)
-       return pieceOrderLess(
-               pieceOrderInput{
-                       PieceRequestOrderState: me.state,
-                       PieceRequestOrderKey:   me.key,
-               },
-               pieceOrderInput{
-                       PieceRequestOrderState: otherConcrete.state,
-                       PieceRequestOrderKey:   otherConcrete.key,
-               },
-       ).Less()
+func (me *pieceRequestOrderItem) Less(otherConcrete *pieceRequestOrderItem) bool {
+       return pieceOrderLess(me, otherConcrete).Less()
 }
 
 func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrderState) {
        if _, ok := me.keys[key]; ok {
                panic(key)
        }
-       if me.tree.ReplaceOrInsert(pieceRequestOrderItem{
-               key:   key,
-               state: state,
-       }) != nil {
-               panic("shouldn't already have this")
-       }
+       me.tree.Add(pieceRequestOrderItem{key, state})
        me.keys[key] = state
 }
 
-func (me *PieceRequestOrder) Update(key PieceRequestOrderKey, state PieceRequestOrderState) {
-       if me.tree.Delete(me.existingItemForKey(key)) == nil {
-               panic(key)
+func (me *PieceRequestOrder) Update(
+       key PieceRequestOrderKey,
+       state PieceRequestOrderState,
+) {
+       oldState, ok := me.keys[key]
+       if !ok {
+               panic("key should have been added already")
        }
-       if me.tree.ReplaceOrInsert(pieceRequestOrderItem{
-               key:   key,
-               state: state,
-       }) != nil {
-               panic(key)
+       if state == oldState {
+               return
        }
+       me.tree.Delete(pieceRequestOrderItem{key, oldState})
+       me.tree.Add(pieceRequestOrderItem{key, state})
        me.keys[key] = state
 }
 
@@ -81,8 +72,10 @@ func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceR
 }
 
 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
-       if me.tree.Delete(me.existingItemForKey(key)) == nil {
-               panic(key)
-       }
+       me.tree.Delete(pieceRequestOrderItem{key, me.keys[key]})
        delete(me.keys, key)
 }
+
+func (me *PieceRequestOrder) Len() int {
+       return len(me.keys)
+}