From da1221dd50432c072a5aca0f1ddab9f329bc63ad Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 5 Oct 2021 17:48:34 +1100
Subject: [PATCH] Change Peer._peerPieces to use raw roaring Bitmap type

The wrapper type was from when roaring didn't support zero-alloc initialization.
---
 peerconn.go           | 22 ++++++++++++++--------
 test/transfer_test.go |  2 +-
 torrent.go            |  4 ++--
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/peerconn.go b/peerconn.go
index 6851d1d6..2c309bd7 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -13,6 +13,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/RoaringBitmap/roaring"
 	"github.com/anacrolix/log"
 	"github.com/anacrolix/missinggo/iter"
 	"github.com/anacrolix/missinggo/v2/bitmap"
@@ -110,7 +111,7 @@ type Peer struct {
 	PeerPrefersEncryption bool // as indicated by 'e' field in extension handshake
 	PeerListenPort        int
 	// The pieces the peer has claimed to have.
-	_peerPieces bitmap.Bitmap
+	_peerPieces roaring.Bitmap
 	// The peer has everything. This can occur due to a special message, when
 	// we may not even know the number of pieces in the torrent yet.
 	peerSentHaveAll bool
@@ -238,7 +239,7 @@ func (cn *Peer) peerHasAllPieces() (all bool, known bool) {
 	if !cn.t.haveInfo() {
 		return false, false
 	}
-	return bitmap.Flip(cn._peerPieces, 0, bitmap.BitRange(cn.t.numPieces())).IsEmpty(), true
+	return roaring.Flip(&cn._peerPieces, 0, bitmap.BitRange(cn.t.numPieces())).IsEmpty(), true
 }
 
 func (cn *PeerConn) locker() *lockWithDeferreds {
@@ -259,7 +260,7 @@ func (cn *Peer) bestPeerNumPieces() pieceIndex {
 }
 
 func (cn *Peer) completedString() string {
-	have := pieceIndex(cn._peerPieces.Len())
+	have := pieceIndex(cn._peerPieces.GetCardinality())
 	if cn.peerSentHaveAll {
 		have = cn.bestPeerNumPieces()
 	}
@@ -758,7 +759,7 @@ func (cn *PeerConn) peerSentHave(piece pieceIndex) error {
 	if !cn.peerHasPiece(piece) {
 		cn.t.incPieceAvailability(piece)
 	}
-	cn._peerPieces.Set(bitmap.BitIndex(piece), true)
+	cn._peerPieces.Add(uint32(piece))
 	cn.t.maybeDropMutuallyCompletePeer(&cn.Peer)
 	if cn.updatePiecePriority(piece) {
 		cn.updateRequests()
@@ -789,7 +790,11 @@ func (cn *PeerConn) peerSentBitfield(bf []bool) error {
 				cn.t.decPieceAvailability(i)
 			}
 		}
-		cn._peerPieces.Set(bitmap.BitIndex(i), have)
+		if have {
+			cn._peerPieces.Add(uint32(i))
+		} else {
+			cn._peerPieces.Remove(uint32(i))
+		}
 	}
 	cn.peerPiecesChanged()
 	return nil
@@ -1596,15 +1601,16 @@ func (l connectionTrust) Less(r connectionTrust) bool {
 
 // Returns the pieces the peer could have based on their claims. If we don't know how many pieces
 // are in the torrent, it could be a very large range the peer has sent HaveAll.
-func (cn *PeerConn) PeerPieces() bitmap.Bitmap {
+func (cn *PeerConn) PeerPieces() *roaring.Bitmap {
 	cn.locker().RLock()
 	defer cn.locker().RUnlock()
 	return cn.newPeerPieces()
 }
 
 // Returns a new Bitmap that includes bits for all pieces the peer could have based on their claims.
-func (cn *Peer) newPeerPieces() bitmap.Bitmap {
-	ret := cn._peerPieces.Copy()
+func (cn *Peer) newPeerPieces() *roaring.Bitmap {
+	// TODO: Can we use copy on write?
+	ret := cn._peerPieces.Clone()
 	if cn.peerSentHaveAll {
 		if cn.t.haveInfo() {
 			ret.AddRange(0, bitmap.BitRange(cn.t.numPieces()))
diff --git a/test/transfer_test.go b/test/transfer_test.go
index 71de58b4..b70cb46c 100644
--- a/test/transfer_test.go
+++ b/test/transfer_test.go
@@ -168,7 +168,7 @@ func testClientTransfer(t *testing.T, ps testClientTransferParams) {
 	}
 	foundSeeder := false
 	for _, pc := range leecherPeerConns {
-		completed := pc.PeerPieces().Len()
+		completed := pc.PeerPieces().GetCardinality()
 		t.Logf("peer conn %v has %v completed pieces", pc, completed)
 		if completed == bitmap.BitRange(leecherTorrent.Info().NumPieces()) {
 			foundSeeder = true
diff --git a/torrent.go b/torrent.go
index 3981651a..d184afd3 100644
--- a/torrent.go
+++ b/torrent.go
@@ -1408,8 +1408,8 @@ func (t *Torrent) decPeerPieceAvailability(p *Peer) {
 	if !t.haveInfo() {
 		return
 	}
-	p.newPeerPieces().IterTyped(func(i int) bool {
-		p.t.decPieceAvailability(i)
+	p.newPeerPieces().Iterate(func(i uint32) bool {
+		p.t.decPieceAvailability(pieceIndex(i))
 		return true
 	})
 }
-- 
2.44.0