github.com/anacrolix/torrent@v1.61.0/internal/request-strategy/piece-request-order.go (about)

     1  package requestStrategy
     2  
     3  import (
     4  	"unique"
     5  
     6  	g "github.com/anacrolix/generics"
     7  	"github.com/anacrolix/missinggo/v2/panicif"
     8  
     9  	"github.com/anacrolix/torrent/metainfo"
    10  )
    11  
    12  type Btree interface {
    13  	Delete(PieceRequestOrderItem)
    14  	Add(PieceRequestOrderItem)
    15  	// This is an iterator. Go can't guarantee to not allocate if you return an iterator.
    16  	Scan(func(PieceRequestOrderItem) bool)
    17  	Contains(PieceRequestOrderItem) bool
    18  }
    19  
    20  func NewPieceOrder(btree Btree, cap int) *PieceRequestOrder {
    21  	return &PieceRequestOrder{
    22  		tree: btree,
    23  		keys: make(map[PieceRequestOrderKey]PieceRequestOrderState, cap),
    24  	}
    25  }
    26  
    27  type PieceRequestOrder struct {
    28  	tree Btree
    29  	keys map[PieceRequestOrderKey]PieceRequestOrderState
    30  }
    31  
    32  type PieceRequestOrderKey struct {
    33  	// It should be the canonical short hash.
    34  	InfoHash unique.Handle[metainfo.Hash]
    35  	Index    int
    36  }
    37  
    38  type PieceRequestOrderState struct {
    39  	Availability int
    40  	Priority     piecePriority
    41  	Partial      bool
    42  }
    43  
    44  type PieceRequestOrderItem struct {
    45  	Key   PieceRequestOrderKey
    46  	State PieceRequestOrderState
    47  }
    48  
    49  func (me *PieceRequestOrderItem) Less(otherConcrete *PieceRequestOrderItem) bool {
    50  	return pieceOrderLess(me, otherConcrete).Less()
    51  }
    52  
    53  // Returns the old state if the key was already present. The Update method needs to look at it.
    54  func (me *PieceRequestOrder) Add(
    55  	key PieceRequestOrderKey,
    56  	state PieceRequestOrderState,
    57  ) (old g.Option[PieceRequestOrderState]) {
    58  	if old.Value, old.Ok = me.keys[key]; old.Ok {
    59  		if state == old.Value {
    60  			return
    61  		}
    62  		me.tree.Delete(PieceRequestOrderItem{key, old.Value})
    63  	}
    64  	me.tree.Add(PieceRequestOrderItem{key, state})
    65  	me.keys[key] = state
    66  	return
    67  }
    68  
    69  func (me *PieceRequestOrder) Update(
    70  	key PieceRequestOrderKey,
    71  	state PieceRequestOrderState,
    72  ) (changed bool) {
    73  	old := me.Add(key, state)
    74  	if !old.Ok {
    75  		panic("Key should have been added already")
    76  	}
    77  	return old.Value != state
    78  }
    79  
    80  func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) (deleted bool) {
    81  	state, ok := me.keys[key]
    82  	if !ok {
    83  		return false
    84  	}
    85  	me.tree.Delete(PieceRequestOrderItem{key, state})
    86  	delete(me.keys, key)
    87  	return true
    88  }
    89  
    90  func (me *PieceRequestOrder) Len() int {
    91  	return len(me.keys)
    92  }
    93  
    94  func (me *PieceRequestOrder) Iter(yield func(PieceRequestOrderItem) bool) {
    95  	me.tree.Scan(func(item PieceRequestOrderItem) bool {
    96  		return yield(item)
    97  	})
    98  }
    99  
   100  func (me *PieceRequestOrder) Get(key PieceRequestOrderKey) (ret g.Option[PieceRequestOrderState]) {
   101  	ret.Value, ret.Ok = me.keys[key]
   102  	panicif.NotEq(ret.Ok, me.tree.Contains(PieceRequestOrderItem{key, ret.Value}))
   103  	return
   104  }