Add request-strategy.Btree.Contains

We'll need it later when we use piece request order instead of Torrent._pendingPieces.

diff --git a/internal/request-strategy/ajwerner-btree.go b/internal/request-strategy/ajwerner-btree.go
index 73ace30370253d7e3ebc562a40fbd034c3fed54d..81e2d6f3d8c3e40927f61d2503cd84c8483268f7 100644 (file)
--- a/internal/request-strategy/ajwerner-btree.go
+++ b/internal/request-strategy/ajwerner-btree.go
@@ -8,6 +8,11 @@ type ajwernerBtree struct {
        btree btree.Set[PieceRequestOrderItem]
 }
 
+func (a *ajwernerBtree) Contains(item PieceRequestOrderItem) bool {
+       _, ok := a.btree.Get(item)
+       return ok
+}
+
 var _ Btree = (*ajwernerBtree)(nil)
 
 func NewAjwernerBtree() *ajwernerBtree {

diff --git a/internal/request-strategy/piece-request-order.go b/internal/request-strategy/piece-request-order.go
index a9cd07df204d071763c633c4b088350d20cf4104..ccd647fe3fb0b2d7aebb44a8c4d2d074fd8d5eda 100644 (file)
--- a/internal/request-strategy/piece-request-order.go
+++ b/internal/request-strategy/piece-request-order.go
@@ -5,6 +5,7 @@ import (
        "unique"
 
        g "github.com/anacrolix/generics"
+       "github.com/anacrolix/missinggo/v2/panicif"
 
        "github.com/anacrolix/torrent/metainfo"
 )
@@ -12,7 +13,9 @@ import (
 type Btree interface {
        Delete(PieceRequestOrderItem)
        Add(PieceRequestOrderItem)
+       // TODO: Add an iterator variant of this and benchmark.
        Scan(func(PieceRequestOrderItem) bool)
+       Contains(PieceRequestOrderItem) bool
 }
 
 func NewPieceOrder(btree Btree, cap int) *PieceRequestOrder {
@@ -95,3 +98,9 @@ func (me *PieceRequestOrder) Iter() iter.Seq[PieceRequestOrderItem] {
                })
        }
 }
+
+func (me *PieceRequestOrder) Get(key PieceRequestOrderKey) (ret g.Option[PieceRequestOrderState]) {
+       ret.Value, ret.Ok = me.keys[key]
+       panicif.NotEq(ret.Ok, me.tree.Contains(PieceRequestOrderItem{key, ret.Value}))
+       return
+}

diff --git a/internal/request-strategy/tidwall-btree.go b/internal/request-strategy/tidwall-btree.go
index 88bbafd2a2dfdad42ad7be4935def1aa88d469bc..b3efe1c0d74b1969994deffd05d6b8a784784e74 100644 (file)
--- a/internal/request-strategy/tidwall-btree.go
+++ b/internal/request-strategy/tidwall-btree.go
@@ -35,3 +35,8 @@ func (me *tidwallBtree) Delete(item PieceRequestOrderItem) {
        _, deleted := me.tree.DeleteHint(item, me.PathHint)
        mustValue(deleted, item)
 }
+
+func (me *tidwallBtree) Contains(item PieceRequestOrderItem) bool {
+       _, ok := me.tree.Get(item)
+       return ok
+}