From ba009bccde863ea1984263cd8009938179e82a81 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 5 Aug 2025 15:49:09 +1000
Subject: [PATCH] Add request-strategy.Btree.Contains

We'll need it later when we use piece request order instead of Torrent._pendingPieces.
---
 internal/request-strategy/ajwerner-btree.go      | 5 +++++
 internal/request-strategy/piece-request-order.go | 9 +++++++++
 internal/request-strategy/tidwall-btree.go       | 5 +++++
 3 files changed, 19 insertions(+)

diff --git a/internal/request-strategy/ajwerner-btree.go b/internal/request-strategy/ajwerner-btree.go
index 73ace303..81e2d6f3 100644
--- 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 a9cd07df..ccd647fe 100644
--- 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 88bbafd2..b3efe1c0 100644
--- 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
+}
-- 
2.51.0