Use missinggo.bitmap for tracking dirty chunks

diff --git a/piece.go b/piece.go
index a56225c1ec7ca947ce97002676abc9460abc4377..9f22d718b9f918ec309ffc835f500aefa25112bf 100644 (file)
--- a/piece.go
+++ b/piece.go
@@ -4,6 +4,7 @@ import (
        "math/rand"
        "sync"
 
+       "github.com/anacrolix/missinggo/bitmap"
        "github.com/bradfitz/iter"
 
        pp "github.com/anacrolix/torrent/peer_protocol"
@@ -32,7 +33,7 @@ type piece struct {
        Hash pieceSum
        // Chunks we've written to since the last check. The chunk offset and
        // length can be determined by the request chunkSize in use.
-       DirtyChunks      []bool
+       DirtyChunks      bitmap.Bitmap
        Hashing          bool
        QueuedForHash    bool
        EverHashed       bool
@@ -44,44 +45,28 @@ type piece struct {
        noPendingWrites    sync.Cond
 }
 
+func (p *piece) pendingChunkIndex(chunkIndex int) bool {
+       return !p.DirtyChunks.Contains(chunkIndex)
+}
+
 func (p *piece) pendingChunk(cs chunkSpec, chunkSize pp.Integer) bool {
-       ci := chunkIndex(cs, chunkSize)
-       if ci >= len(p.DirtyChunks) {
-               return true
-       }
-       return !p.DirtyChunks[ci]
+       return p.pendingChunkIndex(chunkIndex(cs, chunkSize))
 }
 
 func (p *piece) hasDirtyChunks() bool {
-       for _, dirty := range p.DirtyChunks {
-               if dirty {
-                       return true
-               }
-       }
-       return false
+       return p.DirtyChunks.Len() != 0
 }
 
 func (p *piece) numDirtyChunks() (ret int) {
-       for _, dirty := range p.DirtyChunks {
-               if dirty {
-                       ret++
-               }
-       }
-       return
+       return p.DirtyChunks.Len()
 }
 
 func (p *piece) unpendChunkIndex(i int) {
-       for i >= len(p.DirtyChunks) {
-               p.DirtyChunks = append(p.DirtyChunks, false)
-       }
-       p.DirtyChunks[i] = true
+       p.DirtyChunks.Add(i)
 }
 
 func (p *piece) pendChunkIndex(i int) {
-       if i >= len(p.DirtyChunks) {
-               return
-       }
-       p.DirtyChunks[i] = false
+       p.DirtyChunks.Remove(i)
 }
 
 func chunkIndexSpec(index int, pieceLength, chunkSize pp.Integer) chunkSpec {
@@ -102,7 +87,7 @@ func (p *piece) shuffledPendingChunkSpecs(t *torrent, piece int) (css []chunkSpe
        }
        css = make([]chunkSpec, 0, numPending)
        for ci := range iter.N(t.pieceNumChunks(piece)) {
-               if ci >= len(p.DirtyChunks) || !p.DirtyChunks[ci] {
+               if !p.DirtyChunks.Contains(ci) {
                        css = append(css, t.chunkIndexSpec(ci, piece))
                }
        }

diff --git a/torrent.go b/torrent.go
index 987b5b8fe03e1443d8f51c31f2457a777f2269ae..f58a5d683b3766638a18cc0bbaa90855c077d069 100644 (file)
--- a/torrent.go
+++ b/torrent.go
@@ -38,11 +38,13 @@ func (t *torrent) pieceNumPendingBytes(index int) (count pp.Integer) {
        if !piece.EverHashed {
                return
        }
-       for i, dirty := range piece.DirtyChunks {
-               if dirty {
-                       count -= t.chunkIndexSpec(i, index).Length
-               }
+       regularDirty := piece.numDirtyChunks()
+       lastChunkIndex := t.pieceNumChunks(index) - 1
+       if piece.pendingChunkIndex(lastChunkIndex) {
+               regularDirty--
+               count -= t.chunkIndexSpec(lastChunkIndex, index).Length
        }
+       count -= pp.Integer(regularDirty) * t.chunkSize
        return
 }
 
@@ -633,7 +635,7 @@ func (t *torrent) pieceNumChunks(piece int) int {
 }
 
 func (t *torrent) pendAllChunkSpecs(pieceIndex int) {
-       t.Pieces[pieceIndex].DirtyChunks = nil
+       t.Pieces[pieceIndex].DirtyChunks.Clear()
 }
 
 type Peer struct {
@@ -826,16 +828,7 @@ func (t *torrent) pieceNumPendingChunks(piece int) int {
 }
 
 func (t *torrent) pieceAllDirty(piece int) bool {
-       p := &t.Pieces[piece]
-       if len(p.DirtyChunks) != t.pieceNumChunks(piece) {
-               return false
-       }
-       for _, dirty := range p.DirtyChunks {
-               if !dirty {
-                       return false
-               }
-       }
-       return true
+       return t.Pieces[piece].DirtyChunks.Len() == t.pieceNumChunks(piece)
 }
 
 func (t *torrent) forUrgentPieces(f func(piece int) (again bool)) (all bool) {