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