├── .gitignore ├── LICENSE ├── README.md ├── doc.go ├── doc_test.go ├── go.mod ├── go.sum ├── reader.go ├── reader_test.go ├── writer.go └── writer_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, build with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Takashi Oguma 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # bitstream-go 2 | 3 | A practical, high-performance, and easy-to-use bit stream reader/writer for golang. 4 | 5 | - Type aware (you don't need type castings to fit into your type) 6 | - Endianess aware (Little endian support is still work in progress) 7 | 8 | ## Usage 9 | 10 | Reader 11 | 12 | ``` 13 | package main 14 | 15 | import ( 16 | "bytes" 17 | "fmt" 18 | "log" 19 | 20 | "github.com/bearmini/bitstream-go" 21 | ) 22 | 23 | func main() { 24 | // binary expression: 25 | // 0000 0001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111 26 | data := []byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef} 27 | 28 | // Reader 29 | r := bitstream.NewReader(bytes.NewReader(data), nil) 30 | 31 | // read a single bit 32 | bit0, err := r.ReadBit() 33 | if err != nil { 34 | log.Fatalf("%+v", err) 35 | } 36 | fmt.Printf("bit: %1b\n", bit0) 37 | 38 | // read 2 bits 39 | bit1to2, err := r.ReadNBitsAsUint8(2) 40 | if err != nil { 41 | log.Fatalf("%+v", err) 42 | } 43 | fmt.Printf("bits: %02b\n", bit1to2) 44 | 45 | // read 10 bits as big endian 46 | bit3to12, err := r.ReadNBitsAsUint16BE(10) 47 | if err != nil { 48 | log.Fatalf("%+v", err) 49 | } 50 | fmt.Printf("bits: %010b\n", bit3to12) 51 | 52 | // read 20 bits as big endian 53 | bit13to32, err := r.ReadNBitsAsUint32BE(20) 54 | if err != nil { 55 | log.Fatalf("%+v", err) 56 | } 57 | fmt.Printf("bits: %020b\n", bit13to32) 58 | 59 | // Output: 60 | // bit: 0 61 | // bits: 00 62 | // bits: 0000100100 63 | // bits: 01101000101011001111 64 | } 65 | ``` 66 | 67 | 68 | Writer 69 | ``` 70 | package main 71 | 72 | import ( 73 | "bytes" 74 | "encoding/hex" 75 | "fmt" 76 | "log" 77 | 78 | "github.com/bearmini/bitstream-go" 79 | ) 80 | 81 | func main() { 82 | dst := bytes.NewBuffer([]byte{}) 83 | 84 | // Writer 85 | w := bitstream.NewWriter(dst) 86 | 87 | // Write a single bit `1` 88 | err := w.WriteBit(1) 89 | if err != nil { 90 | log.Fatalf("%+v", err) 91 | } 92 | 93 | // Write a bool value as a bit (true: 1, false: 0) 94 | err = w.WriteBool(false) 95 | if err != nil { 96 | log.Fatalf("%+v", err) 97 | } 98 | 99 | // Write 2 bits `10` 100 | err = w.WriteNBitsOfUint8(2, 0x02) 101 | if err != nil { 102 | log.Fatalf("%+v", err) 103 | } 104 | 105 | // Write 8 bits `0101 0011` 106 | err = w.WriteUint8(0x53) 107 | if err != nil { 108 | log.Fatalf("%+v", err) 109 | } 110 | 111 | // Write 10 bits `11 0010 1101` 112 | err = w.WriteNBitsOfUint16BE(10, 0x032d) 113 | if err != nil { 114 | log.Fatalf("%+v", err) 115 | } 116 | 117 | // Write 16 bits `0000 1111 0101 1010` 118 | err = w.WriteUint16BE(0x0f5a) 119 | if err != nil { 120 | log.Fatalf("%+v", err) 121 | } 122 | 123 | w.Flush() 124 | 125 | // we have written the following bits: 126 | // 1 127 | // 0 128 | // 10 129 | // 0101 0011 130 | // 1100 1011 01 131 | // 00 0011 1101 0110 10 132 | // 1010 0101 0011 1100 1011 0100 0011 1101 0110 10xx 133 | 134 | fmt.Printf("%s", hex.EncodeToString(dst.Bytes())) 135 | // Output: 136 | // a53cb43d68 137 | } 138 | ``` -------------------------------------------------------------------------------- /doc.go: -------------------------------------------------------------------------------- 1 | // Copyright 2018 bearmini (Takashi Oguma) 2 | /* 3 | bitstream-go is a package to provide a bitstream reader / a bitstream writer. 4 | 5 | 6 | */ 7 | 8 | package bitstream 9 | -------------------------------------------------------------------------------- /doc_test.go: -------------------------------------------------------------------------------- 1 | package bitstream_test 2 | 3 | import ( 4 | "bytes" 5 | "encoding/hex" 6 | "fmt" 7 | "log" 8 | 9 | "github.com/bearmini/bitstream-go" 10 | ) 11 | 12 | func ExampleReader() { 13 | // binary expression: 14 | // 0000 0001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111 15 | data := []byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef} 16 | 17 | // Reader 18 | r := bitstream.NewReader(bytes.NewReader(data), nil) 19 | 20 | // read a single bit 21 | bit0, err := r.ReadBit() 22 | if err != nil { 23 | log.Fatalf("%+v", err) 24 | } 25 | fmt.Printf("bit: %1b\n", bit0) 26 | 27 | // read 2 bits 28 | bit1to2, err := r.ReadNBitsAsUint8(2) 29 | if err != nil { 30 | log.Fatalf("%+v", err) 31 | } 32 | fmt.Printf("bits: %02b\n", bit1to2) 33 | 34 | // read 10 bits as big endian 35 | bit3to12, err := r.ReadNBitsAsUint16BE(10) 36 | if err != nil { 37 | log.Fatalf("%+v", err) 38 | } 39 | fmt.Printf("bits: %010b\n", bit3to12) 40 | 41 | // read 20 bits as big endian 42 | bit13to32, err := r.ReadNBitsAsUint32BE(20) 43 | if err != nil { 44 | log.Fatalf("%+v", err) 45 | } 46 | fmt.Printf("bits: %020b\n", bit13to32) 47 | 48 | // Output: 49 | // bit: 0 50 | // bits: 00 51 | // bits: 0000100100 52 | // bits: 01101000101011001111 53 | } 54 | 55 | func ExampleWriter() { 56 | dst := bytes.NewBuffer([]byte{}) 57 | 58 | // Writer 59 | w := bitstream.NewWriter(dst) 60 | 61 | // Write a single bit `1` 62 | err := w.WriteBit(1) 63 | if err != nil { 64 | log.Fatalf("%+v", err) 65 | } 66 | 67 | // Write a bool value as a bit (true: 1, false: 0) 68 | err = w.WriteBool(false) 69 | if err != nil { 70 | log.Fatalf("%+v", err) 71 | } 72 | 73 | // Write 2 bits `10` 74 | err = w.WriteNBitsOfUint8(2, 0x02) 75 | if err != nil { 76 | log.Fatalf("%+v", err) 77 | } 78 | 79 | // Write 8 bits `0101 0011` 80 | err = w.WriteUint8(0x53) 81 | if err != nil { 82 | log.Fatalf("%+v", err) 83 | } 84 | 85 | // Write 10 bits `11 0010 1101` 86 | err = w.WriteNBitsOfUint16BE(10, 0x032d) 87 | if err != nil { 88 | log.Fatalf("%+v", err) 89 | } 90 | 91 | // Write 16 bits `0000 1111 0101 1010` 92 | err = w.WriteUint16BE(0x0f5a) 93 | if err != nil { 94 | log.Fatalf("%+v", err) 95 | } 96 | 97 | w.Flush() 98 | 99 | // we have written the following bits: 100 | // 1 101 | // 0 102 | // 10 103 | // 0101 0011 104 | // 1100 1011 01 105 | // 00 0011 1101 0110 10 106 | // 1010 0101 0011 1100 1011 0100 0011 1101 0110 10xx 107 | 108 | fmt.Printf("%s", hex.EncodeToString(dst.Bytes())) 109 | // Output: 110 | // a53cb43d68 111 | } 112 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/bearmini/bitstream-go 2 | 3 | require github.com/pkg/errors v0.8.0 4 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw= 2 | github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 3 | -------------------------------------------------------------------------------- /reader.go: -------------------------------------------------------------------------------- 1 | package bitstream 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | 7 | "github.com/pkg/errors" 8 | ) 9 | 10 | const ( 11 | DefaultBufferSize = 1024 12 | ) 13 | 14 | // Reader is a bit stream reader. 15 | // It does not have io.Reader interface. 16 | type Reader struct { 17 | src io.Reader 18 | srcEOF bool 19 | buf []byte 20 | bufLen uint 21 | currByteIndex uint // starts from 0 22 | currBitIndex uint8 // MSB: 7, LSB: 0 23 | consumedBytes uint 24 | opt *ReaderOptions 25 | } 26 | 27 | // ReaderOptions is a set of options for creating a Reader. 28 | type ReaderOptions struct { 29 | BufferSize uint 30 | } 31 | 32 | // GetBufferSize gets configured buffer size. 33 | func (opt *ReaderOptions) GetBufferSize() uint { 34 | if opt == nil || opt.BufferSize == 0 { 35 | return DefaultBufferSize 36 | } 37 | return opt.BufferSize 38 | } 39 | 40 | // NewReader creates a new Reader instance with options. 41 | func NewReader(src io.Reader, opt *ReaderOptions) *Reader { 42 | return &Reader{ 43 | src: src, 44 | srcEOF: false, 45 | buf: nil, 46 | bufLen: 0, 47 | currByteIndex: 0, 48 | currBitIndex: 7, 49 | opt: opt, 50 | } 51 | } 52 | 53 | func (r *Reader) dump() { 54 | fmt.Printf("srcEOF=%t, bufLen=%d, currByteIndex=%d, currBitIndex=%d\n", r.srcEOF, r.bufLen, r.currByteIndex, r.currBitIndex) 55 | } 56 | 57 | func (r *Reader) isBufEmpty() bool { 58 | if r.buf == nil { 59 | return true 60 | } 61 | 62 | if r.currByteIndex >= r.bufLen { 63 | return true 64 | } 65 | 66 | return false 67 | } 68 | 69 | func (r *Reader) fillBuf() error { 70 | buf := make([]byte, r.opt.GetBufferSize()) 71 | n, err := r.src.Read(buf[:]) 72 | if err != nil { 73 | return err 74 | } 75 | 76 | r.buf = buf 77 | r.bufLen = uint(n) 78 | r.currByteIndex = 0 79 | r.currBitIndex = 7 80 | return nil 81 | } 82 | 83 | func (r *Reader) fillBufIfNeeded() error { 84 | if !r.isBufEmpty() { 85 | return nil 86 | } 87 | return r.fillBuf() 88 | } 89 | 90 | func (r *Reader) forwardIndecies(nBits uint8) { 91 | if nBits <= r.currBitIndex { 92 | r.currBitIndex -= nBits 93 | return 94 | } 95 | 96 | nBits = nBits - r.currBitIndex 97 | nBytes := uint(nBits/8) + 1 98 | r.currByteIndex += nBytes 99 | r.consumedBytes += nBytes 100 | 101 | bitsToGo := (nBits % 8) 102 | r.currBitIndex = 8 - bitsToGo 103 | } 104 | 105 | // ConsumedBytes returns a number of bytes that has been consumed. 106 | func (r *Reader) ConsumedBytes() uint { 107 | if r.currBitIndex != 7 { 108 | return r.consumedBytes + 1 109 | } 110 | return r.consumedBytes 111 | } 112 | 113 | // ReadBit reads a single bit from the bit stream. 114 | // The bit read from the stream will be set in the LSB of the return value. 115 | func (r *Reader) ReadBit() (byte, error) { 116 | err := r.fillBufIfNeeded() 117 | if err != nil { 118 | return 0, err 119 | } 120 | 121 | b := r.buf[r.currByteIndex] 122 | mask := uint8(1 << r.currBitIndex) 123 | result := (b & mask) >> r.currBitIndex 124 | r.forwardIndecies(1) 125 | return result, nil 126 | } 127 | 128 | // ReadBool reads a single bit from the bit stream and return it as a bool. 129 | func (r *Reader) ReadBool() (bool, error) { 130 | b, err := r.ReadBit() 131 | if err != nil { 132 | return false, err 133 | } 134 | return b != 0, nil 135 | } 136 | 137 | func (r *Reader) mustReadNBitsInCurrentByte(nBits uint8) byte { 138 | if nBits == 0 { 139 | return 0 140 | } 141 | 142 | if r.currBitIndex < (nBits - 1) { 143 | panic(fmt.Sprintf("%+v", errors.New("insufficient bits to read"))) 144 | } 145 | 146 | b := r.buf[r.currByteIndex] 147 | mask := uint8((1 << (r.currBitIndex + 1)) - 1) 148 | result := (b & mask) >> (r.currBitIndex - (nBits - 1)) 149 | r.forwardIndecies(nBits) 150 | return result 151 | } 152 | 153 | // ReadNBitsAsUint8 reads `nBits` bits as a unsigned integer from the bit stream and returns it in uint8 (LSB aligned). 154 | // `nBits` must be less than or equal to 8, otherwise returns an error. 155 | // If `nBits` == 0, this function always returns 0. 156 | func (r *Reader) ReadNBitsAsUint8(nBits uint8) (uint8, error) { 157 | if nBits == 0 { 158 | return 0, nil 159 | } 160 | 161 | if nBits > 8 { 162 | return 0, errors.New("nBits too large for uint8") 163 | } 164 | 165 | err := r.fillBufIfNeeded() 166 | if err != nil { 167 | return 0, err 168 | } 169 | 170 | // remaining bits in current byte 171 | rb := r.currBitIndex + 1 172 | 173 | if nBits <= rb { // can be read from the current byte 174 | b := r.mustReadNBitsInCurrentByte(nBits) 175 | return b, nil 176 | } 177 | 178 | // 8 bits are distributed in 2 bytes 179 | nBits1 := rb 180 | nBits2 := nBits - rb 181 | 182 | b1 := r.mustReadNBitsInCurrentByte(nBits1) 183 | b2, err := r.ReadNBitsAsUint8(nBits2) 184 | if err != nil { 185 | return 0, err 186 | } 187 | 188 | return (b1 << nBits2) | b2, nil 189 | } 190 | 191 | // ReadUint8 reads 8 bits from the bit stream and returns it in uint8. 192 | func (r *Reader) ReadUint8() (uint8, error) { 193 | return r.ReadNBitsAsUint8(8) 194 | } 195 | 196 | // ReadNBitsAsUint16BE reads `nBits` bits as a big endian unsigned integer from the bit stream and returns it in uint16 (LSB aligned). 197 | // `nBits` must be less than or equal to 16, otherwise returns an error. 198 | // If `nBits` == 0, this function always returns 0. 199 | func (r *Reader) ReadNBitsAsUint16BE(nBits uint8) (uint16, error) { 200 | if nBits == 0 { 201 | return 0, nil 202 | } 203 | 204 | if nBits <= 8 { 205 | v, err := r.ReadNBitsAsUint8(nBits) 206 | return uint16(v), err 207 | } 208 | 209 | if nBits > 16 { 210 | return 0, errors.New("nBits too large for uint16") 211 | } 212 | 213 | err := r.fillBufIfNeeded() 214 | if err != nil { 215 | return 0, err 216 | } 217 | 218 | // remaining bits in current byte 219 | rb := r.currBitIndex + 1 220 | 221 | // 16 bits may be distributed in up to 3 bytes 222 | nBits1 := rb // count of bits in the first byte 223 | nBits2 := nBits - rb // count of bits in the second byte 224 | nBits3 := uint8(0) // count of bits in the third byte 225 | if nBits2 > 8 { 226 | nBits3 = nBits2 - 8 227 | nBits2 = 8 228 | } 229 | 230 | b1 := r.mustReadNBitsInCurrentByte(nBits1) 231 | b2, err := r.ReadNBitsAsUint8(nBits2) 232 | if err != nil { 233 | return 0, err 234 | } 235 | b3, err := r.ReadNBitsAsUint8(nBits3) // expects this function returns 0 if nBits3 == 0 236 | if err != nil { 237 | return 0, err 238 | } 239 | 240 | return (uint16(b1) << (nBits2 + nBits3)) | (uint16(b2) << nBits3) | uint16(b3), nil 241 | } 242 | 243 | // ReadUint16BE reads 16 bits as a big endian unsigned integer from the bit stream and returns it in uint16. 244 | func (r *Reader) ReadUint16BE() (uint16, error) { 245 | return r.ReadNBitsAsUint16BE(16) 246 | } 247 | 248 | // ReadNBitsAsUint32BE reads `nBits` bits as a big endian unsigned integer from the bit stream and returns it in uint32 (LSB aligned). 249 | // `nBits` must be less than or equal to 32, otherwise returns an error. 250 | // If `nBits` == 0, this function always returns 0. 251 | func (r *Reader) ReadNBitsAsUint32BE(nBits uint8) (uint32, error) { 252 | if nBits == 0 { 253 | return 0, nil 254 | } 255 | 256 | if nBits <= 16 { 257 | v, err := r.ReadNBitsAsUint16BE(nBits) 258 | return uint32(v), err 259 | } 260 | 261 | if nBits > 32 { 262 | return 0, errors.New("nBits too large for uint32") 263 | } 264 | 265 | err := r.fillBufIfNeeded() 266 | if err != nil { 267 | return 0, err 268 | } 269 | 270 | // remaining bits in current byte 271 | rb := r.currBitIndex + 1 272 | 273 | // 32 bits may be distributed in up to 5 bytes 274 | nBits1 := rb 275 | nBits2 := uint8(8) 276 | nBits3 := nBits - rb - 8 277 | nBits4 := uint8(0) 278 | nBits5 := uint8(0) 279 | if nBits3 > 8 { 280 | nBits4 = nBits3 - 8 281 | if nBits4 > 8 { 282 | nBits5 = nBits4 - 8 283 | nBits4 = 8 284 | } 285 | nBits3 = 8 286 | } 287 | 288 | b1 := r.mustReadNBitsInCurrentByte(nBits1) 289 | b2, err := r.ReadNBitsAsUint8(nBits2) 290 | if err != nil { 291 | return 0, err 292 | } 293 | b3, err := r.ReadNBitsAsUint8(nBits3) 294 | if err != nil { 295 | return 0, err 296 | } 297 | b4, err := r.ReadNBitsAsUint8(nBits4) 298 | if err != nil { 299 | return 0, err 300 | } 301 | b5, err := r.ReadNBitsAsUint8(nBits5) 302 | if err != nil { 303 | return 0, err 304 | } 305 | 306 | return (uint32(b1) << (nBits2 + nBits3 + nBits4 + nBits5)) | (uint32(b2) << (nBits3 + nBits4 + nBits5)) | (uint32(b3) << (nBits4 + nBits5)) | (uint32(b4) << (nBits5)) | uint32(b5), nil 307 | } 308 | 309 | // ReadUint32BE reads 32 bits as a big endian unsigned integer from the bit stream and returns it in uint32. 310 | func (r *Reader) ReadUint32BE() (uint32, error) { 311 | return r.ReadNBitsAsUint32BE(32) 312 | } 313 | 314 | // ReadNBitsAsInt32BE reads `nBits` bits as a big endian signed integer from the bit stream and returns it in int32 (LSB aligned). 315 | // MSB is a sign bit. 316 | // `nBits` must be less than or equal to 32, otherwise returns an error. 317 | // If `nBits` == 0, this function always returns 0. 318 | func (r *Reader) ReadNBitsAsInt32BE(nBits uint8) (int32, error) { 319 | v, err := r.ReadNBitsAsUint32BE(nBits) 320 | if err != nil { 321 | return 0, err 322 | } 323 | 324 | //fmt.Printf("v == %#08x\n", v) 325 | msb := uint32(1) << (nBits - 1) 326 | //fmt.Printf("msb == %#08x\n", msb) 327 | 328 | if (v & msb) == 0 { 329 | return int32(v), nil 330 | } 331 | 332 | f := 0xffffffff & ^(msb - 1) 333 | //fmt.Printf("f ==%#08x\n", f) 334 | //fmt.Printf("f|v ==%#08x\n", f|v) 335 | return int32(f | v), nil 336 | } 337 | 338 | // ReadNBitsAsUint64BE reads `nBits` bits as a big endian unsigned integer from the bit stream and returns it in uint64 (LSB aligned). 339 | // `nBits` must be less than or equal to 64, otherwise returns an error. 340 | // If `nBits` == 0, this function always returns 0. 341 | func (r *Reader) ReadNBitsAsUint64BE(nBits uint8) (uint64, error) { 342 | if nBits == 0 { 343 | return 0, nil 344 | } 345 | 346 | if nBits <= 32 { 347 | v, err := r.ReadNBitsAsUint32BE(nBits) 348 | return uint64(v), err 349 | } 350 | 351 | if nBits > 64 { 352 | return 0, errors.New("nBits too large for uint64") 353 | } 354 | 355 | err := r.fillBufIfNeeded() 356 | if err != nil { 357 | return 0, err 358 | } 359 | 360 | // remaining bits in current byte 361 | rb := r.currBitIndex + 1 362 | 363 | // 64bit value may be distributed in 9 bytes 364 | nBits1 := rb 365 | nBits2 := uint8(8) 366 | nBits3 := uint8(8) 367 | nBits4 := uint8(8) 368 | nBits5 := nBits - rb - 24 369 | nBits6 := uint8(0) 370 | nBits7 := uint8(0) 371 | nBits8 := uint8(0) 372 | nBits9 := uint8(0) 373 | if nBits5 > 8 { 374 | nBits6 = nBits5 - 8 375 | if nBits6 > 8 { 376 | nBits7 = nBits6 - 8 377 | if nBits7 > 8 { 378 | nBits8 = nBits7 - 8 379 | if nBits8 > 8 { 380 | nBits9 = nBits8 - 8 381 | nBits8 = 8 382 | } 383 | nBits7 = 8 384 | } 385 | nBits6 = 8 386 | } 387 | nBits5 = 8 388 | } 389 | 390 | b1 := r.mustReadNBitsInCurrentByte(nBits1) 391 | b2, err := r.ReadNBitsAsUint8(nBits2) 392 | if err != nil { 393 | return 0, err 394 | } 395 | b3, err := r.ReadNBitsAsUint8(nBits3) 396 | if err != nil { 397 | return 0, err 398 | } 399 | b4, err := r.ReadNBitsAsUint8(nBits4) 400 | if err != nil { 401 | return 0, err 402 | } 403 | b5, err := r.ReadNBitsAsUint8(nBits5) 404 | if err != nil { 405 | return 0, err 406 | } 407 | b6, err := r.ReadNBitsAsUint8(nBits6) 408 | if err != nil { 409 | return 0, err 410 | } 411 | b7, err := r.ReadNBitsAsUint8(nBits7) 412 | if err != nil { 413 | return 0, err 414 | } 415 | b8, err := r.ReadNBitsAsUint8(nBits8) 416 | if err != nil { 417 | return 0, err 418 | } 419 | b9, err := r.ReadNBitsAsUint8(nBits9) 420 | if err != nil { 421 | return 0, err 422 | } 423 | 424 | return (uint64(b1) << (nBits2 + nBits3 + nBits4 + nBits5 + nBits6 + nBits7 + nBits8 + nBits9)) | 425 | (uint64(b2) << (nBits3 + nBits4 + nBits5 + nBits6 + nBits7 + nBits8 + nBits9)) | 426 | (uint64(b3) << (nBits4 + nBits5 + nBits6 + nBits7 + nBits8 + nBits9)) | 427 | (uint64(b4) << (nBits5 + nBits6 + nBits7 + nBits8 + nBits9)) | 428 | (uint64(b5) << (nBits6 + nBits7 + nBits8 + nBits9)) | 429 | (uint64(b6) << (nBits7 + nBits8 + nBits9)) | 430 | (uint64(b7) << (nBits8 + nBits9)) | 431 | (uint64(b8) << (nBits9)) | 432 | uint64(b9), nil 433 | } 434 | 435 | // ReadUint64BE reads 64 bits as a big endian unsigned integer from the bit stream and returns it in uint64. 436 | func (r *Reader) ReadUint64BE() (uint64, error) { 437 | return r.ReadNBitsAsUint64BE(64) 438 | } 439 | 440 | // ReadOptions is a set of options to read bits from the bit stream. 441 | type ReadOptions struct { 442 | AlignRight bool // If true, returned value will be aligned to right (default: align to left) 443 | PadOne bool // If true, returned value will be padded with '1' instead of '0' (default: pad with '0') 444 | } 445 | 446 | // ReadNBits reads `nBits` bits from the bit stream and returns it as a slice of bytes. 447 | // If `nBits` == 0, this function always returns nil. 448 | func (r *Reader) ReadNBits(nBits uint8, opt *ReadOptions) ([]byte, error) { 449 | if nBits == 0 { 450 | return nil, nil 451 | } 452 | 453 | err := r.fillBufIfNeeded() 454 | if err != nil { 455 | return nil, err 456 | } 457 | 458 | padOne := (opt != nil && opt.PadOne) 459 | alignRight := (opt != nil && opt.AlignRight) 460 | 461 | maxByteLen := (nBits / 8) + 1 462 | result := make([]byte, 0, maxByteLen) 463 | 464 | // remaining bits in current byte 465 | rb := r.currBitIndex + 1 466 | var bitsToRead uint8 467 | if nBits <= rb { 468 | bitsToRead = nBits 469 | } else { 470 | bitsToRead = rb 471 | } 472 | 473 | tempByte := r.mustReadNBitsInCurrentByte(bitsToRead) 474 | tempByte = tempByte << (8 - bitsToRead) // left align 475 | tempBit := bitsToRead 476 | nBits -= bitsToRead 477 | 478 | if tempBit == 8 { 479 | result = append(result, tempByte) 480 | tempByte = 0 481 | tempBit = 0 482 | } 483 | 484 | for nBits >= 8 { 485 | err := r.fillBufIfNeeded() 486 | if err != nil { 487 | return nil, err 488 | } 489 | 490 | bitsToRead = 8 491 | b := r.mustReadNBitsInCurrentByte(bitsToRead) 492 | b1 := b >> tempBit 493 | b2 := b << (8 - tempBit) 494 | 495 | tempByte = tempByte | b1 496 | result = append(result, tempByte) 497 | tempByte = b2 498 | 499 | nBits -= 8 500 | } 501 | 502 | if nBits > 0 { 503 | err := r.fillBufIfNeeded() 504 | if err != nil { 505 | return nil, err 506 | } 507 | 508 | bitsToRead = nBits 509 | b := r.mustReadNBitsInCurrentByte(bitsToRead) 510 | b1 := b >> (bitsToRead - (8 - tempBit)) // wants to have (8 - tempBit) bits from b. b has bitsToRead bits 511 | b2 := b << (8 - (bitsToRead - (8 - tempBit))) // wants to have (bitsToRead - ) left aligned. 512 | 513 | tempByte = tempByte | b1 514 | result = append(result, tempByte) 515 | 516 | if nBits > (8 - tempBit) { 517 | if padOne { 518 | b2 = b2 | (0xff >> tempBit) 519 | } 520 | result = append(result, b2) 521 | } 522 | } else { 523 | if tempBit > 0 { 524 | result = append(result, tempByte) 525 | } 526 | } 527 | 528 | if alignRight { 529 | return nil, errors.New("not implemented yet") 530 | } 531 | 532 | return result, nil 533 | } 534 | -------------------------------------------------------------------------------- /reader_test.go: -------------------------------------------------------------------------------- 1 | package bitstream 2 | 3 | import ( 4 | "bytes" 5 | "crypto/rand" 6 | "reflect" 7 | "testing" 8 | ) 9 | 10 | type indecies struct { 11 | BitIndex uint8 12 | ByteIndex uint 13 | } 14 | 15 | func TestForwardIndecies(t *testing.T) { 16 | testData := []struct { 17 | Name string 18 | Data []byte 19 | Start indecies 20 | NumBitsToForward uint8 21 | End indecies 22 | }{ 23 | { 24 | Name: "pattern 1", 25 | Data: []byte{0x01}, // b7654 3210 26 | Start: indecies{BitIndex: 7, ByteIndex: 0}, // ^ 27 | NumBitsToForward: 1, 28 | End: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 29 | }, 30 | { 31 | Name: "pattern 2", 32 | Data: []byte{0x02}, // b7654 3210 33 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 34 | NumBitsToForward: 2, 35 | End: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 36 | }, 37 | { 38 | Name: "pattern 3", 39 | Data: []byte{0x04}, // b7654 3210 40 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 41 | NumBitsToForward: 4, 42 | End: indecies{BitIndex: 0, ByteIndex: 0}, // ^ 43 | }, 44 | { 45 | Name: "pattern 5", 46 | Data: []byte{0x05}, // b7654 3210 | 47 | Start: indecies{BitIndex: 0, ByteIndex: 0}, // ^ 48 | NumBitsToForward: 1, 49 | End: indecies{BitIndex: 7, ByteIndex: 1}, // ^ 50 | }, 51 | { 52 | Name: "pattern 6", 53 | Data: []byte{0x06, 0x06}, // b7654 3210 | 7654 3210 54 | Start: indecies{BitIndex: 0, ByteIndex: 0}, // ^ 55 | NumBitsToForward: 2, 56 | End: indecies{BitIndex: 6, ByteIndex: 1}, // ^ 57 | }, 58 | { 59 | Name: "pattern 7", 60 | Data: []byte{0x07, 0x07, 0x07}, // b7654 3210 | 7654 3210 | 7654 3210 61 | Start: indecies{BitIndex: 1, ByteIndex: 0}, // ^ 62 | NumBitsToForward: 10, 63 | End: indecies{BitIndex: 7, ByteIndex: 2}, // ^ 64 | }, 65 | } 66 | 67 | for _, data := range testData { 68 | data := data // capture 69 | t.Run(data.Name, func(t *testing.T) { 70 | //t.Parallel() 71 | 72 | r := NewReader(bytes.NewReader(data.Data), nil) 73 | r.fillBuf() 74 | r.currBitIndex = data.Start.BitIndex 75 | r.currByteIndex = data.Start.ByteIndex 76 | 77 | r.forwardIndecies(data.NumBitsToForward) 78 | 79 | if data.End.BitIndex != r.currBitIndex { 80 | t.Fatalf("\nunexpected bit index\nExpected: %+v\nActual: %+v\n", data.End.BitIndex, r.currBitIndex) 81 | } 82 | if data.End.ByteIndex != r.currByteIndex { 83 | t.Fatalf("\nunexpected byte index\nExpected: %+v\nActual: %+v\n", data.End.ByteIndex, r.currByteIndex) 84 | } 85 | }) 86 | } 87 | } 88 | 89 | func TestReadBit(t *testing.T) { 90 | testData := []struct { 91 | Name string 92 | Data []byte 93 | ExpectedBits []byte 94 | ExpectedConsumedBytes uint 95 | }{ 96 | { 97 | Name: "pattern 1", 98 | Data: []byte{0xaa}, 99 | ExpectedBits: []byte{1, 0, 1, 0, 1, 0, 1, 0}, 100 | ExpectedConsumedBytes: 1, 101 | }, 102 | { 103 | Name: "pattern 2", 104 | Data: []byte{0x55, 0x12}, 105 | ExpectedBits: []byte{0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0}, 106 | ExpectedConsumedBytes: 2, 107 | }, 108 | { 109 | Name: "pattern 3", 110 | Data: []byte{0x55, 0x12}, 111 | ExpectedBits: []byte{0, 1, 0, 1, 0, 1, 0, 1, 0}, // reads only 9 bits 112 | ExpectedConsumedBytes: 2, 113 | }, 114 | } 115 | 116 | for _, data := range testData { 117 | data := data // capture 118 | t.Run(data.Name, func(t *testing.T) { 119 | //t.Parallel() 120 | 121 | r := NewReader(bytes.NewReader(data.Data), nil) 122 | for i, expectedBit := range data.ExpectedBits { 123 | actualBit, err := r.ReadBit() 124 | if err != nil { 125 | t.Fatalf("unexpected error: %+v\n", err) 126 | } 127 | if expectedBit != actualBit { 128 | t.Fatalf("\nbit %d\nExpected: %+v\nActual: %+v\n", i, expectedBit, actualBit) 129 | } 130 | } 131 | 132 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 133 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 134 | } 135 | 136 | if len(data.ExpectedBits)%8 == 0 { 137 | _, err := r.ReadBit() 138 | if err == nil { 139 | t.Fatal("error should occur but no error\n") 140 | } 141 | } 142 | }) 143 | } 144 | } 145 | 146 | // https://dave.cheney.net/2013/06/30/how-to-write-benchmarks-in-go 147 | var toEliminateCompilerOptimizationByte byte 148 | var toEliminateCompilerOptimizationUint16 uint16 149 | var toEliminateCompilerOptimizationUint32 uint32 150 | var toEliminateCompilerOptimizationUint64 uint64 151 | 152 | func BenchmarkReadBit(b *testing.B) { 153 | var v byte 154 | r := NewReader(rand.Reader, nil) 155 | for n := 0; n < b.N; n++ { 156 | v, _ = r.ReadBit() 157 | } 158 | toEliminateCompilerOptimizationByte = v 159 | } 160 | 161 | func TestReadNBitsAsUint8(t *testing.T) { 162 | testData := []struct { 163 | Name string 164 | Data []byte 165 | Start indecies 166 | NBits uint8 167 | Expected uint8 168 | ExpectedConsumedBytes uint 169 | }{ 170 | { 171 | Name: "pattern 1", // b7654 3210 172 | Data: []byte{0x12}, // 0001 0010 173 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 174 | NBits: 4, // ^^^ ^ 175 | Expected: 2, // 001 0 => 2 176 | ExpectedConsumedBytes: 1, 177 | }, 178 | { 179 | Name: "pattern 2", // b7654 3210 180 | Data: []byte{0x12}, // 0001 0010 181 | Start: indecies{BitIndex: 7, ByteIndex: 0}, // ^ 182 | NBits: 8, // ^^^^ ^^^^ 183 | Expected: 0x12, // 0001 0010 => 0x12 184 | ExpectedConsumedBytes: 1, 185 | }, 186 | { 187 | Name: "pattern 3", // b7654 3210 188 | Data: []byte{0x12}, // 0001 0010 189 | Start: indecies{BitIndex: 0, ByteIndex: 0}, // ^ 190 | NBits: 1, // ^ 191 | Expected: 0, // 0 => 0 192 | ExpectedConsumedBytes: 1, 193 | }, 194 | { 195 | Name: "pattern 4", // b7654 3210 7654 3210 196 | Data: []byte{0x12, 0x34}, // 0001 0010 0011 0100 197 | Start: indecies{BitIndex: 0, ByteIndex: 0}, // ^ 198 | NBits: 6, // ^ ^^^^ ^ 199 | Expected: 6, // 0 0011 0 => 6 200 | ExpectedConsumedBytes: 2, 201 | }, 202 | { 203 | Name: "pattern 5", // b7654 3210 7654 3210 204 | Data: []byte{0xff, 0xff}, // 1111 1111 1111 1111 205 | Start: indecies{BitIndex: 3, ByteIndex: 0}, // ^ 206 | NBits: 8, // ^^^^ ^^^^ 207 | Expected: 0xff, // 1111 1111 => 0xff 208 | ExpectedConsumedBytes: 2, 209 | }, 210 | { 211 | Name: "pattern 6", // b7654 3210 7654 3210 212 | Data: []byte{0xaa, 0x55}, // 1010 1010 0101 0101 213 | Start: indecies{BitIndex: 2, ByteIndex: 0}, // ^ 214 | NBits: 7, // ^^^ ^^^^ 215 | Expected: 0x25, // 010 0101 => 0x25 216 | ExpectedConsumedBytes: 2, 217 | }, 218 | { 219 | Name: "pattern 7", // b7654 3210 7654 3210 220 | Data: []byte{0x98, 0xb1}, // 1001 1000 1011 0001 221 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 222 | NBits: 7, // ^ ^^^^ ^^ 223 | Expected: 0x62, // 1 1000 10 => 110 0010 => 0x62 224 | ExpectedConsumedBytes: 2, 225 | }, 226 | } 227 | 228 | for _, data := range testData { 229 | data := data // capture 230 | t.Run(data.Name, func(t *testing.T) { 231 | //t.Parallel() 232 | 233 | r := NewReader(bytes.NewReader(data.Data), nil) 234 | r.fillBuf() 235 | r.currBitIndex = data.Start.BitIndex 236 | r.currByteIndex = data.Start.ByteIndex 237 | 238 | v, err := r.ReadNBitsAsUint8(data.NBits) 239 | if err != nil { 240 | t.Fatalf("unexpected error: %+v\n", err) 241 | } 242 | if data.Expected != v { 243 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", data.Expected, v) 244 | } 245 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 246 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 247 | } 248 | }) 249 | } 250 | } 251 | 252 | func benchmarkReadNBitsAsUint8(b *testing.B, nBits uint8) { 253 | var v byte 254 | r := NewReader(rand.Reader, nil) 255 | for n := 0; n < b.N; n++ { 256 | v, _ = r.ReadNBitsAsUint8(nBits) 257 | } 258 | toEliminateCompilerOptimizationByte = v 259 | } 260 | 261 | func BenchmarkRead1BitAsUint8(b *testing.B) { 262 | benchmarkReadNBitsAsUint8(b, 1) 263 | } 264 | 265 | func BenchmarkRead2BitsAsUint8(b *testing.B) { 266 | benchmarkReadNBitsAsUint8(b, 2) 267 | } 268 | 269 | func BenchmarkRead3BitsAsUint8(b *testing.B) { 270 | benchmarkReadNBitsAsUint8(b, 3) 271 | } 272 | 273 | func BenchmarkRead4BitsAsUint8(b *testing.B) { 274 | benchmarkReadNBitsAsUint8(b, 4) 275 | } 276 | 277 | func BenchmarkRead5BitsAsUint8(b *testing.B) { 278 | benchmarkReadNBitsAsUint8(b, 5) 279 | } 280 | 281 | func BenchmarkRead6BitsAsUint8(b *testing.B) { 282 | benchmarkReadNBitsAsUint8(b, 6) 283 | } 284 | 285 | func BenchmarkRead7BitsAsUint8(b *testing.B) { 286 | benchmarkReadNBitsAsUint8(b, 7) 287 | } 288 | 289 | func BenchmarkRead8BitsAsUint8(b *testing.B) { 290 | benchmarkReadNBitsAsUint8(b, 8) 291 | } 292 | 293 | func TestReadNBitsAsUint16BE(t *testing.T) { 294 | testData := []struct { 295 | Name string 296 | Data []byte 297 | Start indecies 298 | NBits uint8 299 | Expected uint16 300 | ExpectedConsumedBytes uint 301 | }{ 302 | { 303 | Name: "pattern 1", // b7654 3210 | 7654 3210 | 7654 3210 304 | Data: []byte{0x12, 0x34, 0x56}, // 0001 0010 | 0011 0100 | 0101 0110 305 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 306 | NBits: 4, // ^^^ ^ 307 | Expected: 2, // 001 0 => 2 308 | ExpectedConsumedBytes: 1, 309 | }, 310 | { 311 | Name: "pattern 2", // b7654 3210 | 7654 3210 | 7654 3210 312 | Data: []byte{0x12, 0x34, 0x56}, // 0001 0010 | 0011 0100 | 0101 0110 313 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 314 | NBits: 9, // ^^^ ^^^^ ^^ 315 | Expected: 0x48, // 001 0010 00 => 0x48 316 | ExpectedConsumedBytes: 2, 317 | }, 318 | { 319 | Name: "pattern 3", // b7654 3210 | 7654 3210 | 7654 3210 320 | Data: []byte{0x12, 0x34, 0x56}, // 0001 0010 | 0011 0100 | 0101 0110 321 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 322 | NBits: 16, // ^^^ ^^^^ ^^^^ ^^^^ ^ 323 | Expected: 0x2468, // 001 0010 0011 0100 0 => 0x2468 324 | ExpectedConsumedBytes: 3, 325 | }, 326 | { 327 | Name: "pattern 4", // b7654 3210 | 7654 3210 | 7654 3210 328 | Data: []byte{0x12, 0x34, 0x56}, // 0001 0010 | 0011 0100 | 0101 0110 329 | Start: indecies{BitIndex: 7, ByteIndex: 1}, // ^ 330 | NBits: 16, // ^^^^ ^^^^ ^^^^ ^^^^ 331 | Expected: 0x3456, // 0011 0100 0101 0110 0 => 0x3456 332 | ExpectedConsumedBytes: 2, 333 | }, 334 | } 335 | 336 | for _, data := range testData { 337 | data := data // capture 338 | t.Run(data.Name, func(t *testing.T) { 339 | //t.Parallel() 340 | 341 | r := NewReader(bytes.NewReader(data.Data), nil) 342 | r.fillBuf() 343 | r.currBitIndex = data.Start.BitIndex 344 | r.currByteIndex = data.Start.ByteIndex 345 | 346 | v, err := r.ReadNBitsAsUint16BE(data.NBits) 347 | if err != nil { 348 | t.Fatalf("unexpected error: %+v\n", err) 349 | } 350 | if data.Expected != v { 351 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", data.Expected, v) 352 | } 353 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 354 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 355 | } 356 | }) 357 | } 358 | } 359 | 360 | func benchmarkReadNBitsAsUint16BE(b *testing.B, nBits uint8) { 361 | var v uint16 362 | r := NewReader(rand.Reader, nil) 363 | for n := 0; n < b.N; n++ { 364 | v, _ = r.ReadNBitsAsUint16BE(nBits) 365 | } 366 | toEliminateCompilerOptimizationUint16 = v 367 | } 368 | 369 | func BenchmarkRead1BitAsUint16BE(b *testing.B) { 370 | benchmarkReadNBitsAsUint16BE(b, 1) 371 | } 372 | 373 | func BenchmarkRead2BitsAsUint16BE(b *testing.B) { 374 | benchmarkReadNBitsAsUint16BE(b, 2) 375 | } 376 | 377 | func BenchmarkRead9BitsAsUint16BE(b *testing.B) { 378 | benchmarkReadNBitsAsUint16BE(b, 9) 379 | } 380 | 381 | func BenchmarkRead10BitsAsUint16BE(b *testing.B) { 382 | benchmarkReadNBitsAsUint16BE(b, 10) 383 | } 384 | 385 | func BenchmarkRead11BitsAsUint16BE(b *testing.B) { 386 | benchmarkReadNBitsAsUint16BE(b, 11) 387 | } 388 | 389 | func BenchmarkRead12BitsAsUint16BE(b *testing.B) { 390 | benchmarkReadNBitsAsUint16BE(b, 12) 391 | } 392 | 393 | func BenchmarkRead13BitsAsUint16BE(b *testing.B) { 394 | benchmarkReadNBitsAsUint16BE(b, 13) 395 | } 396 | 397 | func BenchmarkRead14BitsAsUint16BE(b *testing.B) { 398 | benchmarkReadNBitsAsUint16BE(b, 14) 399 | } 400 | 401 | func BenchmarkRead15BitsAsUint16BE(b *testing.B) { 402 | benchmarkReadNBitsAsUint16BE(b, 15) 403 | } 404 | 405 | func BenchmarkRead16BitsAsUint16BE(b *testing.B) { 406 | benchmarkReadNBitsAsUint16BE(b, 16) 407 | } 408 | 409 | func TestReadNBitsAsUint32BE(t *testing.T) { 410 | testData := []struct { 411 | Name string 412 | Data []byte 413 | Start indecies 414 | NBits uint8 415 | Expected uint32 416 | ExpectedConsumedBytes uint 417 | }{ 418 | { 419 | Name: "pattern 1", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 420 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 421 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 422 | NBits: 4, // ^^^ ^ 423 | Expected: 2, // 001 0 => 2 424 | ExpectedConsumedBytes: 1, 425 | }, 426 | { 427 | Name: "pattern 2", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 428 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 429 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 430 | NBits: 9, // ^^^ ^^^^ ^^ 431 | Expected: 0x48, // 001 0010 00 => 0x48 432 | ExpectedConsumedBytes: 2, 433 | }, 434 | { 435 | Name: "pattern 3", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 436 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 437 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 438 | NBits: 17, // ^^^ ^^^^ ^^^^ ^^^^ ^^ 439 | Expected: 0x48D1, // 001 0010 0011 0100 01 => 0x48D1 440 | ExpectedConsumedBytes: 3, 441 | }, 442 | { 443 | Name: "pattern 4", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 444 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 445 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 446 | NBits: 24, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^ 447 | Expected: 0x2468AC, // 001 0010 0011 0100 0101 0110 0 => 0010 0100 0110 1000 1010 1100 => 0x2468AC 448 | ExpectedConsumedBytes: 4, 449 | }, 450 | { 451 | Name: "pattern 5", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 452 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 453 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 454 | NBits: 32, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^ 455 | Expected: 0x2468ACF1, // 001 0010 0011 0100 0101 0110 0111 1000 1 => 0010 0100 0110 1000 1010 1100 1111 0001 => 0x2468ACF1 456 | ExpectedConsumedBytes: 5, 457 | }, 458 | { 459 | Name: "pattern 6", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 460 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 461 | Start: indecies{BitIndex: 7, ByteIndex: 1}, // ^ 462 | NBits: 32, // ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 463 | Expected: 0x3456789A, // 0011 0100 0101 0110 0111 1000 1001 1010 => 0x3456789A 464 | ExpectedConsumedBytes: 4, 465 | }, 466 | } 467 | 468 | for _, data := range testData { 469 | data := data // capture 470 | t.Run(data.Name, func(t *testing.T) { 471 | //t.Parallel() 472 | 473 | r := NewReader(bytes.NewReader(data.Data), nil) 474 | r.fillBuf() 475 | r.currBitIndex = data.Start.BitIndex 476 | r.currByteIndex = data.Start.ByteIndex 477 | 478 | v, err := r.ReadNBitsAsUint32BE(data.NBits) 479 | if err != nil { 480 | t.Fatalf("unexpected error: %+v\n", err) 481 | } 482 | if data.Expected != v { 483 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", data.Expected, v) 484 | } 485 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 486 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 487 | } 488 | }) 489 | } 490 | } 491 | 492 | func benchmarkReadNBitsAsUint32BE(b *testing.B, nBits uint8) { 493 | var v uint32 494 | r := NewReader(rand.Reader, nil) 495 | for n := 0; n < b.N; n++ { 496 | v, _ = r.ReadNBitsAsUint32BE(nBits) 497 | } 498 | toEliminateCompilerOptimizationUint32 = v 499 | } 500 | 501 | func BenchmarkRead1BitAsUint32BE(b *testing.B) { 502 | benchmarkReadNBitsAsUint32BE(b, 1) 503 | } 504 | 505 | func BenchmarkRead2BitsAsUint32BE(b *testing.B) { 506 | benchmarkReadNBitsAsUint32BE(b, 2) 507 | } 508 | 509 | func BenchmarkRead9BitsAsUint32BE(b *testing.B) { 510 | benchmarkReadNBitsAsUint32BE(b, 9) 511 | } 512 | 513 | func BenchmarkRead10BitsAsUint32BE(b *testing.B) { 514 | benchmarkReadNBitsAsUint32BE(b, 10) 515 | } 516 | 517 | func BenchmarkRead11BitsAsUint32BE(b *testing.B) { 518 | benchmarkReadNBitsAsUint32BE(b, 17) 519 | } 520 | 521 | func BenchmarkRead12BitsAsUint32BE(b *testing.B) { 522 | benchmarkReadNBitsAsUint32BE(b, 18) 523 | } 524 | func BenchmarkRead15BitsAsUint32BE(b *testing.B) { 525 | benchmarkReadNBitsAsUint32BE(b, 31) 526 | } 527 | 528 | func BenchmarkRead32BitsAsUint32BE(b *testing.B) { 529 | benchmarkReadNBitsAsUint32BE(b, 32) 530 | } 531 | 532 | func TestReadNBitsAsInt32BE(t *testing.T) { 533 | testData := []struct { 534 | Name string 535 | Data []byte 536 | Start indecies 537 | NBits uint8 538 | Expected int32 539 | ExpectedConsumedBytes uint 540 | }{ 541 | { 542 | Name: "pattern 1", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 543 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 544 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 545 | NBits: 4, // ^^^ ^ 546 | Expected: 2, // 001 0 => 2 547 | ExpectedConsumedBytes: 1, 548 | }, 549 | { 550 | Name: "pattern 2", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 551 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 552 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 553 | NBits: 8, // ^ ^^^^ ^^^ 554 | Expected: -111, // 1 0010 001 => 0x91 => -(^0x11[7bit] + 1) => -(0x6E + 1) => -111 555 | ExpectedConsumedBytes: 2, 556 | }, 557 | { 558 | Name: "pattern 3", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 559 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 560 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 561 | NBits: 9, // ^ ^^^^ ^^^^ 562 | Expected: -221, // 1 0010 0011 => 0x123 => -(^0x23[8bit] + 1) => -(0xDC + 1) => -221 563 | ExpectedConsumedBytes: 2, 564 | }, 565 | { 566 | Name: "pattern 4", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 567 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 568 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 569 | NBits: 17, // ^ ^^^^ ^^^^ ^^^^ ^^^^ 570 | Expected: -56507, // 1 0010 0011 0100 0101 => 0x12345 => -(^0x2345[16bit] + 1) => -(0xDCBA + 1) => -56507 571 | ExpectedConsumedBytes: 3, 572 | }, 573 | { 574 | Name: "pattern 5", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 575 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 576 | Start: indecies{BitIndex: 1, ByteIndex: 0}, // ^ 577 | NBits: 24, // ^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^ 578 | Expected: -7531106, // 10 0011 0100 0101 0110 0111 10 => 1000 1101 0001 0101 1001 1110 => 0x8D159E => -(^0x0D159E[23bit] + 1) => -(0x72EA61 + 1) => -7531106 579 | ExpectedConsumedBytes: 4, 580 | }, 581 | { 582 | Name: "pattern 6", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 583 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 584 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 585 | NBits: 32, // ^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^ 586 | Expected: -1851608124, // 1 0010 0011 0100 0101 0110 0111 1000 100 => 1001 0001 1010 0010 1011 0011 1100 0100 => 0x91A2B3C4 => -(^0x11A2B3C4[31bit] + 1) => -(0x6E5D4C3B + 1) => -1851608124 587 | ExpectedConsumedBytes: 5, 588 | }, 589 | { 590 | Name: "pattern 7", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 591 | Data: []byte{0x12, 0xB4, 0x56, 0x78, 0x9a}, // 0001 0010 | 1011 0100 | 0101 0110 | 0111 1000 | 1001 1010 592 | Start: indecies{BitIndex: 7, ByteIndex: 1}, // ^ 593 | NBits: 32, // ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 594 | Expected: -1269401446, // 1011 0100 0101 0110 0111 1000 1001 1010 => 0xB456789A => -(^0x3456789A[31bit] + 1) => -(0x4BA98765 + 1) => -1269401446 595 | ExpectedConsumedBytes: 4, 596 | }, 597 | { 598 | Name: "pattern 8", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 599 | Data: []byte{0x38, 0xb6, 0x0b, 0x98, 0x9d}, // 0011 1000 | 1011 0110 | 0000 1011 | 1001 1000 | 1001 1101 600 | Start: indecies{BitIndex: 5, ByteIndex: 0}, // ^ 601 | NBits: 24, // ^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^ 602 | Expected: -1910738, // 11 1000 1011 0110 0000 1011 10 => 1110 0010 1101 1000 0010 1110 => 0xE2D82E => -(^0x62D82E[23bit] + 1) => -(0x1D27D1 + 1) => -1910738 603 | ExpectedConsumedBytes: 4, 604 | }, 605 | } 606 | 607 | for _, data := range testData { 608 | data := data // capture 609 | t.Run(data.Name, func(t *testing.T) { 610 | //t.Parallel() 611 | 612 | r := NewReader(bytes.NewReader(data.Data), nil) 613 | r.fillBuf() 614 | r.currBitIndex = data.Start.BitIndex 615 | r.currByteIndex = data.Start.ByteIndex 616 | 617 | v, err := r.ReadNBitsAsInt32BE(data.NBits) 618 | if err != nil { 619 | t.Fatalf("unexpected error: %+v\n", err) 620 | } 621 | if data.Expected != v { 622 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", data.Expected, v) 623 | } 624 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 625 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 626 | } 627 | }) 628 | } 629 | } 630 | 631 | func TestReadNBitsAsUint64BE(t *testing.T) { 632 | testData := []struct { 633 | Name string 634 | Data []byte 635 | Start indecies 636 | NBits uint8 637 | Expected uint64 638 | ExpectedConsumedBytes uint 639 | }{ 640 | { 641 | Name: "pattern 1", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 642 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 643 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 644 | NBits: 4, // ^^^ ^ 645 | Expected: 2, // 001 0 => 2 646 | ExpectedConsumedBytes: 1, 647 | }, 648 | { 649 | Name: "pattern 2", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 650 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 651 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 652 | NBits: 9, // ^^^ ^^^^ ^^ 653 | Expected: 0x48, // 001 0010 00 => 0x48 654 | ExpectedConsumedBytes: 2, 655 | }, 656 | { 657 | Name: "pattern 3", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 658 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 659 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 660 | NBits: 17, // ^^^ ^^^^ ^^^^ ^^^^ ^^ 661 | Expected: 0x48D1, // 001 0010 0011 0100 01 => 0x48D1 662 | ExpectedConsumedBytes: 3, 663 | }, 664 | { 665 | Name: "pattern 4", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 666 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 667 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 668 | NBits: 33, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^ 669 | Expected: 0x48D159E2, // 001 0010 0011 0100 0101 0110 0111 1000 10 => 0 0100 1000 1101 0001 0101 1001 1110 0010 => 0x48D159E2 670 | ExpectedConsumedBytes: 5, 671 | }, 672 | { 673 | Name: "pattern 5", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 674 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 675 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 676 | NBits: 42, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^ 677 | Expected: 0x91A2B3C4D5, // 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 101 => 00 1001 0001 1010 0010 1011 0011 1100 0100 1101 0101 => 0x91A2B3C4D5 678 | ExpectedConsumedBytes: 6, 679 | }, 680 | { 681 | Name: "pattern 6", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 682 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 683 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 684 | NBits: 51, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 685 | Expected: 0x123456789ABCD, // 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 => 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 => 0x123456789ABCD 686 | ExpectedConsumedBytes: 7, 687 | }, 688 | { 689 | Name: "pattern 7", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 690 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 691 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 692 | NBits: 60, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^ 693 | Expected: 0x2468ACF13579BDE, // 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111 0 => 0010 0100 0110 1000 1010 1100 1111 0001 0011 0101 0111 1001 1011 1101 1110 => 0x2468ACF13578BDE 694 | ExpectedConsumedBytes: 8, 695 | }, 696 | { 697 | Name: "pattern 6", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 698 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 699 | Start: indecies{BitIndex: 7, ByteIndex: 1}, // ^ 700 | NBits: 64, // ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 701 | Expected: 0x3456789ABCDEF012, // 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111 0000 0001 0010 => 0x3456789ABCDEF012 702 | ExpectedConsumedBytes: 8, 703 | }, 704 | } 705 | 706 | for _, data := range testData { 707 | data := data // capture 708 | t.Run(data.Name, func(t *testing.T) { 709 | //t.Parallel() 710 | 711 | r := NewReader(bytes.NewReader(data.Data), nil) 712 | r.fillBuf() 713 | r.currBitIndex = data.Start.BitIndex 714 | r.currByteIndex = data.Start.ByteIndex 715 | 716 | v, err := r.ReadNBitsAsUint64BE(data.NBits) 717 | if err != nil { 718 | t.Fatalf("unexpected error: %+v\n", err) 719 | } 720 | if data.Expected != v { 721 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", data.Expected, v) 722 | } 723 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 724 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 725 | } 726 | }) 727 | } 728 | } 729 | 730 | func benchmarkReadNBitsAsUint64BE(b *testing.B, nBits uint8) { 731 | var v uint64 732 | r := NewReader(rand.Reader, nil) 733 | for n := 0; n < b.N; n++ { 734 | v, _ = r.ReadNBitsAsUint64BE(nBits) 735 | } 736 | toEliminateCompilerOptimizationUint64 = v 737 | } 738 | 739 | func BenchmarkRead1BitAsUint64BE(b *testing.B) { 740 | benchmarkReadNBitsAsUint64BE(b, 1) 741 | } 742 | 743 | func BenchmarkRead2BitsAsUint64BE(b *testing.B) { 744 | benchmarkReadNBitsAsUint64BE(b, 2) 745 | } 746 | 747 | func BenchmarkRead9BitsAsUint64BE(b *testing.B) { 748 | benchmarkReadNBitsAsUint64BE(b, 9) 749 | } 750 | 751 | func BenchmarkRead10BitsAsUint64BE(b *testing.B) { 752 | benchmarkReadNBitsAsUint64BE(b, 10) 753 | } 754 | 755 | func BenchmarkRead11BitsAsUint64BE(b *testing.B) { 756 | benchmarkReadNBitsAsUint64BE(b, 17) 757 | } 758 | 759 | func BenchmarkRead12BitsAsUint64BE(b *testing.B) { 760 | benchmarkReadNBitsAsUint64BE(b, 18) 761 | } 762 | func BenchmarkRead15BitsAsUint64BE(b *testing.B) { 763 | benchmarkReadNBitsAsUint64BE(b, 31) 764 | } 765 | 766 | func BenchmarkRead64BitsAsUint64BE(b *testing.B) { 767 | benchmarkReadNBitsAsUint64BE(b, 64) 768 | } 769 | 770 | func TestReadNBits(t *testing.T) { 771 | testData := []struct { 772 | Name string 773 | Data []byte 774 | Start indecies 775 | NBits uint8 776 | AlignRight bool 777 | PadOne bool 778 | Expected []byte 779 | ExpectedConsumedBytes uint 780 | }{ 781 | { 782 | Name: "pattern 1", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 783 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 784 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 785 | NBits: 4, // ^^^ ^ 786 | Expected: []byte{0x20}, // 001 0 => 0010 0000 => 0x20 787 | ExpectedConsumedBytes: 1, 788 | }, 789 | { 790 | Name: "pattern 2", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 791 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 792 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 793 | NBits: 8, // ^^^ ^^^^ ^ 794 | Expected: []byte{0x24}, // 001 0010 0 => 00100100 => 0x24 795 | ExpectedConsumedBytes: 2, 796 | }, 797 | { 798 | Name: "pattern 3", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 799 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 800 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 801 | NBits: 9, // ^^^ ^^^^ ^^ 802 | Expected: []byte{0x24, 0x00}, // 001 0010 00 => 00100100 0 => 0x24 0x00 803 | ExpectedConsumedBytes: 2, 804 | }, 805 | { 806 | Name: "pattern 4", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 807 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 808 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 809 | NBits: 16, // ^^^ ^^^^ ^^^^ ^^^^ ^ 810 | Expected: []byte{0x24, 0x68}, // 001 0010 0011 0100 0 => 0010 0100 0110 1000 => 0x24 0x68 811 | ExpectedConsumedBytes: 3, 812 | }, 813 | { 814 | Name: "pattern 5", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 815 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 816 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 817 | NBits: 17, // ^^^ ^^^^ ^^^^ ^^^^ ^^ 818 | Expected: []byte{0x24, 0x68, 0x80}, // 001 0010 0011 0100 01 => 0010 0100 0110 1000 1 => 0x24 0x68 0x80 819 | ExpectedConsumedBytes: 3, 820 | }, 821 | { 822 | Name: "pattern 6", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 823 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 824 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 825 | NBits: 33, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^ 826 | Expected: []byte{0x24, 0x68, 0xac, 0xf1, 0x00}, // 001 0010 0011 0100 0101 0110 0111 1000 10 => 0010 0100 0110 1000 1010 1100 1111 0001 0 => 0x24 0x68 0xac 0xf1 0x00 827 | ExpectedConsumedBytes: 5, 828 | }, 829 | { 830 | Name: "pattern 7", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 831 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 832 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 833 | NBits: 42, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^ 834 | Expected: []byte{0x24, 0x68, 0xac, 0xf1, 0x35, 0x40}, // 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 101 => 0010 0100 0110 1000 1010 1100 1111 0001 0011 0101 01 => 0x24 0x68 0xac 0xf1 0x35 0x40 835 | ExpectedConsumedBytes: 6, 836 | }, 837 | { 838 | Name: "pattern 8", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 839 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 840 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 841 | NBits: 51, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 842 | Expected: []byte{0x24, 0x68, 0xac, 0xf1, 0x35, 0x79, 0xa0}, // 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 => 0010 0100 0110 1000 1010 1100 1111 0001 0011 0101 0111 1001 101 => 0x24 0x68 0xac 0xf1 0x35 0x79 0xa0 843 | ExpectedConsumedBytes: 7, 844 | }, 845 | { 846 | Name: "pattern 9", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 847 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 848 | Start: indecies{BitIndex: 6, ByteIndex: 0}, // ^ 849 | NBits: 60, // ^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^ 850 | Expected: []byte{0x24, 0x68, 0xac, 0xf1, 0x35, 0x79, 0xbd, 0xe0}, // 001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111 0 => 0010 0100 0110 1000 1010 1100 1111 0001 0011 0101 0111 1001 1011 1101 1110 => 0x24 0x68 0xAC 0xF1 0x35 0x79 0xBD 0xE0 851 | ExpectedConsumedBytes: 8, 852 | }, 853 | { 854 | Name: "pattern 10", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 855 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 856 | Start: indecies{BitIndex: 7, ByteIndex: 1}, // ^ 857 | NBits: 64, // ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 858 | Expected: []byte{0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0, 0x12}, // 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111 0000 0001 0010 => 0x34 56 78 9A BC DE F0 12 859 | ExpectedConsumedBytes: 8, 860 | }, 861 | { 862 | Name: "pattern 11", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 863 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 864 | Start: indecies{BitIndex: 5, ByteIndex: 0}, // ^ 865 | NBits: 17, // ^^ ^^^^ ^^^^ ^^^^ ^^^ 866 | Expected: []byte{0x48, 0xd1, 0x00}, // 01 0010 0011 0100 010 => 0100 1000 1101 0001 0 => 0x48 0xD1 0x00 867 | ExpectedConsumedBytes: 3, 868 | }, 869 | { 870 | Name: "pattern 12", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 871 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 872 | Start: indecies{BitIndex: 4, ByteIndex: 0}, // ^ 873 | NBits: 17, // ^ ^^^^ ^^^^ ^^^^ ^^^^ 874 | Expected: []byte{0x91, 0xa2, 0x80}, // 1 0010 0011 0100 0101 => 1001 0001 1010 0010 1 => 0x91 0xa2 0x80 875 | ExpectedConsumedBytes: 3, 876 | }, 877 | { 878 | Name: "pattern 13", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 879 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 880 | Start: indecies{BitIndex: 3, ByteIndex: 0}, // ^ 881 | NBits: 17, // ^^^^ ^^^^ ^^^^ ^^^^ ^ 882 | Expected: []byte{0x23, 0x45, 0x00}, // 0010 0011 0100 0101 0 => 0010 0011 0100 0101 0 => 0x23 0x45 0x00 883 | ExpectedConsumedBytes: 3, 884 | }, 885 | { 886 | Name: "pattern 14", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 887 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 888 | Start: indecies{BitIndex: 2, ByteIndex: 0}, // ^ 889 | NBits: 17, // ^^^ ^^^^ ^^^^ ^^^^ ^^ 890 | Expected: []byte{0x46, 0x8a, 0x80}, // 010 0011 0100 0101 01 => 0100 0110 1000 1010 1 => 0x46 0x8a 0x80 891 | ExpectedConsumedBytes: 3, 892 | }, 893 | { 894 | Name: "pattern 15", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 895 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 896 | Start: indecies{BitIndex: 1, ByteIndex: 0}, // ^ 897 | NBits: 17, // ^^ ^^^^ ^^^^ ^^^^ ^^^ 898 | Expected: []byte{0x8d, 0x15, 0x80}, // 10 0011 0100 0101 011 => 1000 1101 0001 0101 1 => 0x8D 0x15 0x80 899 | ExpectedConsumedBytes: 3, 900 | }, 901 | { 902 | Name: "pattern 16", // b7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 | 7654 3210 903 | Data: []byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12}, // 0001 0010 | 0011 0100 | 0101 0110 | 0111 1000 | 1001 1010 | 1011 1100 | 1101 1110 | 1111 0000 | 0001 0010 904 | Start: indecies{BitIndex: 0, ByteIndex: 0}, // ^ 905 | NBits: 17, // ^ ^^^^ ^^^^ ^^^^ ^^^^ 906 | Expected: []byte{0x1a, 0x2b, 0x00}, // 0 0011 0100 0101 0110 => 0001 1010 0010 1011 0 => 0x1A 0x2B 0x00 907 | ExpectedConsumedBytes: 3, 908 | }, 909 | } 910 | 911 | for _, data := range testData { 912 | data := data // capture 913 | t.Run(data.Name, func(t *testing.T) { 914 | //t.Parallel() 915 | 916 | r := NewReader(bytes.NewReader(data.Data), nil) 917 | r.fillBuf() 918 | r.currBitIndex = data.Start.BitIndex 919 | r.currByteIndex = data.Start.ByteIndex 920 | 921 | v, err := r.ReadNBits(data.NBits, &ReadOptions{AlignRight: data.AlignRight, PadOne: data.PadOne}) 922 | if err != nil { 923 | t.Fatalf("unexpected error: %+v\n", err) 924 | } 925 | if !reflect.DeepEqual(data.Expected, v) { 926 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", data.Expected, v) 927 | } 928 | if data.ExpectedConsumedBytes != r.ConsumedBytes() { 929 | t.Fatalf("\nExpected consumed bytes: %d\nActual consumed bytes: %d\n", data.ExpectedConsumedBytes, r.ConsumedBytes()) 930 | } 931 | }) 932 | } 933 | } 934 | 935 | func benchmarkReadNBits(b *testing.B, nBits uint8) { 936 | var v uint64 937 | r := NewReader(rand.Reader, nil) 938 | for n := 0; n < b.N; n++ { 939 | v, _ = r.ReadNBitsAsUint64BE(nBits) 940 | } 941 | toEliminateCompilerOptimizationUint64 = v 942 | } 943 | 944 | func BenchmarkRead1Bit(b *testing.B) { 945 | benchmarkReadNBits(b, 1) 946 | } 947 | 948 | func BenchmarkRead2Bits(b *testing.B) { 949 | benchmarkReadNBits(b, 2) 950 | } 951 | 952 | func BenchmarkRead9Bits(b *testing.B) { 953 | benchmarkReadNBits(b, 9) 954 | } 955 | 956 | func BenchmarkRead10Bits(b *testing.B) { 957 | benchmarkReadNBits(b, 10) 958 | } 959 | 960 | func BenchmarkRead11Bits(b *testing.B) { 961 | benchmarkReadNBits(b, 17) 962 | } 963 | 964 | func BenchmarkRead12Bits(b *testing.B) { 965 | benchmarkReadNBits(b, 18) 966 | } 967 | func BenchmarkRead15Bits(b *testing.B) { 968 | benchmarkReadNBits(b, 31) 969 | } 970 | 971 | func BenchmarkRead64Bits(b *testing.B) { 972 | benchmarkReadNBits(b, 64) 973 | } 974 | -------------------------------------------------------------------------------- /writer.go: -------------------------------------------------------------------------------- 1 | package bitstream 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | 7 | "github.com/pkg/errors" 8 | ) 9 | 10 | // Writer is a bit stream writer. 11 | // It does not have io.Writer interface 12 | type Writer struct { 13 | dst io.Writer 14 | currByte []uint8 15 | currBitIndex uint8 // MSB: 7, LSB: 0 16 | writtenBits uint 17 | } 18 | 19 | // NewWriter creates a new Writer instance. 20 | func NewWriter(dst io.Writer) *Writer { 21 | return &Writer{ 22 | dst: dst, 23 | currByte: []byte{0}, 24 | currBitIndex: 7, 25 | writtenBits: 0, 26 | } 27 | } 28 | 29 | func (w *Writer) dump() string { 30 | return fmt.Sprintf("currByte: %02x, currBitIndex: %d", w.currByte[0], w.currBitIndex) 31 | } 32 | 33 | func (w *Writer) WrittenBits() uint { 34 | return w.writtenBits 35 | } 36 | 37 | // WriteBit writes a single bit to the bit stream. 38 | // Uses the LSB bit in `bit`. 39 | func (w *Writer) WriteBit(bit uint8) error { 40 | if bit&0x01 != 0 { 41 | w.currByte[0] |= ((bit & 0x01) << w.currBitIndex) 42 | } 43 | w.writtenBits++ 44 | 45 | if w.currBitIndex > 0 { 46 | w.currBitIndex-- 47 | return nil 48 | } 49 | 50 | return w.Flush() 51 | } 52 | 53 | // WriteBool writes a single bit to the bit stream. 54 | // Write 1 if b is `true`, 0 otherwise. 55 | func (w *Writer) WriteBool(b bool) error { 56 | bit := uint8(0) 57 | if b { 58 | bit = 1 59 | } 60 | 61 | return w.WriteBit(bit) 62 | } 63 | 64 | // WriteNBitsOfUint8 writes `nBits` bits to the bit stream. 65 | // `nBits` must be less than or equal to 8, otherwise returns an error. 66 | // 67 | // This function uses n bits from `val`'s LSB. 68 | // i.e.) 69 | // if you have the following status of bit stream before calling WriteNBitsOfUint8, 70 | // currByte: 0101xxxxb 71 | // currBitIndex: 3 72 | // 73 | // and if you calls WriteNBitsOfUint8(3, 0xaa), 74 | // where nBits == 3, val == 0xaa (10101010b) 75 | // 76 | // WriteNBitsOfUint8 uses the 3 bits from `val`'s LSB, i.e.) xxxxx010b and as a result, status of the bit stream become: 77 | // currByte: 0101010xb (0101xxxxb | xxxx010xb) 78 | // currBitIndex: 0 79 | func (w *Writer) WriteNBitsOfUint8(nBits, val uint8) error { 80 | defer func() { w.writtenBits += uint(nBits) }() 81 | 82 | if nBits == 0 { 83 | return nil 84 | } 85 | 86 | if nBits > 8 { 87 | return errors.New("nBits too large for uint8") 88 | } 89 | 90 | // wb: bits can be written in currByte 91 | wb := w.currBitIndex + 1 92 | 93 | if nBits <= wb { // all the bits can be written in the currByte 94 | mask := uint8(1<<(nBits) - 1) // create a mask to make sure val has exactly n bits (to set 0's to upper bits) 95 | w.currByte[0] |= (val & mask) << (wb - nBits) 96 | if nBits == wb { 97 | return w.Flush() 98 | } 99 | w.currBitIndex -= nBits 100 | return nil 101 | } 102 | 103 | // need to separate val into 2 parts 104 | b1 := val >> (nBits - wb) // part 1: can be written in the currByte 105 | b2 := val << (8 - (nBits - wb)) // part 2: should be written in the next byte (MSB aligned) 106 | b1Mask := uint8((1 << (w.currBitIndex + 1)) - 1) 107 | w.currByte[0] |= (b1 & b1Mask) 108 | err := w.Flush() 109 | if err != nil { 110 | return err 111 | } 112 | w.currByte[0] = b2 113 | w.currBitIndex = 7 - (nBits - wb) 114 | 115 | return nil 116 | } 117 | 118 | // WriteUint8 writes a uint8 value to the bit stream. 119 | func (w *Writer) WriteUint8(val uint8) error { 120 | return w.WriteNBitsOfUint8(8, val) 121 | } 122 | 123 | // WriteNBitsOfUint16 writes `nBits` bits to the bit stream. 124 | // `nBits` must be less than or equal to 16, otherwise returns an error. 125 | func (w *Writer) WriteNBitsOfUint16BE(nBits uint8, val uint16) error { 126 | if nBits == 0 { 127 | return nil 128 | } 129 | 130 | if nBits <= 8 { 131 | return w.WriteNBitsOfUint8(nBits, uint8(val)) 132 | } 133 | 134 | if nBits > 16 { 135 | return errors.New("nBits too large for uint16") 136 | } 137 | 138 | defer func() { w.writtenBits += uint(nBits) }() 139 | 140 | // wb: bits can be written in currByte 141 | wb := w.currBitIndex + 1 142 | 143 | // 16 bits may be distributed in 3 bytes 144 | b1Bits := wb 145 | b2Bits := uint8(nBits - b1Bits) 146 | b3Bits := uint8(0) 147 | if b2Bits > 8 { 148 | b3Bits = b2Bits - 8 149 | b2Bits = 8 150 | } 151 | 152 | b1Mask := uint16(((1 << b1Bits) - 1) << (b2Bits + b3Bits)) 153 | b2Mask := uint16(((1 << b2Bits) - 1) << b3Bits) 154 | b3Mask := uint16((1 << b3Bits) - 1) 155 | 156 | b1 := uint8((val & b1Mask) >> (b2Bits + b3Bits)) 157 | b2 := uint8(((val & b2Mask) >> b3Bits) << (8 - b2Bits)) // left aligned 158 | b3 := uint8((val & b3Mask) << (8 - b3Bits)) // left aligned 159 | 160 | w.currByte[0] |= b1 161 | err := w.Flush() 162 | if err != nil { 163 | return err 164 | } 165 | 166 | if b3Bits == 0 { 167 | w.currByte[0] = b2 168 | if b2Bits == 8 { 169 | return w.Flush() 170 | } 171 | w.currBitIndex = 7 - b2Bits 172 | return nil 173 | } 174 | 175 | w.currByte[0] = b2 176 | err = w.Flush() 177 | if err != nil { 178 | return err 179 | } 180 | w.currByte[0] = b3 181 | w.currBitIndex = 7 - b3Bits 182 | 183 | return nil 184 | } 185 | 186 | // WriteUint16 writes a uint16 value to the bit stream. 187 | func (w *Writer) WriteUint16BE(val uint16) error { 188 | return w.WriteNBitsOfUint16BE(16, val) 189 | } 190 | 191 | // WriteNBitsOfUint32 writes `nBits` bits to the bit stream. 192 | // `nBits` must be less than or equal to 32, otherwise returns an error. 193 | func (w *Writer) WriteNBitsOfUint32BE(nBits uint8, val uint32) error { 194 | if nBits == 0 { 195 | return nil 196 | } 197 | 198 | if nBits <= 16 { 199 | return w.WriteNBitsOfUint16BE(nBits, uint16(val)) 200 | } 201 | 202 | if nBits > 32 { 203 | return errors.New("nBits too large for uint32") 204 | } 205 | 206 | defer func() { w.writtenBits += uint(nBits) }() 207 | 208 | // wb: bits can be written in currByte 209 | wb := w.currBitIndex + 1 210 | 211 | // 32 bits may be distributed in 5 bytes 212 | b1Bits := wb 213 | b2Bits := uint8(8) 214 | b3Bits := uint8(nBits - 8 - wb) 215 | b4Bits := uint8(0) 216 | b5Bits := uint8(0) 217 | if b3Bits > 8 { 218 | b4Bits = b3Bits - 8 219 | if b4Bits > 8 { 220 | b5Bits = b4Bits - 8 221 | b4Bits = 8 222 | } 223 | b3Bits = 8 224 | } 225 | 226 | b1Mask := uint32(((1 << b1Bits) - 1) << (b2Bits + b3Bits + b4Bits + b5Bits)) 227 | b2Mask := uint32(((1 << b2Bits) - 1) << (b3Bits + b4Bits + b5Bits)) 228 | b3Mask := uint32(((1 << b3Bits) - 1) << (b4Bits + b5Bits)) 229 | b4Mask := uint32(((1 << b4Bits) - 1) << b5Bits) 230 | b5Mask := uint32((1 << b5Bits) - 1) 231 | 232 | b1 := uint8((val & b1Mask) >> (b2Bits + b3Bits + b4Bits + b5Bits)) 233 | b2 := uint8(((val & b2Mask) >> (b3Bits + b4Bits + b5Bits)) << (8 - b2Bits)) // left aligned 234 | b3 := uint8(((val & b3Mask) >> (b4Bits + b5Bits)) << (8 - b3Bits)) // left aligned 235 | b4 := uint8(((val & b4Mask) >> b5Bits) << (8 - b4Bits)) // left aligned 236 | b5 := uint8((val & b5Mask) << (8 - b5Bits)) // left aligned 237 | 238 | w.currByte[0] |= b1 239 | err := w.Flush() 240 | if err != nil { 241 | return err 242 | } 243 | 244 | w.currByte[0] = b2 245 | err = w.Flush() 246 | if err != nil { 247 | return err 248 | } 249 | 250 | w.currByte[0] = b3 251 | if b3Bits == 8 { 252 | err = w.Flush() 253 | if err != nil { 254 | return err 255 | } 256 | } 257 | if b4Bits == 0 { 258 | if b3Bits != 8 { 259 | w.currBitIndex = 7 - b3Bits 260 | } 261 | return nil 262 | } 263 | 264 | w.currByte[0] = b4 265 | if b4Bits == 8 { 266 | err = w.Flush() 267 | if err != nil { 268 | return err 269 | } 270 | } 271 | if b5Bits == 0 { 272 | if b4Bits != 8 { 273 | w.currBitIndex = 7 - b4Bits 274 | } 275 | return nil 276 | } 277 | 278 | w.currByte[0] = b5 279 | w.currBitIndex = 7 - b5Bits 280 | 281 | return nil 282 | } 283 | 284 | // WriteUint32 writes a uint32 value to the bit stream. 285 | func (w *Writer) WriteUint32BE(val uint32) error { 286 | return w.WriteNBitsOfUint32BE(32, val) 287 | } 288 | 289 | // WriteNBits writes specified number of bits of the bytes to the bit stream. 290 | func (w *Writer) WriteNBits(nBits uint, data []byte) error { 291 | if nBits == 0 { 292 | return nil 293 | } 294 | 295 | for nBits > 8 { 296 | if len(data) == 0 { 297 | return errors.New("insufficient data") 298 | } 299 | 300 | b := data[0] 301 | err := w.WriteNBitsOfUint8(8, b) 302 | if err != nil { 303 | return err 304 | } 305 | data = data[1:] 306 | nBits -= 8 307 | } 308 | 309 | if nBits > 0 { 310 | if len(data) == 0 { 311 | return errors.New("insufficient data") 312 | } 313 | b := data[0] 314 | b = b >> (8 - nBits) 315 | err := w.WriteNBitsOfUint8(uint8(nBits), b) 316 | if err != nil { 317 | return err 318 | } 319 | } 320 | 321 | return nil 322 | } 323 | 324 | // Flush ensures the bufferred bits (bits not writen to the stream because it has less than 8 bits) to the destination writer. 325 | func (w *Writer) Flush() error { 326 | nWritten, err := w.dst.Write(w.currByte) 327 | if err != nil { 328 | return err 329 | } 330 | if nWritten != 1 { 331 | return errors.New("unable to write 1 byte") 332 | } 333 | 334 | w.currByte[0] = 0x00 335 | w.currBitIndex = 7 336 | 337 | return nil 338 | } 339 | -------------------------------------------------------------------------------- /writer_test.go: -------------------------------------------------------------------------------- 1 | package bitstream 2 | 3 | import ( 4 | "bytes" 5 | "math/rand" 6 | "reflect" 7 | "testing" 8 | "time" 9 | ) 10 | 11 | func TestWriteBit(t *testing.T) { 12 | buf := bytes.NewBuffer([]byte{}) 13 | bw := NewWriter(buf) 14 | 15 | bw.WriteBit(0) 16 | bw.WriteBit(1) 17 | bw.WriteBit(0) 18 | bw.WriteBit(1) 19 | bw.WriteBit(0) 20 | bw.WriteBit(0) 21 | bw.WriteBit(1) 22 | bw.WriteBit(1) 23 | bw.WriteBit(1) 24 | bw.WriteBit(0) 25 | bw.WriteBit(1) 26 | bw.WriteBit(0) 27 | bw.WriteBit(1) 28 | bw.WriteBit(1) 29 | bw.WriteBit(0) 30 | bw.WriteBit(0) 31 | 32 | expected := []byte{0x53, 0xac} 33 | if !reflect.DeepEqual(buf.Bytes(), expected) { 34 | t.Fatalf("\nExpected: %+v\nActual: %+v\n", expected, buf.Bytes()) 35 | } 36 | if uint(16) != bw.WrittenBits() { 37 | t.Fatalf("\nunexpected writtenBits\nExpected: %+v\nActual: %+v\n", 16, bw.WrittenBits()) 38 | } 39 | } 40 | 41 | func BenchmarkWriteBit(b *testing.B) { 42 | rand.Seed(time.Now().UnixNano()) 43 | buf := bytes.NewBuffer([]byte{}) 44 | bw := NewWriter(buf) 45 | for n := 0; n < b.N; n++ { 46 | _ = bw.WriteBit(uint8(rand.Intn(256))) 47 | } 48 | } 49 | 50 | type writerStatus struct { 51 | currByte byte 52 | currBitIndex uint8 53 | buf []byte 54 | } 55 | 56 | func TestWriteNBitsOfUint8(t *testing.T) { 57 | testData := []struct { 58 | Name string 59 | NBits uint8 60 | Value uint8 61 | Start writerStatus 62 | Expected writerStatus 63 | }{ 64 | { 65 | Name: "pattern 1", 66 | NBits: 1, 67 | Value: 0xff, // 1 68 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, 69 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, 70 | }, 71 | { 72 | Name: "pattern 2", 73 | NBits: 4, 74 | Value: 0x02, // 0010 75 | Start: writerStatus{currByte: 0x40, currBitIndex: 5, buf: []byte{}}, // 01xx xxxx 76 | Expected: writerStatus{currByte: 0x48, currBitIndex: 1, buf: []byte{}}, // 0100 10xx 77 | }, 78 | { 79 | Name: "pattern 3", 80 | NBits: 8, 81 | Value: 0xff, 82 | Start: writerStatus{currByte: 0x00, currBitIndex: 3, buf: []byte{}}, 83 | Expected: writerStatus{currByte: 0xf0, currBitIndex: 3, buf: []byte{0x0f}}, 84 | }, 85 | { 86 | Name: "pattern 4", 87 | NBits: 8, 88 | Value: 0xaa, 89 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, 90 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0xaa}}, 91 | }, 92 | } 93 | 94 | for _, data := range testData { 95 | data := data // capture 96 | t.Run(data.Name, func(t *testing.T) { 97 | //t.Parallel() 98 | 99 | buf := bytes.NewBuffer(data.Start.buf) 100 | bw := NewWriter(buf) 101 | 102 | bw.currByte[0] = data.Start.currByte 103 | bw.currBitIndex = data.Start.currBitIndex 104 | 105 | err := bw.WriteNBitsOfUint8(data.NBits, data.Value) 106 | if err != nil { 107 | t.Fatalf("unexpected error: %+v\n", err) 108 | } 109 | if uint(data.NBits) != bw.WrittenBits() { 110 | t.Fatalf("\nunexpected writtenBits\nExpected: %+v\nActual: %+v\n", data.NBits, bw.WrittenBits()) 111 | } 112 | if data.Expected.currByte != bw.currByte[0] { 113 | t.Fatalf("\nunexpected currByte\nExpected: %+v\nActual: %+v\n", data.Expected.currByte, bw.currByte[0]) 114 | } 115 | if data.Expected.currBitIndex != bw.currBitIndex { 116 | t.Fatalf("\nunexpected currBitIndex\nExpected: %+v\nActual: %+v\n", data.Expected.currBitIndex, bw.currBitIndex) 117 | } 118 | if !reflect.DeepEqual(data.Expected.buf, buf.Bytes()) { 119 | t.Fatalf("\nunexpected flushed data\nExpected: %+v\nActual: %+v\n", data.Expected.buf, buf.Bytes()) 120 | } 121 | 122 | }) 123 | } 124 | 125 | } 126 | 127 | func benchmarkWriteNBitsOfUint8(nBits uint8, b *testing.B) { 128 | rand.Seed(time.Now().UnixNano()) 129 | buf := bytes.NewBuffer([]byte{}) 130 | bw := NewWriter(buf) 131 | for n := 0; n < b.N; n++ { 132 | _ = bw.WriteNBitsOfUint8(nBits, uint8(rand.Intn(256))) 133 | } 134 | } 135 | 136 | func BenchmarkWrite1BitsOfUint8(b *testing.B) { 137 | benchmarkWriteNBitsOfUint8(1, b) 138 | } 139 | 140 | func BenchmarkWrite2BitsOfUint8(b *testing.B) { 141 | benchmarkWriteNBitsOfUint8(2, b) 142 | } 143 | 144 | func BenchmarkWrite3BitsOfUint8(b *testing.B) { 145 | benchmarkWriteNBitsOfUint8(3, b) 146 | } 147 | 148 | func BenchmarkWrite4BitsOfUint8(b *testing.B) { 149 | benchmarkWriteNBitsOfUint8(4, b) 150 | } 151 | 152 | func BenchmarkWrite5BitsOfUint8(b *testing.B) { 153 | benchmarkWriteNBitsOfUint8(5, b) 154 | } 155 | 156 | func BenchmarkWrite6BitsOfUint8(b *testing.B) { 157 | benchmarkWriteNBitsOfUint8(6, b) 158 | } 159 | 160 | func BenchmarkWrite7BitsOfUint8(b *testing.B) { 161 | benchmarkWriteNBitsOfUint8(7, b) 162 | } 163 | 164 | func BenchmarkWrite8BitsOfUint8(b *testing.B) { 165 | benchmarkWriteNBitsOfUint8(8, b) 166 | } 167 | 168 | func TestWriteNBitsOfUint16BE(t *testing.T) { 169 | testData := []struct { 170 | Name string 171 | NBits uint8 172 | Value uint16 173 | Start writerStatus 174 | Expected writerStatus 175 | }{ 176 | { 177 | Name: "pattern 1", 178 | NBits: 1, 179 | Value: 0xffff, // 1 180 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, 181 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, 182 | }, 183 | { 184 | Name: "pattern 2", 185 | NBits: 10, 186 | Value: 0x2222, // 10 0010 0010 187 | Start: writerStatus{currByte: 0x40, currBitIndex: 5, buf: []byte{}}, // 01xx xxxx 188 | Expected: writerStatus{currByte: 0x20, currBitIndex: 3, buf: []byte{0x62}}, // 0110 0010 0010 xxxx 189 | }, 190 | { 191 | Name: "pattern 3", 192 | NBits: 13, 193 | Value: 0xffff, // 1111 1111 1111 1 194 | Start: writerStatus{currByte: 0x00, currBitIndex: 3, buf: []byte{}}, // 0000 xxxx 195 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{0x0f, 0xff}}, // 0000 1111 1111 1111 1xxx 196 | }, 197 | { 198 | Name: "pattern 4", 199 | NBits: 16, 200 | Value: 0xabcd, // 1010 1011 1100 1101 201 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 202 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0xab, 0xcd}}, // 1010 1011 1100 1101 203 | }, 204 | { 205 | Name: "pattern 5", 206 | NBits: 16, 207 | Value: 0xabcd, // 10 1010 1111 0011 01 208 | Start: writerStatus{currByte: 0x88, currBitIndex: 1, buf: []byte{}}, // 1000 10xx 209 | Expected: writerStatus{currByte: 0x34, currBitIndex: 1, buf: []byte{0x8a, 0xaf}}, // 1000 1010 1010 1111 0011 01xx 210 | }, 211 | } 212 | 213 | for _, data := range testData { 214 | data := data // capture 215 | t.Run(data.Name, func(t *testing.T) { 216 | //t.Parallel() 217 | 218 | buf := bytes.NewBuffer(data.Start.buf) 219 | bw := NewWriter(buf) 220 | 221 | bw.currByte[0] = data.Start.currByte 222 | bw.currBitIndex = data.Start.currBitIndex 223 | 224 | err := bw.WriteNBitsOfUint16BE(data.NBits, data.Value) 225 | if err != nil { 226 | t.Fatalf("unexpected error: %+v\n", err) 227 | } 228 | if uint(data.NBits) != bw.WrittenBits() { 229 | t.Fatalf("\nunexpected writtenBits\nExpected: %+v\nActual: %+v\n", data.NBits, bw.WrittenBits()) 230 | } 231 | if data.Expected.currByte != bw.currByte[0] { 232 | t.Fatalf("\nunexpected currByte\nExpected: %+v\nActual: %+v\n", data.Expected.currByte, bw.currByte[0]) 233 | } 234 | if data.Expected.currBitIndex != bw.currBitIndex { 235 | t.Fatalf("\nunexpected currBitIndex\nExpected: %+v\nActual: %+v\n", data.Expected.currBitIndex, bw.currBitIndex) 236 | } 237 | if !reflect.DeepEqual(data.Expected.buf, buf.Bytes()) { 238 | t.Fatalf("\nunexpected flushed data\nExpected: %+v\nActual: %+v\n", data.Expected.buf, buf.Bytes()) 239 | } 240 | 241 | }) 242 | } 243 | 244 | } 245 | 246 | func benchmarkWriteNBitsOfUint16BE(nBits uint8, b *testing.B) { 247 | rand.Seed(time.Now().UnixNano()) 248 | buf := bytes.NewBuffer([]byte{}) 249 | bw := NewWriter(buf) 250 | for n := 0; n < b.N; n++ { 251 | _ = bw.WriteNBitsOfUint16BE(nBits, uint16(rand.Intn(65536))) 252 | } 253 | } 254 | 255 | func BenchmarkWrite1BitsOfUint16BE(b *testing.B) { 256 | benchmarkWriteNBitsOfUint16BE(1, b) 257 | } 258 | 259 | func BenchmarkWrite2BitsOfUint16BE(b *testing.B) { 260 | benchmarkWriteNBitsOfUint16BE(2, b) 261 | } 262 | 263 | func BenchmarkWrite9BitsOfUint16BE(b *testing.B) { 264 | benchmarkWriteNBitsOfUint16BE(9, b) 265 | } 266 | 267 | func BenchmarkWrite10BitsOfUint16BE(b *testing.B) { 268 | benchmarkWriteNBitsOfUint16BE(10, b) 269 | } 270 | 271 | func BenchmarkWrite15BitsOfUint16BE(b *testing.B) { 272 | benchmarkWriteNBitsOfUint16BE(15, b) 273 | } 274 | 275 | func BenchmarkWrite16BitsOfUint16BE(b *testing.B) { 276 | benchmarkWriteNBitsOfUint16BE(16, b) 277 | } 278 | 279 | func TestWriteNBitsOfUint32BE(t *testing.T) { 280 | testData := []struct { 281 | Name string 282 | NBits uint8 283 | Value uint32 284 | Start writerStatus 285 | Expected writerStatus 286 | }{ 287 | { 288 | Name: "pattern 1", 289 | NBits: 1, 290 | Value: 0xffff, // 1 291 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, 292 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, 293 | }, 294 | { 295 | Name: "pattern 2", 296 | NBits: 16, 297 | Value: 0xffffabcd, // 10 1010 1111 0011 01 298 | Start: writerStatus{currByte: 0x88, currBitIndex: 1, buf: []byte{}}, // 1000 10xx 299 | Expected: writerStatus{currByte: 0x34, currBitIndex: 1, buf: []byte{0x8a, 0xaf}}, // 1000 1010 1010 1111 0011 01xx 300 | }, 301 | { 302 | Name: "pattern 3", 303 | NBits: 17, 304 | Value: 0xffffabcd, // 11 0101 0111 1001 101 305 | Start: writerStatus{currByte: 0x88, currBitIndex: 1, buf: []byte{}}, // 1000 10xx 306 | Expected: writerStatus{currByte: 0x9a, currBitIndex: 0, buf: []byte{0x8b, 0x57}}, // 1000 1011 0101 0111 1001 101x 307 | }, 308 | { 309 | Name: "pattern 4", 310 | NBits: 24, 311 | Value: 0xffabcdef, // 1010 1011 1100 1101 1110 1111 312 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 313 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0xab, 0xcd, 0xef}}, // 1010 1011 1100 1101 1110 1111 314 | }, 315 | { 316 | Name: "pattern 5", 317 | NBits: 24, 318 | Value: 0xffabcdef, // 1010 1011 1100 1101 1110 1111 319 | Start: writerStatus{currByte: 0xf0, currBitIndex: 3, buf: []byte{}}, // 1111 xxxx 320 | Expected: writerStatus{currByte: 0xf0, currBitIndex: 3, buf: []byte{0xfa, 0xbc, 0xde}}, // 1111 1010 1011 1100 1101 1110 1111 xxxx 321 | }, 322 | { 323 | Name: "pattern 6", 324 | NBits: 24, 325 | Value: 0xffabcdef, // 10 1010 1111 0011 0111 1011 11 326 | Start: writerStatus{currByte: 0xfc, currBitIndex: 1, buf: []byte{}}, // 1111 11xx 327 | Expected: writerStatus{currByte: 0xbc, currBitIndex: 1, buf: []byte{0xfe, 0xaf, 0x37}}, // 1111 1110 1010 1111 0011 0111 1011 11xx 328 | }, 329 | { 330 | Name: "pattern 7", 331 | NBits: 31, 332 | Value: 0x89abcdef, // 0001 0011 0101 0111 1001 1011 1101 111 333 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 334 | Expected: writerStatus{currByte: 0xde, currBitIndex: 0, buf: []byte{0x13, 0x57, 0x9b}}, // 0001 0011 0101 0111 1001 1011 1101 111x 335 | }, 336 | { 337 | Name: "pattern 8", 338 | NBits: 31, 339 | Value: 0x89abcdef, // 000 1001 1010 1011 1100 1101 1110 1111 340 | Start: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, // 1xxx xxxx 341 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0x89, 0xab, 0xcd, 0xef}}, // 1000 1001 1010 1011 1100 1101 1110 1111 342 | }, 343 | { 344 | Name: "pattern 9", 345 | NBits: 31, 346 | Value: 0x89abcdef, // 00 0100 1101 0101 1110 0110 1111 0111 1 347 | Start: writerStatus{currByte: 0xc0, currBitIndex: 5, buf: []byte{}}, // 11xx xxxx 348 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{0xc4, 0xd5, 0xe6, 0xf7}}, // 1100 0100 1101 0101 1110 0110 1111 0111 1xxx xxxx 349 | }, 350 | { 351 | Name: "pattern 10", 352 | NBits: 32, 353 | Value: 0x89abcdef, // 1000 1001 1010 1011 1100 1101 1110 1111 354 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 355 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0x89, 0xab, 0xcd, 0xef}}, // 1000 1001 1010 1011 1100 1101 1110 1111 xxxx xxxx 356 | }, 357 | { 358 | Name: "pattern 11", 359 | NBits: 32, 360 | Value: 0x89abcdef, // 100 0100 1101 0101 1110 0110 1111 0111 1 361 | Start: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, // 1xxx xxxx 362 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{0xc4, 0xd5, 0xe6, 0xf7}}, // 1100 0100 1101 0101 1110 0110 1111 0111 1xxx xxxx 363 | }, 364 | } 365 | 366 | for _, data := range testData { 367 | data := data // capture 368 | t.Run(data.Name, func(t *testing.T) { 369 | //t.Parallel() 370 | 371 | buf := bytes.NewBuffer(data.Start.buf) 372 | bw := NewWriter(buf) 373 | 374 | bw.currByte[0] = data.Start.currByte 375 | bw.currBitIndex = data.Start.currBitIndex 376 | 377 | err := bw.WriteNBitsOfUint32BE(data.NBits, data.Value) 378 | if err != nil { 379 | t.Fatalf("unexpected error: %+v\n", err) 380 | } 381 | if uint(data.NBits) != bw.WrittenBits() { 382 | t.Fatalf("\nunexpected writtenBits\nExpected: %+v\nActual: %+v\n", data.NBits, bw.WrittenBits()) 383 | } 384 | if data.Expected.currByte != bw.currByte[0] { 385 | t.Fatalf("\nunexpected currByte\nExpected: %+v\nActual: %+v\n", data.Expected.currByte, bw.currByte[0]) 386 | } 387 | if data.Expected.currBitIndex != bw.currBitIndex { 388 | t.Fatalf("\nunexpected currBitIndex\nExpected: %+v\nActual: %+v\n", data.Expected.currBitIndex, bw.currBitIndex) 389 | } 390 | if !reflect.DeepEqual(data.Expected.buf, buf.Bytes()) { 391 | t.Fatalf("\nunexpected flushed data\nExpected: %+v\nActual: %+v\n", data.Expected.buf, buf.Bytes()) 392 | } 393 | 394 | }) 395 | } 396 | 397 | } 398 | 399 | func benchmarkWriteNBitsOfUint32BE(nBits uint8, b *testing.B) { 400 | rand.Seed(time.Now().UnixNano()) 401 | buf := bytes.NewBuffer([]byte{}) 402 | bw := NewWriter(buf) 403 | for n := 0; n < b.N; n++ { 404 | _ = bw.WriteNBitsOfUint32BE(nBits, uint32(rand.Uint32())) 405 | } 406 | } 407 | 408 | func BenchmarkWrite1BitsOfUint32BE(b *testing.B) { 409 | benchmarkWriteNBitsOfUint32BE(1, b) 410 | } 411 | 412 | func BenchmarkWrite16BitsOfUint32BE(b *testing.B) { 413 | benchmarkWriteNBitsOfUint32BE(16, b) 414 | } 415 | 416 | func BenchmarkWrite17BitsOfUint32BE(b *testing.B) { 417 | benchmarkWriteNBitsOfUint32BE(17, b) 418 | } 419 | 420 | func BenchmarkWrite23BitsOfUint32BE(b *testing.B) { 421 | benchmarkWriteNBitsOfUint32BE(23, b) 422 | } 423 | 424 | func BenchmarkWrite31BitsOfUint32BE(b *testing.B) { 425 | benchmarkWriteNBitsOfUint32BE(31, b) 426 | } 427 | 428 | func BenchmarkWrite32BitsOfUint32BE(b *testing.B) { 429 | benchmarkWriteNBitsOfUint32BE(32, b) 430 | } 431 | 432 | func TestWriteNBits(t *testing.T) { 433 | testData := []struct { 434 | Name string 435 | NBits uint 436 | Value []byte 437 | Start writerStatus 438 | Expected writerStatus 439 | }{ 440 | { 441 | Name: "pattern 1", 442 | NBits: 1, 443 | Value: []byte{0xff}, // 1 444 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, 445 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, 446 | }, 447 | { 448 | Name: "pattern 2", 449 | NBits: 16, 450 | Value: []byte{0xab, 0xcd}, // 10 1010 1111 0011 01 451 | Start: writerStatus{currByte: 0x88, currBitIndex: 1, buf: []byte{}}, // 1000 10xx 452 | Expected: writerStatus{currByte: 0x34, currBitIndex: 1, buf: []byte{0x8a, 0xaf}}, // 1000 1010 1010 1111 0011 01xx 453 | }, 454 | { 455 | Name: "pattern 3", 456 | NBits: 17, 457 | Value: []byte{0xab, 0xcd, 0xef}, // 10 1010 1111 0011 011 458 | Start: writerStatus{currByte: 0x88, currBitIndex: 1, buf: []byte{}}, // 1000 10xx 459 | Expected: writerStatus{currByte: 0x36, currBitIndex: 0, buf: []byte{0x8a, 0xaf}}, // 1000 1010 1010 1111 0011 011x 460 | }, 461 | { 462 | Name: "pattern 4", 463 | NBits: 24, 464 | Value: []byte{0xab, 0xcd, 0xef, 0xff}, // 1010 1011 1100 1101 1110 1111 465 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 466 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0xab, 0xcd, 0xef}}, // 1010 1011 1100 1101 1110 1111 467 | }, 468 | { 469 | Name: "pattern 5", 470 | NBits: 24, 471 | Value: []byte{0xab, 0xcd, 0xef, 0xff}, // 1010 1011 1100 1101 1110 1111 472 | Start: writerStatus{currByte: 0xf0, currBitIndex: 3, buf: []byte{}}, // 1111 xxxx 473 | Expected: writerStatus{currByte: 0xf0, currBitIndex: 3, buf: []byte{0xfa, 0xbc, 0xde}}, // 1111 1010 1011 1100 1101 1110 1111 xxxx 474 | }, 475 | { 476 | Name: "pattern 6", 477 | NBits: 24, 478 | Value: []byte{0xab, 0xcd, 0xef, 0xff}, // 10 1010 1111 0011 0111 1011 11 479 | Start: writerStatus{currByte: 0xfc, currBitIndex: 1, buf: []byte{}}, // 1111 11xx 480 | Expected: writerStatus{currByte: 0xbc, currBitIndex: 1, buf: []byte{0xfe, 0xaf, 0x37}}, // 1111 1110 1010 1111 0011 0111 1011 11xx 481 | }, 482 | { 483 | Name: "pattern 7", 484 | NBits: 31, 485 | Value: []byte{0x89, 0xab, 0xcd, 0xef}, // 1000 1001 1010 1011 1100 1101 1110 111 486 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 487 | Expected: writerStatus{currByte: 0xee, currBitIndex: 0, buf: []byte{0x89, 0xab, 0xcd}}, // 1000 1001 1010 1011 1100 1101 1110 111x 488 | }, 489 | { 490 | Name: "pattern 8", 491 | NBits: 31, 492 | Value: []byte{0x89, 0xab, 0xcd, 0xef}, // 100 0100 1101 0101 1110 0110 1111 0111 493 | Start: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, // 1xxx xxxx 494 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0xc4, 0xd5, 0xe6, 0xf7}}, // 1100 0100 1101 0101 1110 0110 1111 0111 xxxx xxxx 495 | }, 496 | { 497 | Name: "pattern 9", 498 | NBits: 31, 499 | Value: []byte{0x89, 0xab, 0xcd, 0xef}, // 10 0010 0110 1010 1111 0011 0111 1011 1 500 | Start: writerStatus{currByte: 0xc0, currBitIndex: 5, buf: []byte{}}, // 11xx xxxx 501 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{0xe2, 0x6a, 0xf3, 0x7b}}, // 1110 0010 0110 1010 1111 0011 0111 1011 1xxx xxxx 502 | }, 503 | { 504 | Name: "pattern 10", 505 | NBits: 32, 506 | Value: []byte{0x89, 0xab, 0xcd, 0xef}, // 1000 1001 1010 1011 1100 1101 1110 1111 507 | Start: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{}}, // xxxx xxxx 508 | Expected: writerStatus{currByte: 0x00, currBitIndex: 7, buf: []byte{0x89, 0xab, 0xcd, 0xef}}, // 1000 1001 1010 1011 1100 1101 1110 1111 xxxx xxxx 509 | }, 510 | { 511 | Name: "pattern 11", 512 | NBits: 32, 513 | Value: []byte{0x89, 0xab, 0xcd, 0xef}, // 100 0100 1101 0101 1110 0110 1111 0111 1 514 | Start: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{}}, // 1xxx xxxx 515 | Expected: writerStatus{currByte: 0x80, currBitIndex: 6, buf: []byte{0xc4, 0xd5, 0xe6, 0xf7}}, // 1100 0100 1101 0101 1110 0110 1111 0111 1xxx xxxx 516 | }, 517 | } 518 | 519 | for _, data := range testData { 520 | data := data // capture 521 | t.Run(data.Name, func(t *testing.T) { 522 | //t.Parallel() 523 | 524 | buf := bytes.NewBuffer(data.Start.buf) 525 | bw := NewWriter(buf) 526 | 527 | bw.currByte[0] = data.Start.currByte 528 | bw.currBitIndex = data.Start.currBitIndex 529 | 530 | err := bw.WriteNBits(data.NBits, data.Value) 531 | if err != nil { 532 | t.Fatalf("unexpected error: %+v\n", err) 533 | } 534 | if data.NBits != bw.WrittenBits() { 535 | t.Fatalf("\nunexpected writtenBits\nExpected: %+v\nActual: %+v\n", data.NBits, bw.WrittenBits()) 536 | } 537 | if data.Expected.currByte != bw.currByte[0] { 538 | t.Fatalf("\nunexpected currByte\nExpected: %+v\nActual: %+v\n", data.Expected.currByte, bw.currByte[0]) 539 | } 540 | if data.Expected.currBitIndex != bw.currBitIndex { 541 | t.Fatalf("\nunexpected currBitIndex\nExpected: %+v\nActual: %+v\n", data.Expected.currBitIndex, bw.currBitIndex) 542 | } 543 | if !reflect.DeepEqual(data.Expected.buf, buf.Bytes()) { 544 | t.Fatalf("\nunexpected flushed data\nExpected: %+v\nActual: %+v\n", data.Expected.buf, buf.Bytes()) 545 | } 546 | 547 | }) 548 | } 549 | 550 | } 551 | --------------------------------------------------------------------------------