From 45d0473b5f1361907f2db4beec38b23b99f935bc Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 24 Apr 2025 23:08:38 +1000
Subject: [PATCH] Strongly type client piece request order keys

---
 client-piece-request-order.go  | 19 +++++++++++++++++++
 client.go                      |  2 +-
 torrent-piece-request-order.go |  7 ++++---
 3 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 client-piece-request-order.go

diff --git a/client-piece-request-order.go b/client-piece-request-order.go
new file mode 100644
index 00000000..787fc7ef
--- /dev/null
+++ b/client-piece-request-order.go
@@ -0,0 +1,19 @@
+package torrent
+
+import (
+	"github.com/anacrolix/torrent/storage"
+)
+
+type clientPieceRequestOrderKeyTypes interface {
+	storage.TorrentCapacity | *Torrent
+}
+
+type clientPieceRequestOrderKey[T clientPieceRequestOrderKeyTypes] struct {
+	inner T
+}
+
+func (me clientPieceRequestOrderKey[T]) isAClientPieceRequestOrderKeyType() {}
+
+type clientPieceRequestOrderKeySumType interface {
+	isAClientPieceRequestOrderKeyType()
+}
diff --git a/client.go b/client.go
index 0561b007..1d618f97 100644
--- a/client.go
+++ b/client.go
@@ -84,7 +84,7 @@ type Client struct {
 	// info has been obtained, there's no knowing if an infohash belongs to v1 or v2.
 	torrentsByShortHash map[InfoHash]*Torrent
 
-	pieceRequestOrder map[interface{}]*request_strategy.PieceRequestOrder
+	pieceRequestOrder map[clientPieceRequestOrderKeySumType]*request_strategy.PieceRequestOrder
 
 	acceptLimiter map[ipStr]int
 	numHalfOpen   int
diff --git a/torrent-piece-request-order.go b/torrent-piece-request-order.go
index b2bde8dc..02044f91 100644
--- a/torrent-piece-request-order.go
+++ b/torrent-piece-request-order.go
@@ -2,6 +2,7 @@ package torrent
 
 import (
 	g "github.com/anacrolix/generics"
+	"github.com/anacrolix/torrent/storage"
 
 	request_strategy "github.com/anacrolix/torrent/request-strategy"
 )
@@ -27,11 +28,11 @@ func (t *Torrent) updatePieceRequestOrderPiece(pieceIndex int) {
 	}
 }
 
-func (t *Torrent) clientPieceRequestOrderKey() interface{} {
+func (t *Torrent) clientPieceRequestOrderKey() clientPieceRequestOrderKeySumType {
 	if t.storage.Capacity == nil {
-		return t
+		return clientPieceRequestOrderKey[*Torrent]{t}
 	}
-	return t.storage.Capacity
+	return clientPieceRequestOrderKey[storage.TorrentCapacity]{t.storage.Capacity}
 }
 
 func (t *Torrent) deletePieceRequestOrder() {
-- 
2.50.0