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 fse
7 8 import "fmt"
9 10 // bitWriter will write bits.
11 // First bit will be LSB of the first byte of output.
12 type bitWriter struct {
13 bitContainer uint64
14 nBits uint8
15 out []byte
16 }
17 18 // bitMask16 is bitmasks. Has extra to avoid bounds check.
19 var bitMask16 = [32]uint16{
20 0, 1, 3, 7, 0xF, 0x1F,
21 0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF, 0x7FF,
22 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF, 0xFFFF,
23 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF,
24 0xFFFF, 0xFFFF} /* up to 16 bits */
25 26 // addBits16NC will add up to 16 bits.
27 // It will not check if there is space for them,
28 // so the caller must ensure that it has flushed recently.
29 func (b *bitWriter) addBits16NC(value uint16, bits uint8) {
30 b.bitContainer |= uint64(value&bitMask16[bits&31]) << (b.nBits & 63)
31 b.nBits += bits
32 }
33 34 // addBits16Clean will add up to 16 bits. value may not contain more set bits than indicated.
35 // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
36 func (b *bitWriter) addBits16Clean(value uint16, bits uint8) {
37 b.bitContainer |= uint64(value) << (b.nBits & 63)
38 b.nBits += bits
39 }
40 41 // addBits16ZeroNC will add up to 16 bits.
42 // It will not check if there is space for them,
43 // so the caller must ensure that it has flushed recently.
44 // This is fastest if bits can be zero.
45 func (b *bitWriter) addBits16ZeroNC(value uint16, bits uint8) {
46 if bits == 0 {
47 return
48 }
49 value <<= (16 - bits) & 15
50 value >>= (16 - bits) & 15
51 b.bitContainer |= uint64(value) << (b.nBits & 63)
52 b.nBits += bits
53 }
54 55 // flush will flush all pending full bytes.
56 // There will be at least 56 bits available for writing when this has been called.
57 // Using flush32 is faster, but leaves less space for writing.
58 func (b *bitWriter) flush() {
59 v := b.nBits >> 3
60 switch v {
61 case 0:
62 case 1:
63 b.out = append(b.out,
64 byte(b.bitContainer),
65 )
66 case 2:
67 b.out = append(b.out,
68 byte(b.bitContainer),
69 byte(b.bitContainer>>8),
70 )
71 case 3:
72 b.out = append(b.out,
73 byte(b.bitContainer),
74 byte(b.bitContainer>>8),
75 byte(b.bitContainer>>16),
76 )
77 case 4:
78 b.out = append(b.out,
79 byte(b.bitContainer),
80 byte(b.bitContainer>>8),
81 byte(b.bitContainer>>16),
82 byte(b.bitContainer>>24),
83 )
84 case 5:
85 b.out = append(b.out,
86 byte(b.bitContainer),
87 byte(b.bitContainer>>8),
88 byte(b.bitContainer>>16),
89 byte(b.bitContainer>>24),
90 byte(b.bitContainer>>32),
91 )
92 case 6:
93 b.out = append(b.out,
94 byte(b.bitContainer),
95 byte(b.bitContainer>>8),
96 byte(b.bitContainer>>16),
97 byte(b.bitContainer>>24),
98 byte(b.bitContainer>>32),
99 byte(b.bitContainer>>40),
100 )
101 case 7:
102 b.out = append(b.out,
103 byte(b.bitContainer),
104 byte(b.bitContainer>>8),
105 byte(b.bitContainer>>16),
106 byte(b.bitContainer>>24),
107 byte(b.bitContainer>>32),
108 byte(b.bitContainer>>40),
109 byte(b.bitContainer>>48),
110 )
111 case 8:
112 b.out = append(b.out,
113 byte(b.bitContainer),
114 byte(b.bitContainer>>8),
115 byte(b.bitContainer>>16),
116 byte(b.bitContainer>>24),
117 byte(b.bitContainer>>32),
118 byte(b.bitContainer>>40),
119 byte(b.bitContainer>>48),
120 byte(b.bitContainer>>56),
121 )
122 default:
123 panic(fmt.Errorf("bits (%d) > 64", b.nBits))
124 }
125 b.bitContainer >>= v << 3
126 b.nBits &= 7
127 }
128 129 // flush32 will flush out, so there are at least 32 bits available for writing.
130 func (b *bitWriter) flush32() {
131 if b.nBits < 32 {
132 return
133 }
134 b.out = append(b.out,
135 byte(b.bitContainer),
136 byte(b.bitContainer>>8),
137 byte(b.bitContainer>>16),
138 byte(b.bitContainer>>24))
139 b.nBits -= 32
140 b.bitContainer >>= 32
141 }
142 143 // flushAlign will flush remaining full bytes and align to next byte boundary.
144 func (b *bitWriter) flushAlign() {
145 nbBytes := (b.nBits + 7) >> 3
146 for i := range nbBytes {
147 b.out = append(b.out, byte(b.bitContainer>>(i*8)))
148 }
149 b.nBits = 0
150 b.bitContainer = 0
151 }
152 153 // close will write the alignment bit and write the final byte(s)
154 // to the output.
155 func (b *bitWriter) close() {
156 // End mark
157 b.addBits16Clean(1, 1)
158 // flush until next byte.
159 b.flushAlign()
160 }
161 162 // reset and continue writing by appending to out.
163 func (b *bitWriter) reset(out []byte) {
164 b.bitContainer = 0
165 b.nBits = 0
166 b.out = out
167 }
168