Change Peer._peerPieces to use raw roaring Bitmap type

[btrtrc.git] / peerconn.go

diff --git a/peerconn.go b/peerconn.go
index 6851d1d602b62b887449121e570529826006f728..2c309bd7bb9bbd3b523a0895a8a464b0ed8ca6f2 100644 (file)
--- 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()))