From 1f2d896bf9e180684804b4a6e13ac00aa02cb3aa Mon Sep 17 00:00:00 2001
From: Peter Waller <p@pwaller.net>
Date: Tue, 13 Dec 2016 19:44:48 +0000
Subject: [PATCH] 	Add (*Reader).CopyNext(w) (int64, error) (#167)

* Add (*Reader).CopyNext(w) (int64, error)

It is useful to be able to efficiently copy objects without
decoding them.

My use case is filtering when I already know the indices of
the objects I want to keep, and for rewriting a dictionary
of objects as a column of objects.

* Remove ReadNext

* Add opportunistic optimization with m.R.Next

* Remove unused ReadNextError

* Remove commented code

* small fixup

- only call (*Reader).Next() when we're sure it won't realloc its buffer
- promote io.ErrUnexpectedEOF to msgp.ErrShortBytes
---
 msgp/read.go      | 52 +++++++++++++++++++++++++++++++++++++++++++++++
 msgp/read_test.go | 45 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/msgp/read.go b/msgp/read.go
index 6760cd7a..b3aaa350 100644
--- a/msgp/read.go
+++ b/msgp/read.go
@@ -146,6 +146,56 @@ func (m *Reader) Read(p []byte) (int, error) {
 	return m.R.Read(p)
 }
 
+// CopyNext reads the next object from m without decoding it and writes it to w.
+// It avoids unnecessary copies internally.
+func (m *Reader) CopyNext(w io.Writer) (int64, error) {
+	sz, o, err := getNextSize(m.R)
+	if err != nil {
+		return 0, err
+	}
+
+	var n int64
+	// Opportunistic optimization: if we can fit the whole thing in the m.R
+	// buffer, then just get a pointer to that, and pass it to w.Write,
+	// avoiding an allocation.
+	if int(sz) <= m.R.BufferSize() {
+		var nn int
+		var buf []byte
+		buf, err = m.R.Next(int(sz))
+		if err != nil {
+			if err == io.ErrUnexpectedEOF {
+				err = ErrShortBytes
+			}
+			return 0, err
+		}
+		nn, err = w.Write(buf)
+		n += int64(nn)
+	} else {
+		// Fall back to io.CopyN.
+		// May avoid allocating if w is a ReaderFrom (e.g. bytes.Buffer)
+		n, err = io.CopyN(w, m.R, int64(sz))
+		if err == io.ErrUnexpectedEOF {
+			err = ErrShortBytes
+		}
+	}
+	if err != nil {
+		return n, err
+	} else if n < int64(sz) {
+		return n, io.ErrShortWrite
+	}
+
+	// for maps and slices, read elements
+	for x := uintptr(0); x < o; x++ {
+		var n2 int64
+		n2, err = m.CopyNext(w)
+		if err != nil {
+			return n, err
+		}
+		n += n2
+	}
+	return n, nil
+}
+
 // ReadFull implements `io.ReadFull`
 func (m *Reader) ReadFull(p []byte) (int, error) {
 	return m.R.ReadFull(p)
@@ -194,8 +244,10 @@ func (m *Reader) IsNil() bool {
 	return err == nil && p[0] == mnil
 }
 
+// getNextSize returns the size of the next object on the wire.
 // returns (obj size, obj elements, error)
 // only maps and arrays have non-zero obj elements
+// for maps and arrays, obj size does not include elements
 //
 // use uintptr b/c it's guaranteed to be large enough
 // to hold whatever we can fit in memory.
diff --git a/msgp/read_test.go b/msgp/read_test.go
index 12b4dcfd..fd6a9922 100644
--- a/msgp/read_test.go
+++ b/msgp/read_test.go
@@ -745,3 +745,48 @@ func BenchmarkSkip(b *testing.B) {
 		}
 	}
 }
+
+func TestCopyNext(t *testing.T) {
+	var buf bytes.Buffer
+	en := NewWriter(&buf)
+
+	en.WriteMapHeader(6)
+
+	en.WriteString("thing_one")
+	en.WriteString("value_one")
+
+	en.WriteString("thing_two")
+	en.WriteFloat64(3.14159)
+
+	en.WriteString("some_bytes")
+	en.WriteBytes([]byte("nkl4321rqw908vxzpojnlk2314rqew098-s09123rdscasd"))
+
+	en.WriteString("the_time")
+	en.WriteTime(time.Now())
+
+	en.WriteString("what?")
+	en.WriteBool(true)
+
+	en.WriteString("ext")
+	en.WriteExtension(&RawExtension{Type: 55, Data: []byte("raw data!!!")})
+
+	en.Flush()
+
+	// Read from a copy of the original buf.
+	de := NewReader(bytes.NewReader(buf.Bytes()))
+
+	w := new(bytes.Buffer)
+
+	n, err := de.CopyNext(w)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != int64(buf.Len()) {
+		t.Fatalf("CopyNext returned the wrong value (%d != %d)",
+			n, buf.Len())
+	}
+
+	if !bytes.Equal(buf.Bytes(), w.Bytes()) {
+		t.Fatalf("not equal! %v, %v", buf.Bytes(), w.Bytes())
+	}
+}