1 // Copyright 2018 Klaus Post. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4 // Based on work Copyright (c) 2013, Yann Collet, released under BSD License.
5 6 package zstd
7 8 // bitWriter will write bits.
9 // First bit will be LSB of the first byte of output.
10 type bitWriter struct {
11 bitContainer uint64
12 nBits uint8
13 out []byte
14 }
15 16 // bitMask16 is bitmasks. Has extra to avoid bounds check.
17 var bitMask16 = [32]uint16{
18 0, 1, 3, 7, 0xF, 0x1F,
19 0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF, 0x7FF,
20 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF, 0xFFFF,
21 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF,
22 0xFFFF, 0xFFFF} /* up to 16 bits */
23 24 var bitMask32 = [32]uint32{
25 0, 1, 3, 7, 0xF, 0x1F, 0x3F, 0x7F, 0xFF,
26 0x1FF, 0x3FF, 0x7FF, 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF,
27 0x1ffff, 0x3ffff, 0x7FFFF, 0xfFFFF, 0x1fFFFF, 0x3fFFFF, 0x7fFFFF, 0xffFFFF,
28 0x1ffFFFF, 0x3ffFFFF, 0x7ffFFFF, 0xfffFFFF, 0x1fffFFFF, 0x3fffFFFF, 0x7fffFFFF,
29 } // up to 32 bits
30 31 // addBits16NC will add up to 16 bits.
32 // It will not check if there is space for them,
33 // so the caller must ensure that it has flushed recently.
34 func (b *bitWriter) addBits16NC(value uint16, bits uint8) {
35 b.bitContainer |= uint64(value&bitMask16[bits&31]) << (b.nBits & 63)
36 b.nBits += bits
37 }
38 39 // addBits32NC will add up to 31 bits.
40 // It will not check if there is space for them,
41 // so the caller must ensure that it has flushed recently.
42 func (b *bitWriter) addBits32NC(value uint32, bits uint8) {
43 b.bitContainer |= uint64(value&bitMask32[bits&31]) << (b.nBits & 63)
44 b.nBits += bits
45 }
46 47 // addBits64NC will add up to 64 bits.
48 // There must be space for 32 bits.
49 func (b *bitWriter) addBits64NC(value uint64, bits uint8) {
50 if bits <= 31 {
51 b.addBits32Clean(uint32(value), bits)
52 return
53 }
54 b.addBits32Clean(uint32(value), 32)
55 b.flush32()
56 b.addBits32Clean(uint32(value>>32), bits-32)
57 }
58 59 // addBits32Clean will add up to 32 bits.
60 // It will not check if there is space for them.
61 // The input must not contain more bits than specified.
62 func (b *bitWriter) addBits32Clean(value uint32, bits uint8) {
63 b.bitContainer |= uint64(value) << (b.nBits & 63)
64 b.nBits += bits
65 }
66 67 // addBits16Clean will add up to 16 bits. value may not contain more set bits than indicated.
68 // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
69 func (b *bitWriter) addBits16Clean(value uint16, bits uint8) {
70 b.bitContainer |= uint64(value) << (b.nBits & 63)
71 b.nBits += bits
72 }
73 74 // flush32 will flush out, so there are at least 32 bits available for writing.
75 func (b *bitWriter) flush32() {
76 if b.nBits < 32 {
77 return
78 }
79 b.out = append(b.out,
80 byte(b.bitContainer),
81 byte(b.bitContainer>>8),
82 byte(b.bitContainer>>16),
83 byte(b.bitContainer>>24))
84 b.nBits -= 32
85 b.bitContainer >>= 32
86 }
87 88 // flushAlign will flush remaining full bytes and align to next byte boundary.
89 func (b *bitWriter) flushAlign() {
90 nbBytes := (b.nBits + 7) >> 3
91 for i := range nbBytes {
92 b.out = append(b.out, byte(b.bitContainer>>(i*8)))
93 }
94 b.nBits = 0
95 b.bitContainer = 0
96 }
97 98 // close will write the alignment bit and write the final byte(s)
99 // to the output.
100 func (b *bitWriter) close() {
101 // End mark
102 b.addBits16Clean(1, 1)
103 // flush until next byte.
104 b.flushAlign()
105 }
106 107 // reset and continue writing by appending to out.
108 func (b *bitWriter) reset(out []byte) {
109 b.bitContainer = 0
110 b.nBits = 0
111 b.out = out
112 }
113