cubefs

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
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
var 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.
36
func (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 31 bits.
42
// It will not check if there is space for them,
43
// so the caller must ensure that it has flushed recently.
44
func (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
// addBits64NC will add up to 64 bits.
50
// There must be space for 32 bits.
51
func (b *bitWriter) addBits64NC(value uint64, bits uint8) {
52
	if bits <= 31 {
53
		b.addBits32Clean(uint32(value), bits)
54
		return
55
	}
56
	b.addBits32Clean(uint32(value), 32)
57
	b.flush32()
58
	b.addBits32Clean(uint32(value>>32), bits-32)
59
}
60

61
// addBits32Clean will add up to 32 bits.
62
// It will not check if there is space for them.
63
// The input must not contain more bits than specified.
64
func (b *bitWriter) addBits32Clean(value uint32, bits uint8) {
65
	b.bitContainer |= uint64(value) << (b.nBits & 63)
66
	b.nBits += bits
67
}
68

69
// addBits16Clean will add up to 16 bits. value may not contain more set bits than indicated.
70
// It will not check if there is space for them, so the caller must ensure that it has flushed recently.
71
func (b *bitWriter) addBits16Clean(value uint16, bits uint8) {
72
	b.bitContainer |= uint64(value) << (b.nBits & 63)
73
	b.nBits += bits
74
}
75

76
// flush will flush all pending full bytes.
77
// There will be at least 56 bits available for writing when this has been called.
78
// Using flush32 is faster, but leaves less space for writing.
79
func (b *bitWriter) flush() {
80
	v := b.nBits >> 3
81
	switch v {
82
	case 0:
83
	case 1:
84
		b.out = append(b.out,
85
			byte(b.bitContainer),
86
		)
87
	case 2:
88
		b.out = append(b.out,
89
			byte(b.bitContainer),
90
			byte(b.bitContainer>>8),
91
		)
92
	case 3:
93
		b.out = append(b.out,
94
			byte(b.bitContainer),
95
			byte(b.bitContainer>>8),
96
			byte(b.bitContainer>>16),
97
		)
98
	case 4:
99
		b.out = append(b.out,
100
			byte(b.bitContainer),
101
			byte(b.bitContainer>>8),
102
			byte(b.bitContainer>>16),
103
			byte(b.bitContainer>>24),
104
		)
105
	case 5:
106
		b.out = append(b.out,
107
			byte(b.bitContainer),
108
			byte(b.bitContainer>>8),
109
			byte(b.bitContainer>>16),
110
			byte(b.bitContainer>>24),
111
			byte(b.bitContainer>>32),
112
		)
113
	case 6:
114
		b.out = append(b.out,
115
			byte(b.bitContainer),
116
			byte(b.bitContainer>>8),
117
			byte(b.bitContainer>>16),
118
			byte(b.bitContainer>>24),
119
			byte(b.bitContainer>>32),
120
			byte(b.bitContainer>>40),
121
		)
122
	case 7:
123
		b.out = append(b.out,
124
			byte(b.bitContainer),
125
			byte(b.bitContainer>>8),
126
			byte(b.bitContainer>>16),
127
			byte(b.bitContainer>>24),
128
			byte(b.bitContainer>>32),
129
			byte(b.bitContainer>>40),
130
			byte(b.bitContainer>>48),
131
		)
132
	case 8:
133
		b.out = append(b.out,
134
			byte(b.bitContainer),
135
			byte(b.bitContainer>>8),
136
			byte(b.bitContainer>>16),
137
			byte(b.bitContainer>>24),
138
			byte(b.bitContainer>>32),
139
			byte(b.bitContainer>>40),
140
			byte(b.bitContainer>>48),
141
			byte(b.bitContainer>>56),
142
		)
143
	default:
144
		panic(fmt.Errorf("bits (%d) > 64", b.nBits))
145
	}
146
	b.bitContainer >>= v << 3
147
	b.nBits &= 7
148
}
149

150
// flush32 will flush out, so there are at least 32 bits available for writing.
151
func (b *bitWriter) flush32() {
152
	if b.nBits < 32 {
153
		return
154
	}
155
	b.out = append(b.out,
156
		byte(b.bitContainer),
157
		byte(b.bitContainer>>8),
158
		byte(b.bitContainer>>16),
159
		byte(b.bitContainer>>24))
160
	b.nBits -= 32
161
	b.bitContainer >>= 32
162
}
163

164
// flushAlign will flush remaining full bytes and align to next byte boundary.
165
func (b *bitWriter) flushAlign() {
166
	nbBytes := (b.nBits + 7) >> 3
167
	for i := uint8(0); i < nbBytes; i++ {
168
		b.out = append(b.out, byte(b.bitContainer>>(i*8)))
169
	}
170
	b.nBits = 0
171
	b.bitContainer = 0
172
}
173

174
// close will write the alignment bit and write the final byte(s)
175
// to the output.
176
func (b *bitWriter) close() error {
177
	// End mark
178
	b.addBits16Clean(1, 1)
179
	// flush until next byte.
180
	b.flushAlign()
181
	return nil
182
}
183

184
// reset and continue writing by appending to out.
185
func (b *bitWriter) reset(out []byte) {
186
	b.bitContainer = 0
187
	b.nBits = 0
188
	b.out = out
189
}
190