From 73f35ecad8e06864758758b9a8915a12816a7bb0 Mon Sep 17 00:00:00 2001
From: Jordan Orelli <jordanorelli@gmail.com>
Date: Mon, 1 Aug 2016 23:41:25 -0400
Subject: [PATCH] bit twiddling

---
 bit/reader.go      | 26 ++++++++++++++++
 bit/reader_test.go | 76 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 93 insertions(+), 9 deletions(-)

diff --git a/bit/reader.go b/bit/reader.go
index 3e39f99..06fe274 100644
--- a/bit/reader.go
+++ b/bit/reader.go
@@ -51,6 +51,9 @@ func (r *Reader) ReadBits(bits uint) (n uint64) {
 
 // ReadByte reads a single byte, regardless of alignment.
 func (r *Reader) ReadByte() (byte, error) {
+	if r.bits == 0 {
+		return r.src.ReadByte()
+	}
 	b := byte(r.ReadBits(8))
 	if err := r.Err(); err != nil {
 		return 0, err
@@ -70,4 +73,27 @@ func (r *Reader) Read(buf []byte) (int, error) {
 	return len(buf), nil
 }
 
+// ReadUbitVar reads a prefixed uint value. A prefix is 2 bits wide, followed
+// by the 4 least-significant bits, then a variable number of most-significant
+// bits based on the prefix.
+//
+// 00 - 0
+// 01 - 4
+// 10 - 8
+// 11 - 28
+func (r *Reader) ReadUBitVar() uint64 {
+	switch prefix := r.ReadBits(2); prefix {
+	case 0:
+		return r.ReadBits(4)
+	case 1:
+		return r.ReadBits(4) | r.ReadBits(4)<<4
+	case 2:
+		return r.ReadBits(4) | r.ReadBits(8)<<4
+	case 3:
+		return r.ReadBits(4) | r.ReadBits(28)<<4
+	default:
+		panic("not reached")
+	}
+}
+
 func (r *Reader) Err() error { return r.err }
diff --git a/bit/reader_test.go b/bit/reader_test.go
index 31398db..ab9bad1 100644
--- a/bit/reader_test.go
+++ b/bit/reader_test.go
@@ -9,13 +9,17 @@ import (
 var (
 	// 1000 1011 1010 1101 1111 0000 0000 1101
 	badFood = []byte{0x8b, 0xad, 0xf0, 0x0d}
+
+	// 0000 1110 0001 1110 1110 0111 1011 1110 1110 1111
+	eLeetBeef = []byte{0x0e, 0x1e, 0xe7, 0xbe, 0xef}
 )
 
 // test the bit-level reads
 func TestReadBits(t *testing.T) {
-	assert := assert.New(t)
-
-	var r *Reader
+	var (
+		assert = assert.New(t)
+		r      *Reader
+	)
 
 	// aligned reading
 	r = NewBytesReader(badFood)
@@ -44,29 +48,83 @@ func TestReadBits(t *testing.T) {
 	assert.Equal(uint64(0x15), r.ReadBits(5))
 
 	// 1000 1011 1010 1101 1111 0000 0000 1101
-	//                 ^---------^
+	//                 ^----+----^
 	assert.Equal(uint64(0x17c), r.ReadBits(9))
 
 	// 1000 1011 1010 1101 1111 0000 0000 1101
-	//                            ^----------^
+	//                            ^----+----+^
 	assert.Equal(uint64(0xd), r.ReadBits(10))
 }
 
 // test the Read calls, satisfying io.Reader
 func TestRead(t *testing.T) {
-	assert := assert.New(t)
-	var r *Reader
+	var (
+		assert = assert.New(t)
+		r      *Reader
+		buf    []byte
+		n      int
+		err    error
+	)
+
+	r = NewBytesReader(badFood)
+	buf = make([]byte, 4)
+	n, err = r.Read(buf)
+	assert.NoError(err)
+	assert.Equal(4, n)
+	assert.Equal(badFood, buf)
 
 	r = NewBytesReader(badFood)
 	// 1000 1011 1010 1101 1111 0000 0000 1101
 	r.ReadBits(1)
 	// 0001 0111 0101 1011 1110 0000 0001 101
-	buf := make([]byte, 3)
+	buf = make([]byte, 3)
 
-	n, err := r.Read(buf)
+	n, err = r.Read(buf)
 	assert.NoError(err)
 	assert.Equal(3, n)
 
 	expected := []byte{0x17, 0x5b, 0xe0}
 	assert.Equal(expected, buf)
 }
+
+func TestUbitVar(t *testing.T) {
+	var (
+		assert = assert.New(t)
+		r      *Reader
+		u      uint64
+	)
+
+	// 1000 1011 1010 1101 1111 0000 0000 1101
+	//  ||^---^                                : data (0101)
+	//  ^^                                     : prefix (00)
+	r = NewBytesReader(badFood)
+	r.ReadBits(1)
+	u = r.ReadUBitVar()
+	assert.Equal(uint64(5), u)
+
+	r.ReadBits(2)
+	// 1000 1011 1010 1101 1111 0000 0000 1101
+	//            |||   |  ^--^                : most significant (1111)
+	//            ||^---^                      : least significant (0110)
+	//            ^^                           : prefix (01)
+	u = r.ReadUBitVar()
+	assert.Equal(uint64(0xf6), u)
+
+	r = NewBytesReader(badFood)
+	// 1000 1011 1010 1101 1111 0000 0000 1101
+	// |||   |^----+---^                       : most significant (1110 1011)
+	// ||^---^                                 : least significant (0010)
+	// ^^                                      : prefix (10)
+	u = r.ReadUBitVar()
+	assert.Equal(uint64(0xeb2), u)
+
+	r = NewBytesReader(eLeetBeef)
+	r.ReadBits(4)
+	// 0000 1110 0001 1110 1110 0111 1011 1110 1110 1111
+	//      |||   |^----+----+----+----+----+----+---^   : msb (0111 1011 1001 1110
+	//      |||   |                                             1111 1011 1011)
+	//      ||^---^                                      : lsb (1000)
+	//      ^^                                           : prefix
+	u = r.ReadUBitVar()
+	assert.Equal(uint64(0x7b9efbb8), u)
+}