1// Copyright 2015, Joe Tsai. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE.md file.
4
5package prefix
6
7import (
8	"encoding/binary"
9	"io"
10
11	"github.com/dsnet/compress/internal/errors"
12)
13
14// Writer implements a prefix encoder. For performance reasons, Writer will not
15// write bytes immediately to the underlying stream.
16type Writer struct {
17	Offset int64 // Number of bytes written to the underlying io.Writer
18
19	wr        io.Writer
20	bufBits   uint64 // Buffer to hold some bits
21	numBits   uint   // Number of valid bits in bufBits
22	bigEndian bool   // Are bits written in big-endian order?
23
24	buf    [512]byte
25	cntBuf int
26}
27
28// Init initializes the bit Writer to write to w. If bigEndian is true, then
29// bits will be written starting from the most-significant bits of a byte
30// (as done in bzip2), otherwise it will write starting from the
31// least-significant bits of a byte (such as for deflate and brotli).
32func (pw *Writer) Init(w io.Writer, bigEndian bool) {
33	*pw = Writer{wr: w, bigEndian: bigEndian}
34	return
35}
36
37// BitsWritten reports the total number of bits issued to any Write method.
38func (pw *Writer) BitsWritten() int64 {
39	return 8*pw.Offset + 8*int64(pw.cntBuf) + int64(pw.numBits)
40}
41
42// WritePads writes 0-7 bits to the bit buffer to achieve byte-alignment.
43func (pw *Writer) WritePads(v uint) {
44	nb := -pw.numBits & 7
45	pw.bufBits |= uint64(v) << pw.numBits
46	pw.numBits += nb
47}
48
49// Write writes bytes from buf.
50// The bit-ordering mode does not affect this method.
51func (pw *Writer) Write(buf []byte) (cnt int, err error) {
52	if pw.numBits > 0 || pw.cntBuf > 0 {
53		if pw.numBits%8 != 0 {
54			return 0, errorf(errors.Invalid, "non-aligned bit buffer")
55		}
56		if _, err := pw.Flush(); err != nil {
57			return 0, err
58		}
59	}
60	cnt, err = pw.wr.Write(buf)
61	pw.Offset += int64(cnt)
62	return cnt, err
63}
64
65// WriteOffset writes ofs in a (sym, extra) fashion using the provided prefix
66// Encoder and RangeEncoder.
67func (pw *Writer) WriteOffset(ofs uint, pe *Encoder, re *RangeEncoder) {
68	sym := re.Encode(ofs)
69	pw.WriteSymbol(sym, pe)
70	rc := re.rcs[sym]
71	pw.WriteBits(ofs-uint(rc.Base), uint(rc.Len))
72}
73
74// TryWriteBits attempts to write nb bits using the contents of the bit buffer
75// alone. It reports whether it succeeded.
76//
77// This method is designed to be inlined for performance reasons.
78func (pw *Writer) TryWriteBits(v, nb uint) bool {
79	if 64-pw.numBits < nb {
80		return false
81	}
82	pw.bufBits |= uint64(v) << pw.numBits
83	pw.numBits += nb
84	return true
85}
86
87// WriteBits writes nb bits of v to the underlying writer.
88func (pw *Writer) WriteBits(v, nb uint) {
89	if _, err := pw.PushBits(); err != nil {
90		errors.Panic(err)
91	}
92	pw.bufBits |= uint64(v) << pw.numBits
93	pw.numBits += nb
94}
95
96// TryWriteSymbol attempts to encode the next symbol using the contents of the
97// bit buffer alone. It reports whether it succeeded.
98//
99// This method is designed to be inlined for performance reasons.
100func (pw *Writer) TryWriteSymbol(sym uint, pe *Encoder) bool {
101	chunk := pe.chunks[uint32(sym)&pe.chunkMask]
102	nb := uint(chunk & countMask)
103	if 64-pw.numBits < nb {
104		return false
105	}
106	pw.bufBits |= uint64(chunk>>countBits) << pw.numBits
107	pw.numBits += nb
108	return true
109}
110
111// WriteSymbol writes the symbol using the provided prefix Encoder.
112func (pw *Writer) WriteSymbol(sym uint, pe *Encoder) {
113	if _, err := pw.PushBits(); err != nil {
114		errors.Panic(err)
115	}
116	chunk := pe.chunks[uint32(sym)&pe.chunkMask]
117	nb := uint(chunk & countMask)
118	pw.bufBits |= uint64(chunk>>countBits) << pw.numBits
119	pw.numBits += nb
120}
121
122// Flush flushes all complete bytes from the bit buffer to the byte buffer, and
123// then flushes all bytes in the byte buffer to the underlying writer.
124// After this call, the bit Writer is will only withhold 7 bits at most.
125func (pw *Writer) Flush() (int64, error) {
126	if pw.numBits < 8 && pw.cntBuf == 0 {
127		return pw.Offset, nil
128	}
129	if _, err := pw.PushBits(); err != nil {
130		return pw.Offset, err
131	}
132	cnt, err := pw.wr.Write(pw.buf[:pw.cntBuf])
133	pw.cntBuf -= cnt
134	pw.Offset += int64(cnt)
135	return pw.Offset, err
136}
137
138// PushBits pushes as many bytes as possible from the bit buffer to the byte
139// buffer, reporting the number of bits pushed.
140func (pw *Writer) PushBits() (uint, error) {
141	if pw.cntBuf >= len(pw.buf)-8 {
142		cnt, err := pw.wr.Write(pw.buf[:pw.cntBuf])
143		pw.cntBuf -= cnt
144		pw.Offset += int64(cnt)
145		if err != nil {
146			return 0, err
147		}
148	}
149
150	u := pw.bufBits
151	if pw.bigEndian {
152		// Swap all the bits within each byte.
153		u = (u&0xaaaaaaaaaaaaaaaa)>>1 | (u&0x5555555555555555)<<1
154		u = (u&0xcccccccccccccccc)>>2 | (u&0x3333333333333333)<<2
155		u = (u&0xf0f0f0f0f0f0f0f0)>>4 | (u&0x0f0f0f0f0f0f0f0f)<<4
156	}
157	// Starting with Go 1.7, the compiler should use a wide integer
158	// store here if the architecture supports it.
159	binary.LittleEndian.PutUint64(pw.buf[pw.cntBuf:], u)
160
161	nb := pw.numBits / 8 // Number of bytes to copy from bit buffer
162	pw.cntBuf += int(nb)
163	pw.bufBits >>= 8 * nb
164	pw.numBits -= 8 * nb
165	return 8 * nb, nil
166}
167