Switch to non-buffering byte reader

gen.go

@@ -589,16 +589,14 @@ func emitCborUnmarshalStructField(w io.Writer, f Field) error {
 		return doTemplate(w, f, `
 	{
 {{ if .Pointer }}
-		pb, err := br.PeekByte()
+		b, err := br.ReadByte()
 		if err != nil {
 			return err
 		}
-		if pb == cbg.CborNull[0] {
-			var nbuf [1]byte
-			if _, err := br.Read(nbuf[:]); err != nil {
+		if b != cbg.CborNull[0] {
+			if err := br.UnreadByte(); err != nil {
 				return err
 			}
-		} else {
 {{ end }}
 		c, err := cbg.ReadCid(br)
 		if err != nil {
@@ -628,16 +626,14 @@ func emitCborUnmarshalStructField(w io.Writer, f Field) error {
 		return doTemplate(w, f, `
 	{
 {{ if .Pointer }}
-		pb, err := br.PeekByte()
+		b, err := br.ReadByte()
 		if err != nil {
 			return err
 		}
-		if pb == cbg.CborNull[0] {
-			var nbuf [1]byte
-			if _, err := br.Read(nbuf[:]); err != nil {
+		if b != cbg.CborNull[0] {
+			if err := br.UnreadByte(); err != nil {
 				return err
 			}
-		} else {
 			{{ .Name }} = new({{ .TypeName }})
 			if err := {{ .Name }}.UnmarshalCBOR(br); err != nil {
 				return xerrors.Errorf("unmarshaling {{ .Name }} pointer: %w", err)
@@ -685,16 +681,14 @@ func emitCborUnmarshalUint64Field(w io.Writer, f Field) error {
 	return doTemplate(w, f, `
 	{
 {{ if .Pointer }}
-	pb, err := br.PeekByte()
+	b, err := br.ReadByte()
 	if err != nil {
 		return err
 	}
-	if pb == cbg.CborNull[0] {
-		var nbuf [1]byte
-		if _, err := br.Read(nbuf[:]); err != nil {
+	if b != cbg.CborNull[0] {
+		if err := br.UnreadByte(); err != nil {
 			return err
 		}
-	} else {
 		maj, extra, err = {{ ReadHeader "br" }}
 		if err != nil {
 			return err

peeker.go

@@ -0,0 +1,80 @@
+package typegen
+
+import (
+	"bufio"
+	"io"
+)
+
+// BytePeeker combines the Reader and ByteScanner interfaces.
+type BytePeeker interface {
+	io.Reader
+	io.ByteScanner
+}
+
+func GetPeeker(r io.Reader) BytePeeker {
+	if r, ok := r.(BytePeeker); ok {
+		return r
+	}
+	return &peeker{reader: r}
+}
+
+// peeker is a non-buffering BytePeeker.
+type peeker struct {
+	reader    io.Reader
+	peekState int
+	lastByte  byte
+}
+
+const (
+	peekEmpty = iota
+	peekSet
+	peekUnread
+)
+
+func (p *peeker) Read(buf []byte) (n int, err error) {
+	// Read "nothing". I.e., read an error, maybe.
+	if len(buf) == 0 {
+		// There's something pending in the
+		if p.peekState == peekUnread {
+			return 0, nil
+		}
+		return p.reader.Read(nil)
+	}
+
+	if p.peekState == peekUnread {
+		buf[0] = p.lastByte
+		n, err = p.reader.Read(buf[1:])
+		n += 1
+	} else {
+		n, err = p.reader.Read(buf)
+	}
+	if n > 0 {
+		p.peekState = peekSet
+		p.lastByte = buf[n-1]
+	}
+	return n, err
+}
+
+func (p *peeker) ReadByte() (byte, error) {
+	if p.peekState == peekUnread {
+		p.peekState = peekSet
+		return p.lastByte, nil
+	}
+	var buf [1]byte
+	n, err := p.reader.Read(buf[:])
+	if n == 0 {
+		return 0, err
+	}
+	b := buf[0]
+	p.lastByte = b
+	p.peekState = peekSet
+	return b, err
+}
+
+func (p *peeker) UnreadByte() error {
+	if p.peekState != peekSet {
+		return bufio.ErrInvalidUnreadByte
+	}
+	p.peekState = peekUnread
+	return nil
+}

peeker_test.go

@@ -0,0 +1,103 @@
+package typegen
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"testing"
+)
+
+func TestPeeker(t *testing.T) {
+	buf := bytes.NewBuffer([]byte{0, 1, 2, 3})
+	p := peeker{reader: buf}
+	n, err := p.Read(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 0 {
+		t.Fatal(err)
+	}
+
+	err = p.UnreadByte()
+	if err != bufio.ErrInvalidUnreadByte {
+		t.Fatal(err)
+	}
+
+	// read 2 bytes
+	var out [2]byte
+	n, err = p.Read(out[:])
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 2 {
+		t.Fatalf("expected 2 bytes, got %d", n)
+	}
+	if !bytes.Equal(out[:], []byte{0, 1}) {
+		t.Fatalf("unexpected output")
+	}
+
+	// unread that last byte and read it again.
+	err = p.UnreadByte()
+	if err != nil {
+		t.Fatal(err)
+	}
+	b, err := p.ReadByte()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if b != 1 {
+		t.Fatal("expected 1")
+	}
+
+	// unread that last byte then read 2
+	err = p.UnreadByte()
+	if err != nil {
+		t.Fatal(err)
+	}
+	n, err = p.Read(out[:])
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 2 {
+		t.Fatalf("expected 2 bytes, got %d", n)
+	}
+	if !bytes.Equal(out[:], []byte{1, 2}) {
+		t.Fatalf("unexpected output")
+	}
+
+	// read another byte
+	b, err = p.ReadByte()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if b != 3 {
+		t.Fatal("expected 1")
+	}
+
+	// Should read eof at end.
+	n, err = p.Read(out[:])
+	if err != io.EOF {
+		t.Fatal(err)
+	}
+	if n != 0 {
+		t.Fatal("should have been at end")
+	}
+	// should unread eof
+	err = p.UnreadByte()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = p.Read(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b, err = p.ReadByte()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if b != 3 {
+		t.Fatal("expected 1")
+	}
+}