blob: 43e463611b15ccd0f92710da6c6ac8502dbf009a [file] [log] [blame]
khenaidoo106c61a2021-08-11 18:05:46 -04001// 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
6package fse
7
8import "fmt"
9
10// bitWriter will write bits.
11// First bit will be LSB of the first byte of output.
12type bitWriter struct {
13 bitContainer uint64
14 nBits uint8
15 out []byte
16}
17
18// bitMask16 is bitmasks. Has extra to avoid bounds check.
19var 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.
29func (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.
36func (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.
45func (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.
58func (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.
130func (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.
144func (b *bitWriter) flushAlign() {
145 nbBytes := (b.nBits + 7) >> 3
146 for i := uint8(0); i < nbBytes; i++ {
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.
155func (b *bitWriter) close() error {
156 // End mark
157 b.addBits16Clean(1, 1)
158 // flush until next byte.
159 b.flushAlign()
160 return nil
161}
162
163// reset and continue writing by appending to out.
164func (b *bitWriter) reset(out []byte) {
165 b.bitContainer = 0
166 b.nBits = 0
167 b.out = out
168}