bitwriter.go

Documentation: github.com/klauspost/compress/fse

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