bitwriter.go

Documentation: github.com/klauspost/compress/zstd

     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  // bitWriter will write bits.
     9  // First bit will be LSB of the first byte of output.
    10  type bitWriter struct {
    11  	bitContainer uint64
    12  	nBits        uint8
    13  	out          []byte
    14  }
    15  
    16  // bitMask16 is bitmasks. Has extra to avoid bounds check.
    17  var bitMask16 = [32]uint16{
    18  	0, 1, 3, 7, 0xF, 0x1F,
    19  	0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF, 0x7FF,
    20  	0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF, 0xFFFF,
    21  	0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF,
    22  	0xFFFF, 0xFFFF} /* up to 16 bits */
    23  
    24  var bitMask32 = [32]uint32{
    25  	0, 1, 3, 7, 0xF, 0x1F, 0x3F, 0x7F, 0xFF,
    26  	0x1FF, 0x3FF, 0x7FF, 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF,
    27  	0x1ffff, 0x3ffff, 0x7FFFF, 0xfFFFF, 0x1fFFFF, 0x3fFFFF, 0x7fFFFF, 0xffFFFF,
    28  	0x1ffFFFF, 0x3ffFFFF, 0x7ffFFFF, 0xfffFFFF, 0x1fffFFFF, 0x3fffFFFF, 0x7fffFFFF,
    29  } // up to 32 bits
    30  
    31  // addBits16NC will add up to 16 bits.
    32  // It will not check if there is space for them,
    33  // so the caller must ensure that it has flushed recently.
    34  func (b *bitWriter) addBits16NC(value uint16, bits uint8) {
    35  	b.bitContainer |= uint64(value&bitMask16[bits&31]) << (b.nBits & 63)
    36  	b.nBits += bits
    37  }
    38  
    39  // addBits32NC will add up to 31 bits.
    40  // It will not check if there is space for them,
    41  // so the caller must ensure that it has flushed recently.
    42  func (b *bitWriter) addBits32NC(value uint32, bits uint8) {
    43  	b.bitContainer |= uint64(value&bitMask32[bits&31]) << (b.nBits & 63)
    44  	b.nBits += bits
    45  }
    46  
    47  // addBits64NC will add up to 64 bits.
    48  // There must be space for 32 bits.
    49  func (b *bitWriter) addBits64NC(value uint64, bits uint8) {
    50  	if bits <= 31 {
    51  		b.addBits32Clean(uint32(value), bits)
    52  		return
    53  	}
    54  	b.addBits32Clean(uint32(value), 32)
    55  	b.flush32()
    56  	b.addBits32Clean(uint32(value>>32), bits-32)
    57  }
    58  
    59  // addBits32Clean will add up to 32 bits.
    60  // It will not check if there is space for them.
    61  // The input must not contain more bits than specified.
    62  func (b *bitWriter) addBits32Clean(value uint32, bits uint8) {
    63  	b.bitContainer |= uint64(value) << (b.nBits & 63)
    64  	b.nBits += bits
    65  }
    66  
    67  // addBits16Clean will add up to 16 bits. value may not contain more set bits than indicated.
    68  // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
    69  func (b *bitWriter) addBits16Clean(value uint16, bits uint8) {
    70  	b.bitContainer |= uint64(value) << (b.nBits & 63)
    71  	b.nBits += bits
    72  }
    73  
    74  // flush32 will flush out, so there are at least 32 bits available for writing.
    75  func (b *bitWriter) flush32() {
    76  	if b.nBits < 32 {
    77  		return
    78  	}
    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  	b.nBits -= 32
    85  	b.bitContainer >>= 32
    86  }
    87  
    88  // flushAlign will flush remaining full bytes and align to next byte boundary.
    89  func (b *bitWriter) flushAlign() {
    90  	nbBytes := (b.nBits + 7) >> 3
    91  	for i := uint8(0); i < nbBytes; i++ {
    92  		b.out = append(b.out, byte(b.bitContainer>>(i*8)))
    93  	}
    94  	b.nBits = 0
    95  	b.bitContainer = 0
    96  }
    97  
    98  // close will write the alignment bit and write the final byte(s)
    99  // to the output.
   100  func (b *bitWriter) close() {
   101  	// End mark
   102  	b.addBits16Clean(1, 1)
   103  	// flush until next byte.
   104  	b.flushAlign()
   105  }
   106  
   107  // reset and continue writing by appending to out.
   108  func (b *bitWriter) reset(out []byte) {
   109  	b.bitContainer = 0
   110  	b.nBits = 0
   111  	b.out = out
   112  }
   113
View as plain text