2 files changed, 313 insertions, 112 deletions

diff --git a/bitreader.go b/bitreader.go
index 06374b8..3fbc221 100644
--- a/bitreader.go
+++ b/bitreader.go
@@ -1,51 +1,146 @@
+// BitReader is a simple bit reader with big/little-endian support for golang.
+// It can read stream data from an io.Reader; can read from os.File and a byte array with bytes.NewReader(array).
+// Uses bitwise operations for v2.
+// Supports reading up to 64 bits at one time.
+// Includes wrapper functions for most used data types.
+// Error checking on all but wrapper functions.
+// Thanks to github.com/mlugg for the big help!
 package bitreader
 
 import (
         "fmt"
-        "math/bits"
-        "strconv"
-        "strings"
+        "io"
+        "math"
 )
 
+// ReaderType is the main structure of our Reader.
+// Whenever index == 0, we need to read a new byte from stream into curByte
 type ReaderType struct {
-        data  []byte // Reader data from a byte array
-        base  int    // Current reader byte location
-        index int    // Current reader index
-        le    bool   // Little endian or big endian?
+        stream  io.Reader // The underlying stream we're reading bytes from
+        index   uint8     // The current index into the byte [0-7]
+        curByte byte      // The byte we're currently reading from
+        le      bool      // Whether to read in little-endian order
 }
 
-func Reader(data []byte, le bool) *ReaderType {
-        dataClone := data
-        if le {
-                for index, byteValue := range data {
-                        dataClone[index] = bits.Reverse8(byteValue)
-                }
-        }
+// Reader is the main constructor that creates the ReaderType object
+// with stream data and little-endian state.
+func Reader(stream io.Reader, le bool) *ReaderType {
         return &ReaderType{
-                data:  dataClone,
-                base:  0,
-                index: 0,
-                le:    le,
+                stream:  stream,
+                index:   0,
+                curByte: 0, // Initial value doesn't matter, it'll be read as soon as we try to read any bits
+                le:      le,
         }
 }
 
-func (reader *ReaderType) SkipBits(bits int) error {
-        if bits <= 0 {
-                return fmt.Errorf("SkipBits Error: Bits value %d lower or equals than 0.", bits)
+// TryReadBool is a wrapper function that gets the state of 1-bit,
+// returns true if 1, false if 0. Panics on error.
+func (reader *ReaderType) TryReadBool() bool {
+        flag, err := reader.ReadBool()
+        if err != nil {
+                panic(err)
+        }
+        return flag
+}
+
+// TryReadInt1 is a wrapper function that returns the value of 1-bit.
+// Returns type uint8. Panics on error.
+func (reader *ReaderType) TryReadInt1() uint8 {
+        value, err := reader.ReadBits(1)
+        if err != nil {
+                panic(err)
+        }
+        return uint8(value)
+}
+
+// TryReadInt8 is a wrapper function that returns the value of 8-bits.
+// Returns uint8. Panics on error.
+func (reader *ReaderType) TryReadInt8() uint8 {
+        value, err := reader.ReadBits(8)
+        if err != nil {
+                panic(err)
+        }
+        return uint8(value)
+}
+
+// TryReadInt16 is a wrapper function that returns the value of 16-bits.
+// Returns uint16. Panics on error.
+func (reader *ReaderType) TryReadInt16() uint16 {
+        value, err := reader.ReadBits(16)
+        if err != nil {
+                panic(err)
+        }
+        return uint16(value)
+}
+
+// TryReadInt32 is a wrapper function that returns the value of 32-bits.
+// Returns uint32. Panics on error.
+func (reader *ReaderType) TryReadInt32() uint32 {
+        value, err := reader.ReadBits(32)
+        if err != nil {
+                panic(err)
+        }
+        return uint32(value)
+}
+
+// TryReadInt64 is a wrapper function that returns the value of 64-bits.
+// Returns uint64. Panics on error.
+func (reader *ReaderType) TryReadInt64() uint64 {
+        value, err := reader.ReadBits(64)
+        if err != nil {
+                panic(err)
+        }
+        return value
+}
+
+// TryReadFloat32 is a wrapper function that returns the value of 32-bits.
+// Returns float32. Panics on error.
+func (reader *ReaderType) TryReadFloat32() float32 {
+        value, err := reader.ReadBits(32)
+        if err != nil {
+                panic(err)
+        }
+        return math.Float32frombits(uint32(value))
+}
+
+// TryReadFloat64 is a wrapper function that returns the value of 64-bits.
+// Returns float64. Panics on error.
+func (reader *ReaderType) TryReadFloat64() float64 {
+        value, err := reader.ReadBits(64)
+        if err != nil {
+                panic(err)
         }
-        err := reader.checkAvailableBits(bits)
+        return math.Float64frombits(value)
+}
+
+// SkipBits is a function that increases Reader index
+// based on given input bits number. Returns an error
+// if there are no remaining bits.
+func (reader *ReaderType) SkipBits(bits int) error {
+        // Read as many raw bytes as we can
+        bytes := bits / 8
+        buf := make([]byte, bytes)
+        _, err := reader.stream.Read(buf)
         if err != nil {
                 return err
         }
-        for reader.index+bits > 7 {
-                reader.base++
-                reader.index = 0
-                bits -= 8
+        // The final read byte should be the new current byte
+        if bytes > 0 {
+                reader.curByte = buf[bytes-1]
+        }
+        // Read the extra bits
+        for i := bytes * 8; i < bits; i++ {
+                _, err := reader.readBit()
+                if err != nil {
+                        return err
+                }
         }
-        reader.index += bits
         return nil
 }
 
+// SkipBytes is a function that increases Reader index
+// based on given input bytes number. Returns an error
+// if there are no remaining bits.
 func (reader *ReaderType) SkipBytes(bytes int) error {
         err := reader.SkipBits(bytes * 8)
         if err != nil {
@@ -54,81 +149,83 @@ func (reader *ReaderType) SkipBytes(bytes int) error {
         return nil
 }
 
-func (reader *ReaderType) ReadBits(bits int) (int, error) {
-        if bits <= 0 {
-                return -1, fmt.Errorf("ReadBits Error: Bits value %d lower or equals than 0.", bits)
-        }
-        if bits > 64 {
-                return -1, fmt.Errorf("ReadBits Error: Bits value %d higher than 64.", bits)
+// ReadBits is a function that reads the specified amount of bits
+// specified in the parameter and returns the value, error
+// based on the output. It can read up to 64 bits. Returns the read
+// value in type uint64.
+//
+// Returns an error if there are no remaining bits.
+func (reader *ReaderType) ReadBits(bits int) (uint64, error) {
+        if bits < 1 || bits > 64 {
+                return 0, fmt.Errorf("ReadBits(bits) ERROR: Bits number should be between 1 and 64.")
         }
-        err := reader.checkAvailableBits(bits)
-        if err != nil {
-                return -1, err
-        }
-        if reader.le {
-                var output string
-                // Go to last bit and read backwards from there
-                reader.base += bits / 8
-                reader.index += bits % 8
-                if reader.index > 7 {
-                        reader.index -= 8
-                        reader.base++
-                }
-                for i := 0; i < bits; i++ {
-                        reader.index--
-                        if reader.index < 0 {
-                                reader.base--
-                                reader.index = 7
-                        }
-                        binary := fmt.Sprintf("%08b", reader.data[reader.base])
-                        binaryArr := strings.Split(binary, "")
-                        output += binaryArr[reader.index]
-                }
-                // Return to last bit after reading
-                reader.base += bits / 8
-                reader.index += bits % 8
-                if reader.index > 7 {
-                        reader.index -= 8
-                }
-                // Conversion of string binary to int
-                value, err := strconv.ParseUint(output, 2, 64)
+        var val uint64
+        for i := 0; i < bits; i++ {
+                bit, err := reader.readBit()
                 if err != nil {
-                        return -1, fmt.Errorf("%s", err)
-                }
-                return int(value), nil
-        } else {
-                var output string
-                for i := 0; i < bits; i++ {
-                        binary := fmt.Sprintf("%08b", reader.data[reader.base])
-                        binaryArr := strings.Split(binary, "")
-                        output += binaryArr[reader.index]
-                        reader.index++
-                        if reader.index > 7 {
-                                reader.base++
-                                reader.index = 0
-                        }
+                        return 0, err
                 }
-                // Conversion of string binary to int
-                value, err := strconv.ParseUint(output, 2, 64)
-                if err != nil {
-                        return -1, fmt.Errorf("%s", err)
+
+                if reader.le {
+                        val |= uint64(bit) << i
+                } else {
+                        val |= uint64(bit) << (bits - 1 - i)
                 }
-                return int(value), nil
         }
+        return val, nil
 }
 
-func (reader *ReaderType) ReadBit() (bool, error) {
-        value, err := reader.ReadBits(1)
+// ReadBytes is a function that reads the specified amount of bytes
+// specified in the parameter and returns the value, error
+// based on the output. It can read up to 8 bytes. Returns the read
+// value in type uint64.
+//
+// Returns an error if there are no remaining bits.
+func (reader *ReaderType) ReadBytes(bytes int) (uint64, error) {
+        if bytes < 1 || bytes > 8 {
+                return 0, fmt.Errorf("ReadBytes(bytes) ERROR: Bytes number should be between 1 and 8.")
+        }
+        value, err := reader.ReadBits(bytes * 8)
         if err != nil {
-                return false, fmt.Errorf("ReadBit Error: %s", err)
+                return 0, err
         }
-        return value != 0, nil
+        return value, nil
 }
 
-func (reader *ReaderType) checkAvailableBits(bits int) error {
-        availableBits := (len(reader.data)-reader.base)*8 - reader.index
-        if availableBits < bits {
-                return fmt.Errorf("BitReaderOutOfBounds: Wanted to read/skip %d bit(s) but only %d bit(s) is/are available.", bits, availableBits)
+// ReadBool is a function that reads one bit and returns the state, error
+// based on the output. Returns the read value in a bool format.
+//
+// Returns an error if there are no remaining bits.
+func (reader *ReaderType) ReadBool() (bool, error) {
+        val, err := reader.readBit()
+        if err != nil {
+                return false, err
+        }
+        return val == 1, nil
+}
+
+// readBit is a private function that reads a single bit from the stream.
+// This is the main function that makes us read stream data.
+func (reader *ReaderType) readBit() (uint8, error) {
+        if reader.index == 0 {
+                // Read a byte from stream into curByte
+                buf := make([]byte, 1)
+                _, err := reader.stream.Read(buf)
+                if err != nil {
+                        return 0, err
+                }
+                reader.curByte = buf[0]
+        }
+        var val bool
+        if reader.le {
+                val = (reader.curByte & (1 << reader.index)) != 0
+        } else {
+                val = (reader.curByte & (1 << (7 - reader.index))) != 0
+        }
+        reader.index = (reader.index + 1) % 8
+        if val {
+                return 1, nil
+        } else {
+                return 0, nil
         }
-        return nil
 }
diff --git a/bitreader_test.go b/bitreader_test.go
index 5d344d2..85b0c98 100644
--- a/bitreader_test.go
+++ b/bitreader_test.go
@@ -1,17 +1,84 @@
 package bitreader
 
 import (
+        "bytes"
         "testing"
 )
 
 // 01110001, 00001101, 00000000, 00000000, 10100010, 00011011, 00000000, 00000000, 11001100
 var TestArray = [...]byte{113, 13, 0, 0, 162, 27, 0, 0, 204}
 
+func TestTryReadFloat32(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []float32{6.98198182157e+29, -2.10064170919e-18}
+        for i := range expected {
+                value := bitreader.TryReadFloat32()
+                if value != expected[i] {
+                        t.Fatalf("TryReadFloat32 FAIL for index %d: Expected %f, Got %f", i, expected[i], value)
+                }
+        }
+}
+
+func TestTryReadFloat64(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []float64{3.68828741038253948851462939603e+236}
+        for i := range expected {
+                value := bitreader.TryReadFloat64()
+                if value != expected[i] {
+                        t.Fatalf("TryReadFloat64 FAIL for index %d: Expected %f, Got %f", i, expected[i], value)
+                }
+        }
+}
+
+func TestTryReadInt8(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []int{113, 13, 0}
+        for i := range expected {
+                value := bitreader.TryReadInt8()
+                if int(value) != expected[i] {
+                        t.Fatalf("TryReadInt8 FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                }
+        }
+}
+
+func TestTryReadInt16(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []int{28941, 0, 41499, 0}
+        for i := range expected {
+                value := bitreader.TryReadInt16()
+                if int(value) != expected[i] {
+                        t.Fatalf("TryReadInt16 FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                }
+        }
+}
+
+func TestTryReadInt32(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []int{1896677376, 2719678464}
+        for i := range expected {
+                value := bitreader.TryReadInt32()
+                if int(value) != expected[i] {
+                        t.Fatalf("TryReadInt32 FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                }
+        }
+}
+
+func TestTryReadInt64(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []int{8146167303702773760}
+        for i := range expected {
+                value := bitreader.TryReadInt64()
+                if int(value) != expected[i] {
+                        t.Fatalf("TryReadInt64 FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                }
+        }
+}
+
 func TestReadBit(t *testing.T) {
-        bitreader := Reader(TestArray[:], false)
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
         expected := []bool{false, true, true, true}
         for i := range expected {
-                value, err := bitreader.ReadBit()
+                value, err := bitreader.ReadBool()
                 if err != nil {
                         t.Fatal(err)
                 }
@@ -22,27 +89,27 @@ func TestReadBit(t *testing.T) {
 }
 
 func TestReadBitLE(t *testing.T) {
-        bitreader := Reader(TestArray[:], true)
+        bitreader := Reader(bytes.NewReader(TestArray[:]), true)
         expected := []bool{true, false, false, false}
         for i := range expected {
-                value, err := bitreader.ReadBit()
+                value, err := bitreader.ReadBool()
                 if err != nil {
                         t.Fatal(err)
                 }
                 if value != expected[i] {
-                        t.Fatalf("ReadBitLSB FAIL for index %d: Expected %t, Got %t", i, expected[i], value)
+                        t.Fatalf("ReadBitLE FAIL for index %d: Expected %t, Got %t", i, expected[i], value)
                 }
         }
 }
 
-func TestReadBits32(t *testing.T) {
-        bitreader := Reader(TestArray[:], false)
+func TestReadBits(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
         expected := []int{3793354753, 2288779267} // 11100010000110100000000000000001, 10001000011011000000000000000011
         expectedBool := []bool{false, false}
         for i := range expected {
-                bool, err := bitreader.ReadBit()
+                bool, err := bitreader.ReadBool()
                 if bool != expectedBool[i] {
-                        t.Fatalf("ReadBits32 ReadBit FAIL for index %d: Expected %t, Got %t", i, expectedBool[i], bool)
+                        t.Fatalf("ReadBits ReadBit FAIL for index %d: Expected %t, Got %t", i, expectedBool[i], bool)
                 }
                 if err != nil {
                         t.Fatal(err)
@@ -51,36 +118,73 @@ func TestReadBits32(t *testing.T) {
                 if err != nil {
                         t.Fatal(err)
                 }
-                if value != expected[i] {
-                        t.Fatalf("ReadBits32 FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                if int(value) != expected[i] {
+                        t.Fatalf("ReadBits FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
                 }
         }
 }
 
 func TestReadBitsLE(t *testing.T) {
-        bitreader := Reader(TestArray[:], true)
+        bitreader := Reader(bytes.NewReader(TestArray[:]), true)
         expected := []int{1720, 1768} // 11010111000, 11011101000
         for i := range expected {
-                bitreader.ReadBit()
+                bitreader.ReadBool()
                 value, err := bitreader.ReadBits(32)
                 if err != nil {
                         t.Fatal(err)
                 }
-                if value != expected[i] {
-                        t.Fatalf("ReadBits32LSB FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                if int(value) != expected[i] {
+                        t.Fatalf("ReadBitsLE FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                }
+        }
+}
+
+func TestReadBytes(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
+        expected := []int{3793354753, 2288779267} // 11100010000110100000000000000001, 10001000011011000000000000000011
+        expectedBool := []bool{false, false}
+        for i := range expected {
+                bool, err := bitreader.ReadBool()
+                if bool != expectedBool[i] {
+                        t.Fatalf("ReadBytes ReadBit FAIL for index %d: Expected %t, Got %t", i, expectedBool[i], bool)
+                }
+                if err != nil {
+                        t.Fatal(err)
+                }
+                value, err := bitreader.ReadBytes(4)
+                if err != nil {
+                        t.Fatal(err)
+                }
+                if int(value) != expected[i] {
+                        t.Fatalf("ReadBytes FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
+                }
+        }
+}
+
+func TestReadBytesLE(t *testing.T) {
+        bitreader := Reader(bytes.NewReader(TestArray[:]), true)
+        expected := []int{1720, 1768} // 11010111000, 11011101000
+        for i := range expected {
+                bitreader.ReadBool()
+                value, err := bitreader.ReadBytes(4)
+                if err != nil {
+                        t.Fatal(err)
+                }
+                if int(value) != expected[i] {
+                        t.Fatalf("ReadBytesLE FAIL for index %d: Expected %d, Got %d", i, expected[i], value)
                 }
         }
 }
 
 func TestSkipBits(t *testing.T) {
-        bitreader := Reader(TestArray[:], false)
+        bitreader := Reader(bytes.NewReader(TestArray[:]), false)
         expected := []bool{true, true, false, true} //00001101
         err := bitreader.SkipBits(12)
         if err != nil {
                 t.Fatal(err)
         }
         for i := range expected {
-                value, err := bitreader.ReadBit()
+                value, err := bitreader.ReadBool()
                 if err != nil {
                         t.Fatal(err)
                 }
@@ -91,16 +195,16 @@ func TestSkipBits(t *testing.T) {
 }
 
 func TestSkipBitsLE(t *testing.T) {
-        bitreader := Reader(TestArray[:], true)
+        bitreader := Reader(bytes.NewReader(TestArray[:]), true)
         expected := []bool{false, false, false, false} //10110000
         bitreader.SkipBits(12)
         for i := range expected {
-                value, err := bitreader.ReadBit()
+                value, err := bitreader.ReadBool()
                 if err != nil {
                         t.Fatal(err)
                 }
                 if value != expected[i] {
-                        t.Fatalf("SkipBits ReadBit FAIL for index %d: Expected %t, Got %t", i, expected[i], value)
+                        t.Fatalf("SkipBitsLE ReadBit FAIL for index %d: Expected %t, Got %t", i, expected[i], value)
                 }
         }
 }