diff --git a/rlp/encode.go b/rlp/encode.go index d80b66315..6ae4a123a 100644 --- a/rlp/encode.go +++ b/rlp/encode.go @@ -32,6 +32,48 @@ type Encoder interface { EncodeRLP(io.Writer) error } +// Flat wraps a value (which must encode as a list) so +// it encodes as the list's elements. +// +// Example: suppose you have defined a type +// +// type foo struct { A, B uint } +// +// Under normal encoding rules, +// +// rlp.Encode(foo{1, 2}) --> 0xC20102 +// +// This function can help you achieve the following encoding: +// +// rlp.Encode(rlp.Flat(foo{1, 2})) --> 0x0102 +func Flat(val interface{}) Encoder { + return flatenc{val} +} + +type flatenc struct{ val interface{} } + +func (e flatenc) EncodeRLP(out io.Writer) error { + // record current output position + var ( + eb = out.(*encbuf) + prevstrsize = len(eb.str) + prevnheads = len(eb.lheads) + ) + if err := eb.encode(e.val); err != nil { + return err + } + // check that a new list header has appeared + if len(eb.lheads) == prevnheads || eb.lheads[prevnheads].offset == prevstrsize-1 { + return fmt.Errorf("rlp.Flat: %T did not encode as list", e.val) + } + // remove the new list header + newhead := eb.lheads[prevnheads] + copy(eb.lheads[prevnheads:], eb.lheads[prevnheads+1:]) + eb.lheads = eb.lheads[:len(eb.lheads)-1] + eb.lhsize -= newhead.tagsize() + return nil +} + // Encode writes the RLP encoding of val to w. Note that Encode may // perform many small writes in some cases. Consider making w // buffered. @@ -123,6 +165,13 @@ func (head *listhead) encode(buf []byte) []byte { } } +func (head *listhead) tagsize() int { + if head.size < 56 { + return 1 + } + return 1 + intsize(uint64(head.size)) +} + func newencbuf() *encbuf { return &encbuf{sizebuf: make([]byte, 9)} } diff --git a/rlp/encode_test.go b/rlp/encode_test.go index 18b843737..9b3085658 100644 --- a/rlp/encode_test.go +++ b/rlp/encode_test.go @@ -177,6 +177,15 @@ var encTests = []encTest{ {val: &recstruct{5, nil}, output: "C205C0"}, {val: &recstruct{5, &recstruct{4, &recstruct{3, nil}}}, output: "C605C404C203C0"}, + // flat + {val: Flat(uint(1)), error: "rlp.Flat: uint did not encode as list"}, + {val: Flat(simplestruct{A: 3, B: "foo"}), output: "0383666F6F"}, + { + // value generates more list headers after the Flat + val: []interface{}{"foo", []uint{1, 2}, Flat([]uint{3, 4}), []uint{5, 6}, "bar"}, + output: "D083666F6FC201020304C2050683626172", + }, + // nil {val: (*uint)(nil), output: "80"}, {val: (*string)(nil), output: "80"},