network/bitvector: Multibit set/unset + string rep (#1530)

* network/bitvector: Multibit set/unset + string rep

* network/bitvector: Add code comments

* network/bitvector: Make Unset -> Set with bool false

* network/bitvector: Revert to Set/Unset

* network/stream: Update to new bitvector signature

Author: nolash

network/bitvector/bitvector.go

@@ -22,15 +22,20 @@ import (
 
 var errInvalidLength = errors.New("invalid length")
 
+// BitVector is a convenience object for manipulating and representing bit vectors
 type BitVector struct {
 	len int
 	b   []byte
 }
 
+// New creates a new bit vector with the given length
 func New(l int) (bv *BitVector, err error) {
 	return NewFromBytes(make([]byte, l/8+1), l)
 }
 
+// NewFromBytes creates a bit vector from the passed byte slice.
+//
+// Leftmost bit in byte slice becomes leftmost bit in bit vector
 func NewFromBytes(b []byte, l int) (bv *BitVector, err error) {
 	if l <= 0 {
 		return nil, errInvalidLength
@@ -44,19 +49,76 @@ func NewFromBytes(b []byte, l int) (bv *BitVector, err error) {
 	}, nil
 }
 
+// Get gets the corresponding bit, counted from left to right
 func (bv *BitVector) Get(i int) bool {
 	bi := i / 8
 	return bv.b[bi]&(0x1<<uint(i%8)) != 0
 }
 
-func (bv *BitVector) Set(i int, v bool) {
+// Set sets the bit corresponding to the index in the bitvector, counted from left to right
+func (bv *BitVector) set(i int, v bool) {
 	bi := i / 8
 	cv := bv.Get(i)
 	if cv != v {
 		bv.b[bi] ^= 0x1 << uint8(i%8)
 	}
 }
 
+// Set sets the bit corresponding to the index in the bitvector, counted from left to right
+func (bv *BitVector) Set(i int) {
+	bv.set(i, true)
+}
+
+// Unset UNSETS the corresponding bit, counted from left to right
+func (bv *BitVector) Unset(i int) {
+	bv.set(i, false)
+}
+
+// SetBytes sets all bits in the bitvector that are set in the argument
+//
+// The argument must be the same as the bitvector length
+func (bv *BitVector) SetBytes(bs []byte) error {
+	if len(bs) != bv.len {
+		return errors.New("invalid length")
+	}
+	for i := 0; i < bv.len*8; i++ {
+		bi := i / 8
+		if bs[bi]&(0x01<<uint(i%8)) > 0 {
+			bv.set(i, true)
+		}
+	}
+	return nil
+}
+
+// UnsetBytes UNSETS all bits in the bitvector that are set in the argument
+//
+// The argument must be the same as the bitvector length
+func (bv *BitVector) UnsetBytes(bs []byte) error {
+	if len(bs) != bv.len {
+		return errors.New("invalid length")
+	}
+	for i := 0; i < bv.len*8; i++ {
+		bi := i / 8
+		if bs[bi]&(0x01<<uint(i%8)) > 0 {
+			bv.set(i, false)
+		}
+	}
+	return nil
+}
+
+// String implements Stringer interface
+func (bv *BitVector) String() (s string) {
+	for i := 0; i < bv.len*8; i++ {
+		if bv.Get(i) {
+			s += "1"
+		} else {
+			s += "0"
+		}
+	}
+	return s
+}
+
+// Bytes retrieves the underlying bytes of the bitvector
 func (bv *BitVector) Bytes() []byte {
 	return bv.b
 }

network/bitvector/bitvector_test.go

@@ -18,6 +18,7 @@ package bitvector
 
 import "testing"
 
+// TestBitvectorNew checks that enforcements of argument length works in the constructors
 func TestBitvectorNew(t *testing.T) {
 	_, err := New(0)
 	if err != errInvalidLength {
@@ -40,6 +41,7 @@ func TestBitvectorNew(t *testing.T) {
 	}
 }
 
+// TestBitvectorGetSet tests correctness of individual Set and Get commands
 func TestBitvectorGetSet(t *testing.T) {
 	for _, length := range []int{
 		1,
@@ -71,7 +73,7 @@ func TestBitvectorGetSet(t *testing.T) {
 		}()
 
 		for i := 0; i < length; i++ {
-			bv.Set(i, true)
+			bv.Set(i)
 			for j := 0; j < length; j++ {
 				if j == i {
 					if !bv.Get(j) {
@@ -84,7 +86,7 @@ func TestBitvectorGetSet(t *testing.T) {
 				}
 			}
 
-			bv.Set(i, false)
+			bv.Unset(i)
 
 			if bv.Get(i) {
 				t.Errorf("element on index %v is not set to false", i)
@@ -93,6 +95,7 @@ func TestBitvectorGetSet(t *testing.T) {
 	}
 }
 
+// TestBitvectorNewFromBytesGet tests that bit vector is initialized correctly from underlying byte slice
 func TestBitvectorNewFromBytesGet(t *testing.T) {
 	bv, err := NewFromBytes([]byte{8}, 8)
 	if err != nil {
@@ -102,3 +105,36 @@ func TestBitvectorNewFromBytesGet(t *testing.T) {
 		t.Fatalf("element 3 is not set to true: state %08b", bv.b[0])
 	}
 }
+
+// TestBitVectorString tests that string representation of bit vector is correct
+func TestBitVectorString(t *testing.T) {
+	b := []byte{0xa5, 0x81}
+	expect := "1010010110000001"
+	bv, err := NewFromBytes(b, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if bv.String() != expect {
+		t.Fatalf("bitvector string fail: got %s, expect %s", bv.String(), expect)
+	}
+}
+
+// TestBitVectorSetUnsetBytes tests that setting and unsetting by byte slice modifies the bit vector correctly
+func TestBitVectorSetBytes(t *testing.T) {
+	b := []byte{0xff, 0xff}
+	cb := []byte{0xa5, 0x81}
+	expectUnset := "0101101001111110"
+	expectReset := "1111111111111111"
+	bv, err := NewFromBytes(b, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	bv.UnsetBytes(cb)
+	if bv.String() != expectUnset {
+		t.Fatalf("bitvector unset bytes fail: got %s, expect %s", bv.String(), expectUnset)
+	}
+	bv.SetBytes(cb)
+	if bv.String() != expectReset {
+		t.Fatalf("bitvector reset bytes fail: got %s, expect %s", bv.String(), expectReset)
+	}
+}

network/stream/messages.go

@@ -231,7 +231,7 @@ func (p *Peer) handleOfferedHashesMsg(ctx context.Context, req *OfferedHashesMsg
 			ctr++
 
 			// set the bit, so create a request
-			want.Set(i/HashSize, true)
+			want.Set(i / HashSize)
 			log.Trace("need data", "ref", fmt.Sprintf("%x", hash), "request", true)
 
 			// measure how long it takes before we mark chunks for retrieval, and actually send the request