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