2
0
mirror of synced 2025-02-24 14:48:27 +00:00
torrent/internal/pieceordering/pieceordering.go

117 lines
2.0 KiB
Go

package pieceordering
import (
"math/rand"
"github.com/ryszard/goskiplist/skiplist"
)
type Instance struct {
sl *skiplist.SkipList
pieceKeys map[int]int
}
func New() *Instance {
return &Instance{
sl: skiplist.NewIntMap(),
}
}
// Add the piece with the given key. No other piece can have the same key. If
// the piece is already present, change its key.
func (me *Instance) SetPiece(piece, key int) {
if existingKey, ok := me.pieceKeys[piece]; ok {
if existingKey == key {
return
}
me.removeKeyPiece(existingKey, piece)
}
var itemSl []int
if exItem, ok := me.sl.Get(key); ok {
itemSl = exItem.([]int)
}
me.sl.Set(key, append(itemSl, piece))
if me.pieceKeys == nil {
me.pieceKeys = make(map[int]int)
}
me.pieceKeys[piece] = key
me.shuffleItem(key)
}
func (me *Instance) shuffleItem(key int) {
_item, ok := me.sl.Get(key)
if !ok {
return
}
item := _item.([]int)
for i := range item {
j := i + rand.Intn(len(item)-i)
item[i], item[j] = item[j], item[i]
}
me.sl.Set(key, item)
}
func (me *Instance) removeKeyPiece(key, piece int) {
item, ok := me.sl.Get(key)
if !ok {
panic("no item for key")
}
itemSl := item.([]int)
for i, piece1 := range itemSl {
if piece1 == piece {
itemSl[i] = itemSl[len(itemSl)-1]
itemSl = itemSl[:len(itemSl)-1]
break
}
}
if len(itemSl) == 0 {
me.sl.Delete(key)
} else {
me.sl.Set(key, itemSl)
}
}
func (me *Instance) DeletePiece(piece int) {
key, ok := me.pieceKeys[piece]
if !ok {
return
}
me.removeKeyPiece(key, piece)
delete(me.pieceKeys, piece)
}
func (me Instance) First() Element {
i := me.sl.SeekToFirst()
if i == nil {
return nil
}
return &element{i, i.Value().([]int)}
}
type Element interface {
Piece() int
Next() Element
}
type element struct {
i skiplist.Iterator
sl []int
}
func (e *element) Next() Element {
e.sl = e.sl[1:]
if len(e.sl) > 0 {
return e
}
ok := e.i.Next()
if !ok {
return nil
}
e.sl = e.i.Value().([]int)
return e
}
func (e element) Piece() int {
return e.sl[0]
}