github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/recordio/recordioiov/recordioiov.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package recordioiov provides utility functions for dealing with [][]bytes, used
     6  // by recordio transformers.
     7  package recordioiov
     8  
     9  import (
    10  	"io"
    11  )
    12  
    13  // IOVecReader is an io.Reader adapter for [][]byte.
    14  type IOVecReader struct {
    15  	in       [][]byte
    16  	curIndex int
    17  	curBuf   []byte
    18  }
    19  
    20  // NewIOVecReader creates a new io.Reader that reads from the given buf.
    21  func NewIOVecReader(buf [][]byte) IOVecReader {
    22  	return IOVecReader{buf, -1, nil}
    23  }
    24  
    25  func (r *IOVecReader) fillBuf() error {
    26  	for len(r.curBuf) == 0 {
    27  		r.curIndex++
    28  		if r.curIndex >= len(r.in) {
    29  			return io.EOF
    30  		}
    31  		r.curBuf = r.in[r.curIndex]
    32  	}
    33  	return nil
    34  }
    35  
    36  // Read implements the io.Reader interface.
    37  func (r *IOVecReader) Read(p []byte) (int, error) {
    38  	if err := r.fillBuf(); err != nil {
    39  		return 0, err
    40  	}
    41  	bytes := len(p)
    42  	if bytes > len(r.curBuf) {
    43  		bytes = len(r.curBuf)
    44  	}
    45  	copy(p, r.curBuf)
    46  	r.curBuf = r.curBuf[bytes:]
    47  	return bytes, nil
    48  }
    49  
    50  // ReadByte reads one byte from the reader. On EOF, returns (0, io.EOF).
    51  func (r *IOVecReader) ReadByte() (byte, error) {
    52  	if err := r.fillBuf(); err != nil {
    53  		return 0, err
    54  	}
    55  	b := r.curBuf[0]
    56  	r.curBuf = r.curBuf[1:]
    57  	return b, nil
    58  }
    59  
    60  // TotalBytes returns the total size of the iovec.
    61  func TotalBytes(buf [][]byte) int {
    62  	size := 0
    63  	for _, b := range buf {
    64  		size += len(b)
    65  	}
    66  	return size
    67  }
    68  
    69  // Slice returns a byte slice of the given size. If cap(buf)>=size, then it
    70  // returns buf[:size]. Else it allocates a new slice on heap.
    71  func Slice(buf []byte, size int) []byte {
    72  	if cap(buf) >= size {
    73  		return buf[:size]
    74  	}
    75  	return make([]byte, size)
    76  }