github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/segment/writer/metrics/compress/bit_writer.go (about)

     1  /*
     2  Copyright 2023.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package compress
    18  
    19  import (
    20  	"fmt"
    21  	"io"
    22  )
    23  
    24  type bitWriter struct {
    25  	w      io.Writer
    26  	buffer byte
    27  	count  uint8 // How many right-most bits are available for writing in the current byte (the last byte of the buffer).
    28  }
    29  
    30  type bit bool
    31  
    32  const (
    33  	zero bit = false
    34  	one  bit = true
    35  )
    36  
    37  // newBitWriter returns a writer that buffers bits and write the resulting bytes to 'w'
    38  func newBitWriter(w io.Writer) *bitWriter {
    39  	return &bitWriter{
    40  		w: w, count: 8,
    41  	}
    42  }
    43  
    44  // writeBit writes a single bit.
    45  func (b *bitWriter) writeBit(bit bit) error {
    46  	if bit {
    47  		b.buffer |= 1 << (b.count - 1)
    48  	}
    49  
    50  	b.count--
    51  
    52  	if b.count == 0 {
    53  		if _, err := b.w.Write([]byte{b.buffer}); err != nil {
    54  			return fmt.Errorf("failed to write a bit: %w", err)
    55  		}
    56  		b.buffer = 0
    57  		b.count = 8
    58  	}
    59  
    60  	return nil
    61  }
    62  
    63  // writeBits writes the nbits right-most bits of u64 to the buffer in left-to-right order.
    64  func (b *bitWriter) writeBits(u64 uint64, nbits int) error {
    65  	u64 <<= (64 - uint(nbits))
    66  	for nbits >= 8 {
    67  		byt := byte(u64 >> 56)
    68  		err := b.writeByte(byt)
    69  		if err != nil {
    70  			return err
    71  		}
    72  		u64 <<= 8
    73  		nbits -= 8
    74  	}
    75  
    76  	for nbits > 0 {
    77  		err := b.writeBit((u64 >> 63) == 1)
    78  		if err != nil {
    79  			return err
    80  		}
    81  		u64 <<= 1
    82  		nbits--
    83  	}
    84  
    85  	return nil
    86  }
    87  
    88  // writeByte writes a single byte to the stream, regardless of alignment
    89  func (b *bitWriter) writeByte(byt byte) error {
    90  	// Complete the last byte with the leftmost b.buffer bits from byt.
    91  	b.buffer |= byt >> (8 - b.count)
    92  
    93  	if _, err := b.w.Write([]byte{b.buffer}); err != nil {
    94  		return fmt.Errorf("failed to write a byte: %w", err)
    95  	}
    96  	b.buffer = byt << b.count
    97  
    98  	return nil
    99  }
   100  
   101  // flush empties the currently in-process byte by filling it with 'bit'.
   102  func (b *bitWriter) flush(bit bit) error {
   103  	for b.count != 8 {
   104  		err := b.writeBit(bit)
   105  		if err != nil {
   106  			return err
   107  		}
   108  	}
   109  
   110  	return nil
   111  }