github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/segment/writer/metrics/compress/bit_reader.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  	"errors"
    21  	"fmt"
    22  	"io"
    23  )
    24  
    25  // A reader reads bits from an io.reader
    26  type bitReader struct {
    27  	r      io.Reader
    28  	buffer [1]byte
    29  	count  uint8 // The number of right-most bits valid to read (from left) in the current 8 byte buffer.
    30  }
    31  
    32  // newReader returns a reader that returns a single bit at a time from 'r'
    33  func newBitReader(r io.Reader) *bitReader {
    34  	return &bitReader{r: r}
    35  }
    36  
    37  // readBit returns the next bit from the stream, reading a new byte
    38  // from the underlying reader if required.
    39  func (b *bitReader) readBit() (bit, error) {
    40  	if b.count == 0 {
    41  		n, err := b.r.Read(b.buffer[:])
    42  		if err != nil {
    43  			return zero, fmt.Errorf("failed to read a byte: %w", err)
    44  		}
    45  		if n != 1 {
    46  			return zero, errors.New("read more than a byte")
    47  		}
    48  		b.count = 8
    49  	}
    50  	b.count--
    51  	// bitwise AND
    52  	// (e.g.)
    53  	// 11111111 & 10000000 = 10000000
    54  	// 11000011 & 10000000 = 10000000
    55  	d := (b.buffer[0] & 0x80)
    56  	// Left shift to read next bit
    57  	b.buffer[0] <<= 1
    58  	return d != 0, nil
    59  }
    60  
    61  // readBits constructs a uint64 with the nbits right-most bits
    62  // read from the stream, and any other bits 0.
    63  func (b *bitReader) readByte() (byte, error) {
    64  	if b.count == 0 {
    65  		n, err := b.r.Read(b.buffer[:])
    66  		if err != nil {
    67  			return b.buffer[0], fmt.Errorf("failed to read a byte: %w", err)
    68  		}
    69  		if n != 1 {
    70  			return b.buffer[0], errors.New("read more than a byte")
    71  		}
    72  		return b.buffer[0], nil
    73  	}
    74  
    75  	byt := b.buffer[0]
    76  
    77  	n, err := b.r.Read(b.buffer[:])
    78  	if err != nil {
    79  		return 0, fmt.Errorf("failed to read a byte: %w", err)
    80  	}
    81  	if n != 1 {
    82  		return b.buffer[0], errors.New("read more than a byte")
    83  	}
    84  
    85  	byt |= b.buffer[0] >> b.count
    86  	b.buffer[0] <<= (8 - b.count)
    87  
    88  	return byt, nil
    89  }
    90  
    91  // readBits reads nbits from the stream
    92  func (b *bitReader) readBits(nbits int) (uint64, error) {
    93  	var u uint64
    94  
    95  	for 8 <= nbits {
    96  		byt, err := b.readByte()
    97  		if err != nil {
    98  			return 0, err
    99  		}
   100  
   101  		u = (u << 8) | uint64(byt)
   102  		nbits -= 8
   103  	}
   104  
   105  	var err error
   106  	for nbits > 0 && err != io.EOF {
   107  		byt, err := b.readBit()
   108  		if err != nil {
   109  			return 0, err
   110  		}
   111  		u <<= 1
   112  		if byt {
   113  			u |= 1
   114  		}
   115  		nbits--
   116  	}
   117  
   118  	return u, nil
   119  }