/
wantlist.go
142 lines (116 loc) · 3.25 KB
/
wantlist.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
// Package wantlist implements an object for bitswap that contains the keys
// that a given peer wants.
package wantlist
import (
"sort"
pb "github.com/Genon2/ipfs-thesis-bitswap/message/pb"
cid "github.com/ipfs/go-cid"
)
// Wantlist is a raw list of wanted blocks and their priorities
type Wantlist struct {
set map[cid.Cid]Entry
// Re-computing this can get expensive so we memoize it.
cached []Entry
}
// Entry is an entry in a want list, consisting of a cid and its priority
type Entry struct {
Cid cid.Cid
Priority int32
WantType pb.Message_Wantlist_WantType
}
// NewRefEntry creates a new reference tracked wantlist entry.
func NewRefEntry(c cid.Cid, p int32) Entry {
return Entry{
Cid: c,
Priority: p,
WantType: pb.Message_Wantlist_Block,
}
}
type entrySlice []Entry
func (es entrySlice) Len() int { return len(es) }
func (es entrySlice) Swap(i, j int) { es[i], es[j] = es[j], es[i] }
func (es entrySlice) Less(i, j int) bool { return es[i].Priority > es[j].Priority }
// New generates a new raw Wantlist
func New() *Wantlist {
return &Wantlist{
set: make(map[cid.Cid]Entry),
}
}
// Len returns the number of entries in a wantlist.
func (w *Wantlist) Len() int {
return len(w.set)
}
// Add adds an entry in a wantlist from CID & Priority, if not already present.
func (w *Wantlist) Add(c cid.Cid, priority int32, wantType pb.Message_Wantlist_WantType) bool {
e, ok := w.set[c]
// Adding want-have should not override want-block
if ok && (e.WantType == pb.Message_Wantlist_Block || wantType == pb.Message_Wantlist_Have) {
return false
}
w.put(c, Entry{
Cid: c,
Priority: priority,
WantType: wantType,
})
return true
}
// Remove removes the given cid from the wantlist.
func (w *Wantlist) Remove(c cid.Cid) bool {
_, ok := w.set[c]
if !ok {
return false
}
w.delete(c)
return true
}
// Remove removes the given cid from the wantlist, respecting the type:
// Remove with want-have will not remove an existing want-block.
func (w *Wantlist) RemoveType(c cid.Cid, wantType pb.Message_Wantlist_WantType) bool {
e, ok := w.set[c]
if !ok {
return false
}
// Removing want-have should not remove want-block
if e.WantType == pb.Message_Wantlist_Block && wantType == pb.Message_Wantlist_Have {
return false
}
w.delete(c)
return true
}
func (w *Wantlist) delete(c cid.Cid) {
delete(w.set, c)
w.cached = nil
}
func (w *Wantlist) put(c cid.Cid, e Entry) {
w.cached = nil
w.set[c] = e
}
// Contains returns the entry, if present, for the given CID, plus whether it
// was present.
func (w *Wantlist) Contains(c cid.Cid) (Entry, bool) {
e, ok := w.set[c]
return e, ok
}
// Entries returns all wantlist entries for a want list, sorted by priority.
//
// DO NOT MODIFY. The returned list is cached.
func (w *Wantlist) Entries() []Entry {
if w.cached != nil {
return w.cached
}
es := make([]Entry, 0, len(w.set))
for _, e := range w.set {
es = append(es, e)
}
sort.Sort(entrySlice(es))
w.cached = es
return es[0:len(es):len(es)]
}
// Absorb all the entries in other into this want list
func (w *Wantlist) Absorb(other *Wantlist) {
// Invalidate the cache up-front to avoid doing any work trying to keep it up-to-date.
w.cached = nil
for _, e := range other.Entries() {
w.Add(e.Cid, e.Priority, e.WantType)
}
}