]> Sergey Matveev's repositories - btrtrc.git/blob - request-strategy/piece-request-order.go
More Windows test fixes
[btrtrc.git] / request-strategy / piece-request-order.go
1 package requestStrategy
2
3 import "github.com/anacrolix/torrent/metainfo"
4
5 type Btree interface {
6         Delete(pieceRequestOrderItem)
7         Add(pieceRequestOrderItem)
8         Scan(func(pieceRequestOrderItem) bool)
9 }
10
11 func NewPieceOrder(btree Btree, cap int) *PieceRequestOrder {
12         return &PieceRequestOrder{
13                 tree: btree,
14                 keys: make(map[PieceRequestOrderKey]PieceRequestOrderState, cap),
15         }
16 }
17
18 type PieceRequestOrder struct {
19         tree Btree
20         keys map[PieceRequestOrderKey]PieceRequestOrderState
21 }
22
23 type PieceRequestOrderKey struct {
24         InfoHash metainfo.Hash
25         Index    int
26 }
27
28 type PieceRequestOrderState struct {
29         Priority     piecePriority
30         Partial      bool
31         Availability int
32 }
33
34 type pieceRequestOrderItem struct {
35         key   PieceRequestOrderKey
36         state PieceRequestOrderState
37 }
38
39 func (me *pieceRequestOrderItem) Less(otherConcrete *pieceRequestOrderItem) bool {
40         return pieceOrderLess(me, otherConcrete).Less()
41 }
42
43 func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrderState) {
44         if _, ok := me.keys[key]; ok {
45                 panic(key)
46         }
47         me.tree.Add(pieceRequestOrderItem{key, state})
48         me.keys[key] = state
49 }
50
51 func (me *PieceRequestOrder) Update(
52         key PieceRequestOrderKey,
53         state PieceRequestOrderState,
54 ) {
55         oldState, ok := me.keys[key]
56         if !ok {
57                 panic("key should have been added already")
58         }
59         if state == oldState {
60                 return
61         }
62         me.tree.Delete(pieceRequestOrderItem{key, oldState})
63         me.tree.Add(pieceRequestOrderItem{key, state})
64         me.keys[key] = state
65 }
66
67 func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceRequestOrderItem {
68         return pieceRequestOrderItem{
69                 key:   key,
70                 state: me.keys[key],
71         }
72 }
73
74 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
75         me.tree.Delete(pieceRequestOrderItem{key, me.keys[key]})
76         delete(me.keys, key)
77 }
78
79 func (me *PieceRequestOrder) Len() int {
80         return len(me.keys)
81 }