github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/encoding/istream.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package encoding
    22  
    23  import (
    24  	"io"
    25  
    26  	"github.com/m3db/m3/src/dbnode/x/xio"
    27  )
    28  
    29  // IStream encapsulates a readable stream.
    30  type IStream struct {
    31  	r         xio.Reader64
    32  	current   uint64 // current uint64 we are working off of
    33  	index     int    // current index within data slice
    34  	remaining uint8  // bits remaining in current to be read
    35  }
    36  
    37  var (
    38  	_ io.ByteReader = (*IStream)(nil)
    39  	_ io.Reader     = (*IStream)(nil)
    40  )
    41  
    42  // NewIStream creates a new IStream
    43  func NewIStream(reader64 xio.Reader64) *IStream {
    44  	return &IStream{r: reader64}
    45  }
    46  
    47  // Read reads len(b) bytes.
    48  func (is *IStream) Read(b []byte) (int, error) {
    49  	var i int
    50  	for ; i < len(b); i++ {
    51  		res, err := is.ReadBits(8)
    52  		if err != nil {
    53  			return i, err
    54  		}
    55  		b[i] = byte(res)
    56  	}
    57  	return i, nil
    58  }
    59  
    60  // ReadByte reads the next Byte.
    61  func (is *IStream) ReadByte() (byte, error) {
    62  	res, err := is.ReadBits(8)
    63  	return byte(res), err
    64  }
    65  
    66  // ReadBit reads the next Bit.
    67  func (is *IStream) ReadBit() (Bit, error) {
    68  	res, err := is.ReadBits(1)
    69  	return Bit(res), err
    70  }
    71  
    72  // ReadBits reads the next Bits.
    73  func (is *IStream) ReadBits(numBits uint8) (uint64, error) {
    74  	res := is.current >> (64 - numBits)
    75  	remaining := is.remaining
    76  	if numBits <= remaining {
    77  		// Have enough bits buffered.
    78  		is.current <<= numBits
    79  		is.remaining -= numBits
    80  		return res, nil
    81  	}
    82  
    83  	// Not enough bits buffered, read next word from the stream.
    84  	bitsNeeded := numBits - remaining
    85  
    86  	current, n, err := is.r.Read64()
    87  	if err != nil {
    88  		return 0, err
    89  	}
    90  	n *= 8
    91  	if n < bitsNeeded {
    92  		return 0, io.EOF
    93  	}
    94  
    95  	is.current = current << bitsNeeded
    96  	is.remaining = n - bitsNeeded
    97  	return res | current>>(64-bitsNeeded), nil
    98  }
    99  
   100  // PeekBits looks at the next Bits, but doesn't move the pos.
   101  func (is *IStream) PeekBits(numBits uint8) (uint64, error) {
   102  	if numBits <= is.remaining {
   103  		return readBitsInWord(is.current, numBits), nil
   104  	}
   105  	res := readBitsInWord(is.current, numBits)
   106  	bitsNeeded := numBits - is.remaining
   107  	next, bytes, err := is.r.Peek64()
   108  	if err != nil {
   109  		return 0, err
   110  	}
   111  	if rem := 8 * bytes; rem < bitsNeeded {
   112  		return 0, io.EOF
   113  	}
   114  	return res | readBitsInWord(next, bitsNeeded), nil
   115  }
   116  
   117  // RemainingBitsInCurrentByte returns the number of bits remaining to be read in the current byte.
   118  func (is *IStream) RemainingBitsInCurrentByte() uint {
   119  	return uint(is.remaining % 8)
   120  }
   121  
   122  // readBitsInWord reads the first numBits in word w.
   123  func readBitsInWord(w uint64, numBits uint8) uint64 {
   124  	return w >> (64 - numBits)
   125  }
   126  
   127  // Reset resets the IStream.
   128  func (is *IStream) Reset(reader xio.Reader64) {
   129  	is.current = 0
   130  	is.remaining = 0
   131  	is.index = 0
   132  	is.r = reader
   133  }