From cc6614234d382c07eecd0849b85fc8bcaf2524e8 Mon Sep 17 00:00:00 2001
From: q191201771 <191201771@qq.com>
Date: Sun, 10 May 2020 06:52:00 +0800
Subject: [PATCH] [feat] package bits: new func BitReader::ReadBits16,
 ReadBits32, ReadBytes

---
 pkg/nazabits/bits.go      | 32 ++++++++++++++++++++++++++++----
 pkg/nazabits/bits_test.go | 25 +++++++++++++++++++++----
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/pkg/nazabits/bits.go b/pkg/nazabits/bits.go
index ef29e1e..1015a03 100644
--- a/pkg/nazabits/bits.go
+++ b/pkg/nazabits/bits.go
@@ -8,7 +8,10 @@
 
 package nazabits
 
-// TODO chef 这个package的性能可以优化
+// TODO
+// - 这个package的性能可以优化
+// - BitReader考虑增加skip接口
+// - BitReader开率增加返回目前位置接口
 
 // 所有的读写操作，由调用方保证不超出切片的范围
 
@@ -34,16 +37,37 @@ func (br *BitReader) ReadBit() uint8 {
 	return res
 }
 
-// 如果需要读取大于8位，可以分多次读，以后可能加上ReadBits16等接口
 // @param n: 取值范围 [1, 8]
 func (br *BitReader) ReadBits8(n uint) (r uint8) {
-	var i uint
-	for i = 0; i < n; i++ {
+	for i := uint(0); i < n; i++ {
 		r = (r << 1) | br.ReadBit()
 	}
 	return
 }
 
+// @param n: 取值范围 [1, 16]
+func (br *BitReader) ReadBits16(n uint) (r uint16) {
+	for i := uint(0); i < n; i++ {
+		r = (r << 1) | uint16(br.ReadBit())
+	}
+	return
+}
+
+func (br *BitReader) ReadBits32(n uint) (r uint32) {
+	for i := uint(0); i < n; i++ {
+		r = (r << 1) | uint32(br.ReadBit())
+	}
+	return
+}
+
+// @param n: 读取多少个字节
+func (br *BitReader) ReadBytes(n uint) (r []byte) {
+	for i := uint(0); i < n; i++ {
+		r = append(r, br.ReadBits8(8))
+	}
+	return
+}
+
 // ----------------------------------------------------------------------------
 
 type BitWriter struct {
diff --git a/pkg/nazabits/bits_test.go b/pkg/nazabits/bits_test.go
index 114dbc7..737722d 100644
--- a/pkg/nazabits/bits_test.go
+++ b/pkg/nazabits/bits_test.go
@@ -95,8 +95,6 @@ func TestBitReader_ReadBit(t *testing.T) {
 func TestBitReader_ReadBits8(t *testing.T) {
 	v := []byte{48, 57, 48, 57}
 	br := nazabits.NewBitReader(v)
-	// {0 01 1 00 00, 00 11 100 1, 0011 0 000, 0011 1001}
-	//  0,01,1 00,00 00, 11 100,100110,   0000011,
 	assert.Equal(t, uint8(0), br.ReadBits8(1))
 	assert.Equal(t, uint8(1), br.ReadBits8(2))
 	assert.Equal(t, uint8(4), br.ReadBits8(3))
@@ -108,8 +106,27 @@ func TestBitReader_ReadBits8(t *testing.T) {
 	br = nazabits.NewBitReader(v)
 	// {0011 0000, 0011 1001, 0 011 0000, 0011 1001}
 	assert.Equal(t, uint8(48), br.ReadBits8(8))
-	//assert.Equal(t, uint8(114), br.ReadBits8(9))
-	//assert.Equal(t, uint8(385), br.ReadBits8(10))
+}
+
+func TestBitReader_ReadBits16(t *testing.T) {
+	v := []byte{48, 57, 48, 57}
+	br := nazabits.NewBitReader(v)
+	assert.Equal(t, uint16(48), br.ReadBits16(8))
+	assert.Equal(t, uint16(1), br.ReadBits16(3))
+	assert.Equal(t, uint16(25), br.ReadBits16(5))
+	assert.Equal(t, uint16(12345), br.ReadBits16(16))
+}
+
+func TestBitReader_ReadBits32(t *testing.T) {
+	v := []byte{48, 57, 48, 57}
+	br := nazabits.NewBitReader(v)
+	assert.Equal(t, uint32(809054265), br.ReadBits32(32))
+}
+
+func TestBitReader_ReadBytes(t *testing.T) {
+	v := []byte{48, 57}
+	br := nazabits.NewBitReader(v)
+	assert.Equal(t, v, br.ReadBytes(2))
 }
 
 func TestBitWriter_WriteBit(t *testing.T) {