rlp: add Stream.Reset and accept any reader (for p2p)

2 files changed, 66 insertions, 7 deletions

diff --git a/rlp/decode.go b/rlp/decode.go
index 96d912f56..565c84790 100644
--- a/rlp/decode.go
+++ b/rlp/decode.go
@@ -1,6 +1,7 @@
 package rlp
 
 import (
+	"bufio"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -24,8 +25,9 @@ type Decoder interface {
 	DecodeRLP(*Stream) error
 }
 
-// Decode parses RLP-encoded data from r and stores the result
-// in the value pointed to by val. Val must be a non-nil pointer.
+// Decode parses RLP-encoded data from r and stores the result in the
+// value pointed to by val. Val must be a non-nil pointer. If r does
+// not implement ByteReader, Decode will do its own buffering.
 //
 // Decode uses the following type-dependent decoding rules:
 //
@@ -66,7 +68,7 @@ type Decoder interface {
 //
 // Non-empty interface types are not supported, nor are bool, float32,
 // float64, maps, channel types and functions.
-func Decode(r ByteReader, val interface{}) error {
+func Decode(r io.Reader, val interface{}) error {
 	return NewStream(r).Decode(val)
 }
 
@@ -432,8 +434,14 @@ type Stream struct {
 
 type listpos struct{ pos, size uint64 }
 
-func NewStream(r ByteReader) *Stream {
-	return &Stream{r: r, uintbuf: make([]byte, 8), kind: -1}
+// NewStream creates a new stream reading from r.
+// If r does not implement ByteReader, the Stream will
+// introduce its own buffering.
+func NewStream(r io.Reader) *Stream {
+	s := new(Stream)
+	s.Reset(r)
+	return s
+}
 }
 
 // Bytes reads an RLP string and returns its contents as a byte slice.
@@ -543,6 +551,23 @@ func (s *Stream) Decode(val interface{}) error {
 	return info.decoder(s, rval.Elem())
 }
 
+// Reset discards any information about the current decoding context
+// and starts reading from r. If r does not also implement ByteReader,
+// Stream will do its own buffering.
+func (s *Stream) Reset(r io.Reader) {
+	bufr, ok := r.(ByteReader)
+	if !ok {
+		bufr = bufio.NewReader(r)
+	}
+	s.r = bufr
+	s.stack = s.stack[:0]
+	s.size = 0
+	s.kind = -1
+	if s.uintbuf == nil {
+		s.uintbuf = make([]byte, 8)
+	}
+}
+
 // Kind returns the kind and size of the next value in the
 // input stream.
 //
diff --git a/rlp/decode_test.go b/rlp/decode_test.go
index eb1618299..9d320564b 100644
--- a/rlp/decode_test.go
+++ b/rlp/decode_test.go
@@ -286,14 +286,14 @@ var decodeTests = []decodeTest{
 
 func intp(i int) *int { return &i }
 
-func TestDecode(t *testing.T) {
+func runTests(t *testing.T, decode func([]byte, interface{}) error) {
 	for i, test := range decodeTests {
 		input, err := hex.DecodeString(test.input)
 		if err != nil {
 			t.Errorf("test %d: invalid hex input %q", i, test.input)
 			continue
 		}
-		err = Decode(bytes.NewReader(input), test.ptr)
+		err = decode(input, test.ptr)
 		if err != nil && test.error == nil {
 			t.Errorf("test %d: unexpected Decode error: %v\ndecoding into %T\ninput %q",
 				i, err, test.ptr, test.input)
@@ -312,6 +312,40 @@ func TestDecode(t *testing.T) {
 	}
 }
 
+func TestDecodeWithByteReader(t *testing.T) {
+	runTests(t, func(input []byte, into interface{}) error {
+		return Decode(bytes.NewReader(input), into)
+	})
+}
+
+// dumbReader reads from a byte slice but does not
+// implement ReadByte.
+type dumbReader []byte
+
+func (r *dumbReader) Read(buf []byte) (n int, err error) {
+	if len(*r) == 0 {
+		return 0, io.EOF
+	}
+	n = copy(buf, *r)
+	*r = (*r)[n:]
+	return n, nil
+}
+
+func TestDecodeWithNonByteReader(t *testing.T) {
+	runTests(t, func(input []byte, into interface{}) error {
+		r := dumbReader(input)
+		return Decode(&r, into)
+	})
+}
+
+func TestDecodeStreamReset(t *testing.T) {
+	s := NewStream(nil)
+	runTests(t, func(input []byte, into interface{}) error {
+		s.Reset(bytes.NewReader(input))
+		return s.Decode(into)
+	})
+}
+
 type testDecoder struct{ called bool }
 
 func (t *testDecoder) DecodeRLP(s *Stream) error {