[refactor] package nazabits: BitWriter::WriteBits范围从8位增长到16位

pkg/nazabits/bits.go

@@ -10,6 +10,8 @@ package nazabits
 
 // TODO chef 这个package的性能可以优化
 
+// 所有的读写操作，由调用方保证不超出切片的范围
+
 type BitReader struct {
 	core  []byte
 	index uint
@@ -32,7 +34,8 @@ func (br *BitReader) ReadBit() uint8 {
 	return res
 }
 
-// 目前限制一次最多读8位，大于8位可以分多次读，或者以后把限制放开
+// 目前限制一次最多读8位，需调用方自己保证。如果需要读取大于8位，可以分多次读，或者以后把限制放开
+// @param n: 取值范围 [1, 8]
 func (br *BitReader) ReadBits(n uint) (r uint8) {
 	var i uint
 	for i = 0; i < n; i++ {
@@ -65,10 +68,14 @@ func (bw *BitWriter) WriteBit(b uint8) {
 	}
 }
 
-// 目前限制一次最多写8位，大于8位可以分多次写，或者以后把限制放开
-func (bw *BitWriter) WriteBits(n int, v uint8) {
-	for i := n - 1; i >= 0; i-- {
-		bw.WriteBit(v >> uint8(i))
+// @param n: 取值范围 [1, 16]
+// 目前限制一次最多写16位，大于16位可以分多次写，或者以后把限制放开
+func (bw *BitWriter) WriteBits(n uint, v uint16) {
+	for i := n - 1; ; i-- {
+		bw.WriteBit(uint8(v >> i & 0x1))
+		if i == 0 {
+			break
+		}
 	}
 }
 

pkg/nazabits/bits_test.go

@@ -148,6 +148,12 @@ func TestBitWriter_WriteBits(t *testing.T) {
 	bw = nazabits.NewBitWriter(v)
 	bw.WriteBits(3, 1+8+32+128)
 	assert.Equal(t, uint8(1<<5), v[0])
+
+	v = make([]byte, 2)
+	bw = nazabits.NewBitWriter(v)
+	bw.WriteBits(16, 0xFFFF)
+	assert.Equal(t, uint8(0xFF), v[0])
+	assert.Equal(t, uint8(0xFF), v[1])
 }
 
 func BenchmarkGetBits16(b *testing.B) {