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

     1  package morebufio
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/Schaudge/grailbase/ioctx"
     7  )
     8  
     9  // PeekBackReader is a Reader augmented with a function to "peek" backwards at the data that
    10  // was already passed. Peeking does not change the current stream position (that is, PeekBack has
    11  // no effect on the next Read).
    12  type PeekBackReader interface {
    13  	ioctx.Reader
    14  	// PeekBack returns a fixed "window" of the data that Read has already returned, ending at
    15  	// the current Read position. It may be smaller at the start until enough data has been read,
    16  	// but after that it's constant size. PeekBack is allowed after Read returns EOF.
    17  	// The returned slice aliases the internal buffer and is invalidated by the next Read.
    18  	PeekBack() []byte
    19  }
    20  
    21  type peekBackReader struct {
    22  	r   ioctx.Reader
    23  	buf []byte
    24  }
    25  
    26  // NewPeekBackReader returns a PeekBackReader. It doesn't have a "forward" buffer so small
    27  // PeekBackReader.Read operations cause small r.Read operations.
    28  func NewPeekBackReader(r ioctx.Reader, peekBackSize int) PeekBackReader {
    29  	return &peekBackReader{r, make([]byte, 0, peekBackSize)}
    30  }
    31  
    32  func (p *peekBackReader) Read(ctx context.Context, dst []byte) (int, error) {
    33  	nRead, err := p.r.Read(ctx, dst)
    34  	dst = dst[:nRead]
    35  	// First, grow the peek buf until cap, since it starts empty.
    36  	if grow := cap(p.buf) - len(p.buf); grow > 0 {
    37  		if len(dst) < grow {
    38  			grow = len(dst)
    39  		}
    40  		p.buf = append(p.buf, dst[:grow]...)
    41  		dst = dst[grow:]
    42  	}
    43  	if len(dst) == 0 {
    44  		return nRead, err
    45  	}
    46  	// Shift data if any part of the peek buf is still valid.
    47  	updateTail := p.buf
    48  	if len(dst) < len(p.buf) {
    49  		n := copy(p.buf, p.buf[len(dst):])
    50  		updateTail = p.buf[n:]
    51  	}
    52  	_ = copy(updateTail, dst[len(dst)-len(updateTail):])
    53  	return nRead, err
    54  }
    55  
    56  func (p *peekBackReader) PeekBack() []byte { return p.buf }