github.com/apernet/quic-go@v0.43.1-0.20240515053213-5e9e635fd9f0/quicvarint/io.go (about)

     1  package quicvarint
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  )
     7  
     8  // Reader implements both the io.ByteReader and io.Reader interfaces.
     9  type Reader interface {
    10  	io.ByteReader
    11  	io.Reader
    12  }
    13  
    14  var _ Reader = &bytes.Reader{}
    15  
    16  type byteReader struct {
    17  	io.Reader
    18  }
    19  
    20  var _ Reader = &byteReader{}
    21  
    22  // NewReader returns a Reader for r.
    23  // If r already implements both io.ByteReader and io.Reader, NewReader returns r.
    24  // Otherwise, r is wrapped to add the missing interfaces.
    25  func NewReader(r io.Reader) Reader {
    26  	if r, ok := r.(Reader); ok {
    27  		return r
    28  	}
    29  	return &byteReader{r}
    30  }
    31  
    32  func (r *byteReader) ReadByte() (byte, error) {
    33  	var b [1]byte
    34  	n, err := r.Reader.Read(b[:])
    35  	if n == 1 && err == io.EOF {
    36  		err = nil
    37  	}
    38  	return b[0], err
    39  }
    40  
    41  // Writer implements both the io.ByteWriter and io.Writer interfaces.
    42  type Writer interface {
    43  	io.ByteWriter
    44  	io.Writer
    45  }
    46  
    47  var _ Writer = &bytes.Buffer{}
    48  
    49  type byteWriter struct {
    50  	io.Writer
    51  }
    52  
    53  var _ Writer = &byteWriter{}
    54  
    55  // NewWriter returns a Writer for w.
    56  // If r already implements both io.ByteWriter and io.Writer, NewWriter returns w.
    57  // Otherwise, w is wrapped to add the missing interfaces.
    58  func NewWriter(w io.Writer) Writer {
    59  	if w, ok := w.(Writer); ok {
    60  		return w
    61  	}
    62  	return &byteWriter{w}
    63  }
    64  
    65  func (w *byteWriter) WriteByte(c byte) error {
    66  	_, err := w.Writer.Write([]byte{c})
    67  	return err
    68  }