├── README ├── LICENSE ├── bitstream.go └── bitstream_test.go /README: -------------------------------------------------------------------------------- 1 | bitstream: read and write bits from io.Reader and io.Writer streams 2 | 3 | Full package documentation available at: http://godoc.org/github.com/dgryski/go-bitstream 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Damian Gryski 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /bitstream.go: -------------------------------------------------------------------------------- 1 | // Package bitstream is a simple wrapper around a io.Reader and io.Writer to provide bit-level access to the stream. 2 | package bitstream 3 | 4 | import ( 5 | "io" 6 | ) 7 | 8 | // A Bit is a zero or a one 9 | type Bit bool 10 | 11 | const ( 12 | // Zero is our exported type for '0' bits 13 | Zero Bit = false 14 | // One is our exported type for '1' bits 15 | One Bit = true 16 | ) 17 | 18 | // A BitReader reads bits from an io.Reader 19 | type BitReader struct { 20 | r io.Reader 21 | b [1]byte 22 | count uint8 23 | } 24 | 25 | // A BitWriter writes bits to an io.Writer 26 | type BitWriter struct { 27 | w io.Writer 28 | b [1]byte 29 | count uint8 30 | } 31 | 32 | // NewReader returns a BitReader that returns a single bit at a time from 'r' 33 | func NewReader(r io.Reader) *BitReader { 34 | b := new(BitReader) 35 | b.r = r 36 | return b 37 | } 38 | 39 | // ReadBit returns the next bit from the stream, reading a new byte from the underlying reader if required. 40 | func (b *BitReader) ReadBit() (Bit, error) { 41 | if b.count == 0 { 42 | if n, err := b.r.Read(b.b[:]); n != 1 || (err != nil && err != io.EOF) { 43 | return Zero, err 44 | } 45 | b.count = 8 46 | } 47 | b.count-- 48 | d := (b.b[0] & 0x80) 49 | b.b[0] <<= 1 50 | return d != 0, nil 51 | } 52 | 53 | // NewWriter returns a BitWriter that buffers bits and write the resulting bytes to 'w' 54 | func NewWriter(w io.Writer) *BitWriter { 55 | b := new(BitWriter) 56 | b.w = w 57 | b.count = 8 58 | return b 59 | } 60 | 61 | func (b *BitWriter) Pending() (byt byte, vals uint8) { 62 | return b.b[0], b.count 63 | } 64 | 65 | func (b *BitWriter) Resume(data byte, count uint8) { 66 | b.b[0] = data 67 | b.count = count 68 | } 69 | 70 | // WriteBit writes a single bit to the stream, writing a new byte to 'w' if required. 71 | func (b *BitWriter) WriteBit(bit Bit) error { 72 | 73 | if bit { 74 | b.b[0] |= 1 << (b.count - 1) 75 | } 76 | 77 | b.count-- 78 | 79 | if b.count == 0 { 80 | if n, err := b.w.Write(b.b[:]); n != 1 || err != nil { 81 | return err 82 | } 83 | b.b[0] = 0 84 | b.count = 8 85 | } 86 | 87 | return nil 88 | } 89 | 90 | // WriteByte writes a single byte to the stream, regardless of alignment 91 | func (b *BitWriter) WriteByte(byt byte) error { 92 | 93 | // fill up b.b with b.count bits from byt 94 | b.b[0] |= byt >> (8 - b.count) 95 | 96 | if n, err := b.w.Write(b.b[:]); n != 1 || err != nil { 97 | return err 98 | } 99 | b.b[0] = byt << b.count 100 | 101 | return nil 102 | } 103 | 104 | // ReadByte reads a single byte from the stream, regardless of alignment 105 | func (b *BitReader) ReadByte() (byte, error) { 106 | 107 | if b.count == 0 { 108 | n, err := b.r.Read(b.b[:]) 109 | if n != 1 || (err != nil && err != io.EOF) { 110 | b.b[0] = 0 111 | return b.b[0], err 112 | } 113 | // mask io.EOF for the last byte 114 | if err == io.EOF { 115 | err = nil 116 | } 117 | return b.b[0], err 118 | } 119 | 120 | byt := b.b[0] 121 | 122 | var n int 123 | var err error 124 | n, err = b.r.Read(b.b[:]) 125 | if n != 1 || (err != nil && err != io.EOF) { 126 | return 0, err 127 | } 128 | 129 | byt |= b.b[0] >> b.count 130 | 131 | b.b[0] <<= (8 - b.count) 132 | 133 | return byt, err 134 | } 135 | 136 | // ReadBits reads nbits from the stream 137 | func (b *BitReader) ReadBits(nbits int) (uint64, error) { 138 | 139 | var u uint64 140 | 141 | for nbits >= 8 { 142 | byt, err := b.ReadByte() 143 | if err != nil { 144 | return 0, err 145 | } 146 | 147 | u = (u << 8) | uint64(byt) 148 | nbits -= 8 149 | } 150 | 151 | var err error 152 | for nbits > 0 && err != io.EOF { 153 | byt, err := b.ReadBit() 154 | if err != nil { 155 | return 0, err 156 | } 157 | u <<= 1 158 | if byt { 159 | u |= 1 160 | } 161 | nbits-- 162 | } 163 | 164 | return u, nil 165 | } 166 | 167 | // Flush empties the currently in-process byte by filling it with 'bit'. 168 | func (b *BitWriter) Flush(bit Bit) error { 169 | 170 | for b.count != 8 { 171 | err := b.WriteBit(bit) 172 | if err != nil { 173 | return err 174 | } 175 | } 176 | 177 | return nil 178 | } 179 | 180 | // WriteBits writes the nbits least significant bits of u, most-significant-bit first. 181 | func (b *BitWriter) WriteBits(u uint64, nbits int) error { 182 | u <<= (64 - uint(nbits)) 183 | for nbits >= 8 { 184 | byt := byte(u >> 56) 185 | err := b.WriteByte(byt) 186 | if err != nil { 187 | return err 188 | } 189 | u <<= 8 190 | nbits -= 8 191 | } 192 | 193 | for nbits > 0 { 194 | err := b.WriteBit((u >> 63) == 1) 195 | if err != nil { 196 | return err 197 | } 198 | u <<= 1 199 | nbits-- 200 | } 201 | 202 | return nil 203 | } 204 | 205 | // Reset replaces the underlying io.Reader with the provided reader and resets 206 | // all interal state to its zero value, allowing for reuse of the BitReader 207 | // without additional allocations 208 | func (b *BitReader) Reset(r io.Reader) { 209 | b.r = r 210 | b.b[0] = 0x00 211 | b.count = 0 212 | } 213 | 214 | // Reset replaces the underlying io.Writer with the provided writer and resets 215 | // all internal state to its initial value allowing for reuse of the BitWriter 216 | // without additional allocations 217 | func (b *BitWriter) Reset(w io.Writer) { 218 | b.w = w 219 | b.b[0] = 0x00 220 | b.count = 8 221 | } 222 | -------------------------------------------------------------------------------- /bitstream_test.go: -------------------------------------------------------------------------------- 1 | package bitstream 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io" 7 | "strings" 8 | "testing" 9 | ) 10 | 11 | func TestBitStreamEOF(t *testing.T) { 12 | 13 | br := NewReader(strings.NewReader("0")) 14 | 15 | b, err := br.ReadByte() 16 | if b != '0' { 17 | t.Error("ReadByte didn't return first byte") 18 | } 19 | 20 | b, err = br.ReadByte() 21 | if err != io.EOF { 22 | t.Error("ReadByte on empty string didn't return EOF") 23 | } 24 | 25 | // 0 = 0b00110000 26 | br = NewReader(strings.NewReader("0")) 27 | 28 | buf := bytes.NewBuffer(nil) 29 | bw := NewWriter(buf) 30 | 31 | for i := 0; i < 4; i++ { 32 | bit, err := br.ReadBit() 33 | if err == io.EOF { 34 | break 35 | } 36 | if err != nil { 37 | t.Error("GetBit returned error err=", err.Error()) 38 | return 39 | } 40 | err = bw.WriteBit(bit) 41 | if err != nil { 42 | t.Errorf("unexpected writer error") 43 | } 44 | } 45 | 46 | bw.Flush(One) 47 | 48 | err = bw.WriteByte(0xAA) 49 | if err != nil { 50 | t.Error("unable to WriteByte") 51 | } 52 | 53 | c := buf.Bytes() 54 | 55 | if len(c) != 2 || c[1] != 0xAA || c[0] != 0x3f { 56 | t.Error("bad return from 4 read bytes") 57 | } 58 | 59 | _, err = NewReader(strings.NewReader("")).ReadBit() 60 | if err != io.EOF { 61 | t.Error("ReadBit on empty string didn't return EOF") 62 | } 63 | 64 | } 65 | 66 | // reader that returns a single byte with EOF 67 | type eofReader struct { 68 | readCalled bool 69 | } 70 | 71 | func (e *eofReader) Read(p []byte) (int, error) { 72 | if e.readCalled { 73 | return 0, io.EOF 74 | } 75 | e.readCalled = true 76 | p[0] = 'A' 77 | return 1, io.EOF 78 | } 79 | 80 | func TestBitStreamReadEOF(t *testing.T) { 81 | 82 | br := NewReader(&eofReader{}) 83 | 84 | r, _ := br.ReadBits(8) 85 | 86 | if r != 'A' { 87 | t.Errorf("failed to read single byte with EOF: %08b", r) 88 | } 89 | 90 | } 91 | 92 | func TestBitStream(t *testing.T) { 93 | 94 | buf := bytes.NewBuffer(nil) 95 | br := NewReader(strings.NewReader("hello")) 96 | bw := NewWriter(buf) 97 | 98 | for { 99 | bit, err := br.ReadBit() 100 | if err == io.EOF { 101 | break 102 | } 103 | if err != nil { 104 | t.Error("GetBit returned error err=", err.Error()) 105 | return 106 | } 107 | err = bw.WriteBit(bit) 108 | if err != nil { 109 | t.Errorf("unexpected writer error") 110 | } 111 | } 112 | 113 | s := buf.String() 114 | 115 | if s != "hello" { 116 | t.Error("expected 'hello', got=", []byte(s)) 117 | } 118 | } 119 | 120 | func TestByteStream(t *testing.T) { 121 | 122 | buf := bytes.NewBuffer(nil) 123 | br := NewReader(strings.NewReader("hello")) 124 | bw := NewWriter(buf) 125 | 126 | for i := 0; i < 3; i++ { 127 | bit, err := br.ReadBit() 128 | if err == io.EOF { 129 | break 130 | } 131 | if err != nil { 132 | t.Error("GetBit returned error err=", err.Error()) 133 | return 134 | } 135 | err = bw.WriteBit(bit) 136 | if err != nil { 137 | t.Errorf("unexpected writer error") 138 | } 139 | } 140 | 141 | for i := 0; i < 3; i++ { 142 | byt, err := br.ReadByte() 143 | if err == io.EOF { 144 | break 145 | } 146 | if err != nil { 147 | t.Error("GetByte returned error err=", err.Error()) 148 | return 149 | } 150 | bw.WriteByte(byt) 151 | } 152 | 153 | u, err := br.ReadBits(13) 154 | 155 | if err != nil { 156 | t.Error("ReadBits returned error err=", err.Error()) 157 | return 158 | } 159 | 160 | err = bw.WriteBits(u, 13) 161 | if err != nil { 162 | t.Errorf("unexpected writer error") 163 | } 164 | 165 | err = bw.WriteBits(('!'<<12)|('.'<<4)|0x02, 20) 166 | if err != nil { 167 | t.Errorf("unexpected writer error") 168 | } 169 | // 0x2f == '/' 170 | bw.Flush(One) 171 | 172 | s := buf.String() 173 | 174 | if s != "hello!./" { 175 | t.Errorf("expected 'hello!./', got=%x", []byte(s)) 176 | } 177 | } 178 | 179 | var myError error = fmt.Errorf("my error") 180 | 181 | type badWriter struct{} 182 | 183 | func (w *badWriter) Write(p []byte) (n int, err error) { 184 | return 0, myError 185 | } 186 | func TestErrorPropagation(t *testing.T) { 187 | // check WriteBit 188 | w := &badWriter{} 189 | bw := NewWriter(w) 190 | for i := 0; i < 7; i++ { 191 | err := bw.WriteBit(One) 192 | if err != nil { 193 | t.Errorf("unexpected error during buffered write operation") 194 | } 195 | } 196 | err := bw.WriteBit(One) 197 | if err != myError { 198 | t.Errorf("failed to propagate error") 199 | } 200 | 201 | // check WriteBits 202 | w = &badWriter{} 203 | bw = NewWriter(w) 204 | err = bw.WriteBits(256, 8) 205 | if err != myError { 206 | t.Errorf("failed to propagate error") 207 | } 208 | } 209 | 210 | func TestReset(t *testing.T) { 211 | r := NewReader(strings.NewReader("abc")) 212 | b, _ := r.ReadBits(8) 213 | if b != 'a' { 214 | t.Errorf("expected 'a', got=%08b", b) 215 | } 216 | 217 | r.Reset(strings.NewReader("def")) 218 | b, _ = r.ReadBits(8) 219 | 220 | if b != 'd' { 221 | t.Errorf("expected 'd', got=%08b", b) 222 | } 223 | 224 | firstWriter := bytes.NewBuffer(nil) 225 | w := NewWriter(firstWriter) 226 | w.WriteBits(97, 8) 227 | 228 | secondWriter := bytes.NewBuffer(nil) 229 | w.Reset(secondWriter) 230 | w.WriteBits(98, 8) 231 | 232 | if firstWriter.String() != "a" { 233 | t.Errorf("expected 'a', got=%x", firstWriter.Bytes()) 234 | } 235 | 236 | if secondWriter.String() != "b" { 237 | t.Errorf("expected 'b', got=%x", secondWriter.Bytes()) 238 | } 239 | } 240 | --------------------------------------------------------------------------------