rlp: add Stream.Raw

1 files changed, 19 insertions, 11 deletions

diff --git a/rlp/encode.go b/rlp/encode.go
index 7ac74d8fb..289bc4eaa 100644
--- a/rlp/encode.go
+++ b/rlp/encode.go
@@ -70,7 +70,7 @@ func (e flatenc) EncodeRLP(out io.Writer) error {
 	newhead := eb.lheads[prevnheads]
 	copy(eb.lheads[prevnheads:], eb.lheads[prevnheads+1:])
 	eb.lheads = eb.lheads[:len(eb.lheads)-1]
-	eb.lhsize -= newhead.tagsize()
+	eb.lhsize -= headsize(uint64(newhead.size))
 	return nil
 }
 
@@ -155,21 +155,29 @@ type listhead struct {
 // encode writes head to the given buffer, which must be at least
 // 9 bytes long. It returns the encoded bytes.
 func (head *listhead) encode(buf []byte) []byte {
-	if head.size < 56 {
-		buf[0] = 0xC0 + byte(head.size)
-		return buf[:1]
-	} else {
-		sizesize := putint(buf[1:], uint64(head.size))
-		buf[0] = 0xF7 + byte(sizesize)
-		return buf[:sizesize+1]
+	return buf[:puthead(buf, 0xC0, 0xF7, uint64(head.size))]
+}
+
+// headsize returns the size of a list or string header
+// for a value of the given size.
+func headsize(size uint64) int {
+	if size < 56 {
+		return 1
 	}
+	return 1 + intsize(size)
 }
 
-func (head *listhead) tagsize() int {
-	if head.size < 56 {
+// puthead writes a list or string header to buf.
+// buf must be at least 9 bytes long.
+func puthead(buf []byte, smalltag, largetag byte, size uint64) int {
+	if size < 56 {
+		buf[0] = smalltag + byte(size)
 		return 1
+	} else {
+		sizesize := putint(buf[1:], size)
+		buf[0] = largetag + byte(sizesize)
+		return sizesize + 1
 	}
-	return 1 + intsize(uint64(head.size))
 }
 
 func newencbuf() *encbuf {