github.com/coyove/sdss@v0.0.0-20231129015646-c2ec58cca6a2/contrib/roaring/internal/byte_input.go (about)

     1  package internal
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  )
     7  
     8  // ByteInput typed interface around io.Reader or raw bytes
     9  type ByteInput interface {
    10  	// Next returns a slice containing the next n bytes from the buffer,
    11  	// advancing the buffer as if the bytes had been returned by Read.
    12  	Next(n int) ([]byte, error)
    13  	// ReadUInt32 reads uint32 with LittleEndian order
    14  	ReadUInt32() (uint32, error)
    15  	// ReadUInt16 reads uint16 with LittleEndian order
    16  	ReadUInt16() (uint16, error)
    17  	// GetReadBytes returns read bytes
    18  	GetReadBytes() int64
    19  	// SkipBytes skips exactly n bytes
    20  	SkipBytes(n int) error
    21  }
    22  
    23  // NewByteInputFromReader creates reader wrapper
    24  func NewByteInputFromReader(reader io.Reader) ByteInput {
    25  	return &ByteInputAdapter{
    26  		r:         reader,
    27  		readBytes: 0,
    28  	}
    29  }
    30  
    31  // NewByteInput creates raw bytes wrapper
    32  func NewByteInput(buf []byte) ByteInput {
    33  	return &ByteBuffer{
    34  		buf: buf,
    35  		off: 0,
    36  	}
    37  }
    38  
    39  // ByteBuffer raw bytes wrapper
    40  type ByteBuffer struct {
    41  	buf []byte
    42  	off int
    43  }
    44  
    45  // Next returns a slice containing the next n bytes from the reader
    46  // If there are fewer bytes than the given n, io.ErrUnexpectedEOF will be returned
    47  func (b *ByteBuffer) Next(n int) ([]byte, error) {
    48  	m := len(b.buf) - b.off
    49  
    50  	if n > m {
    51  		return nil, io.ErrUnexpectedEOF
    52  	}
    53  
    54  	data := b.buf[b.off : b.off+n]
    55  	b.off += n
    56  
    57  	return data, nil
    58  }
    59  
    60  // ReadUInt32 reads uint32 with LittleEndian order
    61  func (b *ByteBuffer) ReadUInt32() (uint32, error) {
    62  	if len(b.buf)-b.off < 4 {
    63  		return 0, io.ErrUnexpectedEOF
    64  	}
    65  
    66  	v := binary.LittleEndian.Uint32(b.buf[b.off:])
    67  	b.off += 4
    68  
    69  	return v, nil
    70  }
    71  
    72  // ReadUInt16 reads uint16 with LittleEndian order
    73  func (b *ByteBuffer) ReadUInt16() (uint16, error) {
    74  	if len(b.buf)-b.off < 2 {
    75  		return 0, io.ErrUnexpectedEOF
    76  	}
    77  
    78  	v := binary.LittleEndian.Uint16(b.buf[b.off:])
    79  	b.off += 2
    80  
    81  	return v, nil
    82  }
    83  
    84  // GetReadBytes returns read bytes
    85  func (b *ByteBuffer) GetReadBytes() int64 {
    86  	return int64(b.off)
    87  }
    88  
    89  // SkipBytes skips exactly n bytes
    90  func (b *ByteBuffer) SkipBytes(n int) error {
    91  	m := len(b.buf) - b.off
    92  
    93  	if n > m {
    94  		return io.ErrUnexpectedEOF
    95  	}
    96  
    97  	b.off += n
    98  
    99  	return nil
   100  }
   101  
   102  // Reset resets the given buffer with a new byte slice
   103  func (b *ByteBuffer) Reset(buf []byte) {
   104  	b.buf = buf
   105  	b.off = 0
   106  }
   107  
   108  // ByteInputAdapter reader wrapper
   109  type ByteInputAdapter struct {
   110  	r         io.Reader
   111  	readBytes int
   112  	tmpbuf    []byte
   113  }
   114  
   115  // Next returns a slice containing the next n bytes from the buffer,
   116  // advancing the buffer as if the bytes had been returned by Read.
   117  func (b *ByteInputAdapter) Next(n int) ([]byte, error) {
   118  	return b._next(n, false)
   119  }
   120  
   121  func (b *ByteInputAdapter) _next(n int, tmp bool) ([]byte, error) {
   122  	var buf []byte
   123  	if tmp {
   124  		if cap(b.tmpbuf) < n {
   125  			b.tmpbuf = make([]byte, n)
   126  		}
   127  		buf = b.tmpbuf[:n]
   128  	} else {
   129  		buf = make([]byte, n)
   130  	}
   131  
   132  	m, err := io.ReadAtLeast(b.r, buf, n)
   133  	b.readBytes += m
   134  
   135  	if err != nil {
   136  		return nil, err
   137  	}
   138  
   139  	return buf, nil
   140  }
   141  
   142  // ReadUInt32 reads uint32 with LittleEndian order
   143  func (b *ByteInputAdapter) ReadUInt32() (uint32, error) {
   144  	buf, err := b._next(4, true)
   145  
   146  	if err != nil {
   147  		return 0, err
   148  	}
   149  
   150  	return binary.LittleEndian.Uint32(buf), nil
   151  }
   152  
   153  // ReadUInt16 reads uint16 with LittleEndian order
   154  func (b *ByteInputAdapter) ReadUInt16() (uint16, error) {
   155  	buf, err := b._next(2, true)
   156  
   157  	if err != nil {
   158  		return 0, err
   159  	}
   160  
   161  	return binary.LittleEndian.Uint16(buf), nil
   162  }
   163  
   164  // GetReadBytes returns read bytes
   165  func (b *ByteInputAdapter) GetReadBytes() int64 {
   166  	return int64(b.readBytes)
   167  }
   168  
   169  // SkipBytes skips exactly n bytes
   170  func (b *ByteInputAdapter) SkipBytes(n int) error {
   171  	_, err := b._next(n, true)
   172  
   173  	return err
   174  }
   175  
   176  // Reset resets the given buffer with a new stream
   177  func (b *ByteInputAdapter) Reset(stream io.Reader) {
   178  	b.r = stream
   179  	b.readBytes = 0
   180  }