refactor: impl piece state with bitmap

protocol/bt/peer/message/msg_bitfield.go

@@ -1,15 +1,18 @@
 package message
 
+import "github.com/RoaringBitmap/roaring"
+
 // bitfield: <len=0001+X><id=5><bitfield>
 type Bitfield struct {
 	*Message
-	payload []byte
+	payload *roaring.Bitmap
 }
 
-func NewBitfield(payload []byte) *Bitfield {
+func NewBitfield(payload *roaring.Bitmap) *Bitfield {
 	return &Bitfield{
 		Message: &Message{
-			Length: uint32(len(payload) + 1),
+			// TODO 长度处理
+			Length: 1,
 			ID:     IdBitfield,
 		},
 		payload: payload,
@@ -17,44 +20,46 @@ func NewBitfield(payload []byte) *Bitfield {
 }
 
 func (b *Bitfield) Encode() []byte {
-	return append(b.Message.Encode(), b.payload...)
+	buf := b.Message.Encode()
+	for i := 0; i < int(b.Length)-len(buf); i++ {
+		temp := byte(0)
+		for j := 0; j < 8; j++ {
+			// 如果piece存在,位数标识为1
+			if b.payload.ContainsInt(i*8 + j) {
+				temp = temp | (1 << (7 - j))
+			}
+		}
+		buf = append(buf, temp)
+	}
+	return buf
 }
 
 func (b *Bitfield) Decode(buf []byte) Serialize {
 	b.Message = &Message{}
 	b.Message.Decode(buf)
-	b.payload = buf[5:]
+	temp := buf[5:]
+	b.payload = roaring.New()
+	// 初始化bitmap
+	for i := 0; i < len(temp); i++ {
+		for j := 0; j < 8; j++ {
+			if temp[i]&(1<<(7-j)) > 0 {
+				b.payload.AddInt(i*8 + j)
+			}
+		}
+	}
 	return b
 }
 
 // 某个分片是否下载完成
 func (b *Bitfield) IsComplete(i int) bool {
-	index := i / 8
-	if index >= len(b.payload) {
-		return false
-	}
-	return b.payload[index]&(1<<(7-i%8)) > 0
+	return b.payload.ContainsInt(i)
 }
 
-// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标
-func (b *Bitfield) Have(pieces []bool) []int {
-	arr := make([]int, 0)
-	length := len(pieces) / 8
-	if len(pieces)%8 != 0 {
-		length++
-	}
-
-	for i := 0; i < length; i++ {
-		for j := 0; j < 8; j++ {
-			index := i*8 + j
-			// 如果此分片在本地还未下载，检查peer是否能提供该分片下载
-			if index < len(pieces) &&
-				!pieces[index] &&
-				i < len(b.payload) &&
-				b.payload[i]&(1<<(7-j)) > 0 {
-				arr = append(arr, index)
-			}
-		}
-	}
-	return arr
+// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标，如果已经拥有则不用提供
+// had = 0b10010000,has = 0b01111111
+// had & has = 0b10010000 & 0b01111111 = 0b00010000
+// had & has ^ has = 0b00010000 ^ 0b01111111 = 0b01101111
+func (b *Bitfield) Provide(had *roaring.Bitmap) []uint32 {
+	xor := roaring.Xor(roaring.And(had, b.payload), b.payload)
+	return xor.ToArray()
 }

protocol/bt/peer/message/msg_bitfield_test.go

@@ -1,6 +1,7 @@
 package message
 
 import (
+	"github.com/RoaringBitmap/roaring"
 	"reflect"
 	"testing"
 )
@@ -8,7 +9,7 @@ import (
 func TestBitfield_IsComplete(t *testing.T) {
 	type fields struct {
 		Message Message
-		payload []byte
+		payload *roaring.Bitmap
 	}
 	type args struct {
 		i int
@@ -19,20 +20,10 @@ func TestBitfield_IsComplete(t *testing.T) {
 		args   args
 		want   bool
 	}{
-		{
-			"case-0",
-			fields{
-				payload: []byte{0b10000000},
-			},
-			args{
-				0,
-			},
-			true,
-		},
 		{
 			"case-1",
 			fields{
-				payload: []byte{0b01000000},
+				payload: roaring.BitmapOf(1, 2, 3, 4),
 			},
 			args{
 				1,
@@ -42,40 +33,10 @@ func TestBitfield_IsComplete(t *testing.T) {
 		{
 			"case-2",
 			fields{
-				payload: []byte{0b00000000, 0b00100000},
+				payload: roaring.BitmapOf(1, 2, 3, 4),
 			},
 			args{
-				2,
-			},
-			false,
-		},
-		{
-			"case-10",
-			fields{
-				payload: []byte{0b00000000, 0b00100000},
-			},
-			args{
-				10,
-			},
-			true,
-		},
-		{
-			"case-11",
-			fields{
-				payload: []byte{0b00000000, 0b00100000},
-			},
-			args{
-				11,
-			},
-			false,
-		},
-		{
-			"case-20",
-			fields{
-				payload: []byte{0b00000000, 0b00100000},
-			},
-			args{
-				20,
+				5,
 			},
 			false,
 		},
@@ -96,56 +57,46 @@ func TestBitfield_IsComplete(t *testing.T) {
 func TestBitfield_Have(t *testing.T) {
 	type fields struct {
 		Message Message
-		payload []byte
+		payload *roaring.Bitmap
 	}
 	type args struct {
-		pieces []bool
+		had *roaring.Bitmap
 	}
 	tests := []struct {
 		name   string
 		fields fields
 		args   args
-		want   []int
+		want   []uint32
 	}{
 		{
 			"case-1",
 			fields{
-				payload: []byte{0b11111111, 0b00000000},
+				payload: roaring.BitmapOf(0, 1, 2, 3, 4),
 			},
 			args{
-				[]bool{false, false, false, false, false, false, false, false},
+				roaring.BitmapOf(0),
 			},
-			[]int{0, 1, 2, 3, 4, 5, 6, 7},
+			[]uint32{1, 2, 3, 4},
 		},
 		{
 			"case-2",
 			fields{
-				payload: []byte{0b01010101, 0b01000000},
+				payload: roaring.BitmapOf(0, 1, 2, 3, 4, 5, 6, 7, 8),
 			},
 			args{
-				[]bool{false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false},
+				roaring.BitmapOf(0, 1, 3, 5, 6, 8),
 			},
-			[]int{1, 3, 5, 7, 9},
+			[]uint32{2, 4, 7},
 		},
 		{
 			"case-3",
 			fields{
-				payload: []byte{0b11010101, 0b01000000},
+				payload: roaring.BitmapOf(),
 			},
 			args{
-				[]bool{false, true, false},
+				roaring.BitmapOf(0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
 			},
-			[]int{0},
-		},
-		{
-			"case-4",
-			fields{
-				payload: []byte{0b11010101},
-			},
-			args{
-				[]bool{false, true, false, false, true, false, false, true, false},
-			},
-			[]int{0, 3, 5},
+			[]uint32{},
 		},
 	}
 	for _, tt := range tests {
@@ -154,8 +105,8 @@ func TestBitfield_Have(t *testing.T) {
 				Message: &tt.fields.Message,
 				payload: tt.fields.payload,
 			}
-			if got := b.Have(tt.args.pieces); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("Bitfield.Have() = %v, want %v", got, tt.want)
+			if got := b.Provide(tt.args.had); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Bitfield.Provide() = %v, want %v", got, tt.want)
 			}
 		})
 	}