1 files changed, 189 insertions, 92 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
 }