Repository: goccy/go-yaml
Pull request #8 (2019-10-24 15:06:02)

Add Marshaler / Unmarshaler interface compatible with encoding/json
Supported both interfaces

## Marshaler

- `MarshalYAML() ([]byte, error)`  
- `MarshalYAML() (interface{}, error)`

## Unmarshaler

- `UnmarshalYAML(func(interface{})error) error`
- `UnmarshalYAML([]byte) error`



Patch:
diff --git a/ast/ast.go b/ast/ast.go
index 11387e8..611feee 100644
--- a/ast/ast.go
+++ b/ast/ast.go
@@ -55,6 +55,53 @@ const (
 	TagType
 )
 
+// String node type identifier to text
+func (t NodeType) String() string {
+	switch t {
+	case UnknownNodeType:
+		return "UnknownNode"
+	case DocumentType:
+		return "Document"
+	case NullType:
+		return "Null"
+	case BoolType:
+		return "Bool"
+	case IntegerType:
+		return "Integer"
+	case FloatType:
+		return "Float"
+	case InfinityType:
+		return "Infinity"
+	case NanType:
+		return "Nan"
+	case StringType:
+		return "String"
+	case MergeKeyType:
+		return "MergeKey"
+	case LiteralType:
+		return "Literal"
+	case FlowMappingType:
+		return "FlowMapping"
+	case MappingCollectionType:
+		return "MappingCollection"
+	case MappingValueType:
+		return "MappingValue"
+	case FlowSequenceType:
+		return "FlowSequence"
+	case SequenceType:
+		return "Sequence"
+	case AnchorType:
+		return "Anchor"
+	case AliasType:
+		return "Alias"
+	case DirectiveType:
+		return "Directive"
+	case TagType:
+		return "Tag"
+	}
+	return ""
+}
+
 // Node type of node
 type Node interface {
 	// String node to text
@@ -386,6 +433,39 @@ func (n *NanNode) String() string {
 	return n.Token.Value
 }
 
+// MapNode interface of FlowMappingNode / MappingValueNode / MappingCollectionNode
+type MapNode interface {
+	MapRange() *MapNodeIter
+}
+
+// MapNodeIter is an iterator for ranging over a MapNode
+type MapNodeIter struct {
+	values []*MappingValueNode
+	idx    int
+}
+
+const (
+	startRangeIndex = -1
+)
+
+// Next advances the map iterator and reports whether there is another entry.
+// It returns false when the iterator is exhausted.
+func (m *MapNodeIter) Next() bool {
+	m.idx++
+	next := m.idx < len(m.values)
+	return next
+}
+
+// Key returns the key of the iterator's current map node entry.
+func (m *MapNodeIter) Key() Node {
+	return m.values[m.idx].Key
+}
+
+// Value returns the value of the iterator's current map node entry.
+func (m *MapNodeIter) Value() Node {
+	return m.values[m.idx].Value
+}
+
 // FlowMappingNode type of flow mapping node
 type FlowMappingNode struct {
 	Start  *token.Token
@@ -410,6 +490,14 @@ func (n *FlowMappingNode) String() string {
 	return fmt.Sprintf("{%s}", strings.Join(values, ", "))
 }
 
+// MapRange implements MapNode protocol
+func (n *FlowMappingNode) MapRange() *MapNodeIter {
+	return &MapNodeIter{
+		idx:    startRangeIndex,
+		values: n.Values,
+	}
+}
+
 // MappingCollectionNode type of mapping collection node
 type MappingCollectionNode struct {
 	Start  *token.Token
@@ -433,6 +521,28 @@ func (n *MappingCollectionNode) String() string {
 	return strings.Join(values, "\n")
 }
 
+func (n *MappingCollectionNode) toMappingValues() []*MappingValueNode {
+	values := []*MappingValueNode{}
+	for _, value := range n.Values {
+		if mvnode, ok := value.(*MappingValueNode); ok {
+			values = append(values, mvnode)
+		} else if c, ok := value.(*MappingCollectionNode); ok {
+			values = append(values, c.toMappingValues()...)
+		} else if fmnode, ok := value.(*FlowMappingNode); ok {
+			values = append(values, fmnode.Values...)
+		}
+	}
+	return values
+}
+
+// MapRange implements MapNode protocol
+func (n *MappingCollectionNode) MapRange() *MapNodeIter {
+	return &MapNodeIter{
+		idx:    startRangeIndex,
+		values: n.toMappingValues(),
+	}
+}
+
 // MappingValueNode type of mapping value
 type MappingValueNode struct {
 	Start *token.Token
@@ -467,6 +577,43 @@ func (n *MappingValueNode) String() string {
 	return fmt.Sprintf("%s%s:\n%s", space, n.Key.String(), n.Value.String())
 }
 
+// MapRange implements MapNode protocol
+func (n *MappingValueNode) MapRange() *MapNodeIter {
+	return &MapNodeIter{
+		idx:    startRangeIndex,
+		values: []*MappingValueNode{n},
+	}
+}
+
+// ArrayNode interface of FlowSequenceNode / SequenceNode
+type ArrayNode interface {
+	ArrayRange() *ArrayNodeIter
+}
+
+// ArrayNodeIter is an iterator for ranging over a ArrayNode
+type ArrayNodeIter struct {
+	values []Node
+	idx    int
+}
+
+// Next advances the array iterator and reports whether there is another entry.
+// It returns false when the iterator is exhausted.
+func (m *ArrayNodeIter) Next() bool {
+	m.idx++
+	next := m.idx < len(m.values)
+	return next
+}
+
+// Value returns the value of the iterator's current array entry.
+func (m *ArrayNodeIter) Value() Node {
+	return m.values[m.idx]
+}
+
+// Len returns length of array
+func (m *ArrayNodeIter) Len() int {
+	return len(m.values)
+}
+
 // FlowSequenceNode type of sequence node
 type FlowSequenceNode struct {
 	Start  *token.Token
@@ -491,6 +638,14 @@ func (n *FlowSequenceNode) String() string {
 	return fmt.Sprintf("[%s]", strings.Join(values, ", "))
 }
 
+// ArrayRange implements ArrayNode protocol
+func (n *FlowSequenceNode) ArrayRange() *ArrayNodeIter {
+	return &ArrayNodeIter{
+		idx:    startRangeIndex,
+		values: n.Values,
+	}
+}
+
 // SequenceNode type of sequence node
 type SequenceNode struct {
 	Start  *token.Token
@@ -525,6 +680,14 @@ func (n *SequenceNode) String() string {
 	return strings.Join(values, "\n")
 }
 
+// ArrayRange implements ArrayNode protocol
+func (n *SequenceNode) ArrayRange() *ArrayNodeIter {
+	return &ArrayNodeIter{
+		idx:    startRangeIndex,
+		values: n.Values,
+	}
+}
+
 // AnchorNode type of anchor node
 type AnchorNode struct {
 	Start *token.Token
diff --git a/decode.go b/decode.go
index a4c5ff2..51eef5a 100644
--- a/decode.go
+++ b/decode.go
@@ -14,13 +14,14 @@ import (
 	"github.com/goccy/go-yaml/lexer"
 	"github.com/goccy/go-yaml/parser"
 	"github.com/goccy/go-yaml/token"
+	"golang.org/x/xerrors"
 )
 
 // Decoder reads and decodes YAML values from an input stream.
 type Decoder struct {
 	reader              io.Reader
 	referenceReaders    []io.Reader
-	anchorMap           map[string]interface{}
+	anchorMap           map[string]ast.Node
 	opts                []DecodeOption
 	referenceFiles      []string
 	referenceDirs       []string
@@ -32,7 +33,7 @@ type Decoder struct {
 func NewDecoder(r io.Reader, opts ...DecodeOption) *Decoder {
 	return &Decoder{
 		reader:              r,
-		anchorMap:           map[string]interface{}{},
+		anchorMap:           map[string]ast.Node{},
 		opts:                opts,
 		referenceReaders:    []io.Reader{},
 		referenceFiles:      []string{},
@@ -67,11 +68,11 @@ func (d *Decoder) nodeToValue(node ast.Node) interface{} {
 	case *ast.AnchorNode:
 		anchorName := n.Name.GetToken().Value
 		anchorValue := d.nodeToValue(n.Value)
-		d.anchorMap[anchorName] = anchorValue
+		d.anchorMap[anchorName] = n.Value
 		return anchorValue
 	case *ast.AliasNode:
 		aliasName := n.Value.GetToken().Value
-		return d.anchorMap[aliasName]
+		return d.nodeToValue(d.anchorMap[aliasName])
 	case *ast.LiteralNode:
 		return n.Value.GetValue()
 	case *ast.FlowMappingNode:
@@ -118,18 +119,72 @@ func (d *Decoder) nodeToValue(node ast.Node) interface{} {
 	return nil
 }
 
-func (d *Decoder) docToValue(doc *ast.Document) interface{} {
+func (d *Decoder) getMapNode(node ast.Node) (ast.MapNode, error) {
+	if anchor, ok := node.(*ast.AnchorNode); ok {
+		mapNode, ok := anchor.Value.(ast.MapNode)
+		if ok {
+			return mapNode, nil
+		}
+		return nil, xerrors.Errorf("%s node doesn't MapNode", anchor.Value.Type())
+	}
+	if alias, ok := node.(*ast.AliasNode); ok {
+		aliasName := alias.Value.GetToken().Value
+		anchorNode := d.anchorMap[aliasName]
+		mapNode, ok := anchorNode.(ast.MapNode)
+		if ok {
+			return mapNode, nil
+		}
+		return nil, xerrors.Errorf("%s node doesn't MapNode", anchorNode.Type())
+	}
+	mapNode, ok := node.(ast.MapNode)
+	if !ok {
+		return nil, xerrors.Errorf("%s node doesn't MapNode", node.Type())
+	}
+	return mapNode, nil
+}
+
+func (d *Decoder) getArrayNode(node ast.Node) (ast.ArrayNode, error) {
+	if anchor, ok := node.(*ast.AnchorNode); ok {
+		arrayNode, ok := anchor.Value.(ast.ArrayNode)
+		if ok {
+			return arrayNode, nil
+		}
+		return nil, xerrors.Errorf("%s node doesn't ArrayNode", anchor.Value.Type())
+	}
+	if alias, ok := node.(*ast.AliasNode); ok {
+		aliasName := alias.Value.GetToken().Value
+		anchorNode := d.anchorMap[aliasName]
+		arrayNode, ok := anchorNode.(ast.ArrayNode)
+		if ok {
+			return arrayNode, nil
+		}
+		return nil, xerrors.Errorf("%s node doesn't ArrayNode", anchorNode.Type())
+	}
+	arrayNode, ok := node.(ast.ArrayNode)
+	if !ok {
+		return nil, xerrors.Errorf("%s node doesn't ArrayNode", node.Type())
+	}
+	return arrayNode, nil
+}
+
+func (d *Decoder) docToNode(doc *ast.Document) ast.Node {
 	for _, node := range doc.Nodes {
 		if v := d.nodeToValue(node); v != nil {
-			return v
+			return node
 		}
 	}
 	return nil
 }
 
-func (d *Decoder) decodeValue(dst reflect.Value, src interface{}) error {
+func (d *Decoder) decodeValue(dst reflect.Value, src ast.Node) error {
 	valueType := dst.Type()
-	if unmarshaler, ok := dst.Addr().Interface().(Unmarshaler); ok {
+	if unmarshaler, ok := dst.Addr().Interface().(BytesUnmarshaler); ok {
+		b := fmt.Sprintf("%v", src)
+		if err := unmarshaler.UnmarshalYAML([]byte(b)); err != nil {
+			return errors.Wrapf(err, "failed to UnmarshalYAML")
+		}
+		return nil
+	} else if unmarshaler, ok := dst.Addr().Interface().(InterfaceUnmarshaler); ok {
 		if err := unmarshaler.UnmarshalYAML(func(v interface{}) error {
 			rv := reflect.ValueOf(v)
 			if rv.Type().Kind() != reflect.Ptr {
@@ -155,7 +210,7 @@ func (d *Decoder) decodeValue(dst reflect.Value, src interface{}) error {
 		}
 		dst.Set(d.castToAssignableValue(v, dst.Type()))
 	case reflect.Interface:
-		v := reflect.ValueOf(src)
+		v := reflect.ValueOf(d.nodeToValue(src))
 		if v.IsValid() {
 			dst.Set(v)
 		}
@@ -166,7 +221,7 @@ func (d *Decoder) decodeValue(dst reflect.Value, src interface{}) error {
 	case reflect.Struct:
 		return d.decodeStruct(dst, src)
 	}
-	v := reflect.ValueOf(src)
+	v := reflect.ValueOf(d.nodeToValue(src))
 	if v.IsValid() {
 		dst.Set(v.Convert(dst.Type()))
 	}
@@ -203,7 +258,35 @@ func (d *Decoder) castToAssignableValue(value reflect.Value, target reflect.Type
 	return value
 }
 
-func (d *Decoder) decodeStruct(dst reflect.Value, src interface{}) error {
+func (d *Decoder) keyToNodeMap(node ast.Node) (map[string]ast.Node, error) {
+	mapNode, err := d.getMapNode(node)
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to get map node")
+	}
+	mapIter := mapNode.MapRange()
+	keyToNodeMap := map[string]ast.Node{}
+	for mapIter.Next() {
+		keyNode := mapIter.Key()
+		if keyNode.Type() == ast.MergeKeyType {
+			mergeMap, err := d.keyToNodeMap(mapIter.Value())
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to get keyToNodeMap by MergeKey node")
+			}
+			for k, v := range mergeMap {
+				keyToNodeMap[k] = v
+			}
+		} else {
+			key, ok := d.nodeToValue(keyNode).(string)
+			if !ok {
+				return nil, errors.Wrapf(err, "failed to decode map key")
+			}
+			keyToNodeMap[key] = mapIter.Value()
+		}
+	}
+	return keyToNodeMap, nil
+}
+
+func (d *Decoder) decodeStruct(dst reflect.Value, src ast.Node) error {
 	if src == nil {
 		return nil
 	}
@@ -213,9 +296,9 @@ func (d *Decoder) decodeStruct(dst reflect.Value, src interface{}) error {
 	if err != nil {
 		return errors.Wrapf(err, "failed to create struct field map")
 	}
-	srcMap, ok := src.(map[string]interface{})
-	if !ok {
-		return errors.Wrapf(err, "value is not struct type: %s", reflect.TypeOf(src).Name())
+	keyToNodeMap, err := d.keyToNodeMap(src)
+	if err != nil {
+		return errors.Wrapf(err, "failed to get keyToNodeMap")
 	}
 	for i := 0; i < structType.NumField(); i++ {
 		field := structType.Field(i)
@@ -223,7 +306,7 @@ func (d *Decoder) decodeStruct(dst reflect.Value, src interface{}) error {
 			continue
 		}
 		structField := structFieldMap[field.Name]
-		v, exists := srcMap[structField.RenderName]
+		v, exists := keyToNodeMap[structField.RenderName]
 		if !exists {
 			continue
 		}
@@ -238,12 +321,17 @@ func (d *Decoder) decodeStruct(dst reflect.Value, src interface{}) error {
 	return nil
 }
 
-func (d *Decoder) decodeSlice(dst reflect.Value, src interface{}) error {
+func (d *Decoder) decodeSlice(dst reflect.Value, src ast.Node) error {
+	arrayNode, err := d.getArrayNode(src)
+	if err != nil {
+		return errors.Wrapf(err, "failed to get array node")
+	}
+	iter := arrayNode.ArrayRange()
 	sliceType := dst.Type()
-	slice := src.([]interface{})
-	sliceValue := reflect.MakeSlice(sliceType, 0, len(slice))
+	sliceValue := reflect.MakeSlice(sliceType, 0, iter.Len())
 	elemType := sliceType.Elem()
-	for _, v := range slice {
+	for iter.Next() {
+		v := iter.Value()
 		dstValue := d.createDecodableValue(elemType)
 		if err := d.decodeValue(dstValue, v); err != nil {
 			return errors.Wrapf(err, "failed to decode value")
@@ -254,17 +342,24 @@ func (d *Decoder) decodeSlice(dst reflect.Value, src interface{}) error {
 	return nil
 }
 
-func (d *Decoder) decodeMap(dst reflect.Value, src interface{}) error {
+func (d *Decoder) decodeMap(dst reflect.Value, src ast.Node) error {
+	mapNode, err := d.getMapNode(src)
+	if err != nil {
+		return errors.Wrapf(err, "failed to get map node")
+	}
 	mapType := dst.Type()
 	mapValue := reflect.MakeMap(mapType)
 	keyType := mapValue.Type().Key()
 	valueType := mapValue.Type().Elem()
-	for k, v := range src.(map[string]interface{}) {
+	mapIter := mapNode.MapRange()
+	for mapIter.Next() {
+		key := mapIter.Key()
+		value := mapIter.Value()
 		dstValue := d.createDecodableValue(valueType)
-		if err := d.decodeValue(dstValue, v); err != nil {
+		if err := d.decodeValue(dstValue, value); err != nil {
 			return errors.Wrapf(err, "failed to decode value")
 		}
-		castedKey := reflect.ValueOf(k).Convert(keyType)
+		castedKey := reflect.ValueOf(d.nodeToValue(key)).Convert(keyType)
 		mapValue.SetMapIndex(castedKey, d.castToAssignableValue(dstValue, valueType))
 	}
 	dst.Set(mapValue)
@@ -371,7 +466,7 @@ func (d *Decoder) resolveReference() error {
 	return nil
 }
 
-func (d *Decoder) decode(bytes []byte) (interface{}, error) {
+func (d *Decoder) decode(bytes []byte) (ast.Node, error) {
 	var (
 		parser parser.Parser
 	)
@@ -380,7 +475,7 @@ func (d *Decoder) decode(bytes []byte) (interface{}, error) {
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to parse yaml")
 	}
-	return d.docToValue(doc), nil
+	return d.docToNode(doc), nil
 }
 
 // Decode reads the next YAML-encoded value from its input
@@ -402,14 +497,14 @@ func (d *Decoder) Decode(v interface{}) error {
 	if err != nil {
 		return errors.Wrapf(err, "failed to read buffer")
 	}
-	value, err := d.decode(bytes)
+	node, err := d.decode(bytes)
 	if err != nil {
 		return errors.Wrapf(err, "failed to decode")
 	}
-	if value == nil {
+	if node == nil {
 		return nil
 	}
-	if err := d.decodeValue(rv.Elem(), value); err != nil {
+	if err := d.decodeValue(rv.Elem(), node); err != nil {
 		return errors.Wrapf(err, "failed to decode value")
 	}
 	return nil
diff --git a/encode.go b/encode.go
index 53493c4..5a35086 100644
--- a/encode.go
+++ b/encode.go
@@ -11,6 +11,8 @@ import (
 
 	"github.com/goccy/go-yaml/ast"
 	"github.com/goccy/go-yaml/internal/errors"
+	"github.com/goccy/go-yaml/lexer"
+	"github.com/goccy/go-yaml/parser"
 	"github.com/goccy/go-yaml/printer"
 	"github.com/goccy/go-yaml/token"
 	"golang.org/x/xerrors"
@@ -71,8 +73,35 @@ func (e *Encoder) Encode(v interface{}) error {
 	return nil
 }
 
+func (e *Encoder) encodeDocument(doc []byte) (ast.Node, error) {
+	var (
+		parser parser.Parser
+	)
+	tokens := lexer.Tokenize(string(doc))
+	docNode, err := parser.Parse(tokens)
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to parse yaml")
+	}
+	for _, node := range docNode.Nodes {
+		if node != nil {
+			return node, nil
+		}
+	}
+	return nil, nil
+}
+
 func (e *Encoder) encodeValue(v reflect.Value, column int) (ast.Node, error) {
-	if marshaler, ok := v.Interface().(Marshaler); ok {
+	if marshaler, ok := v.Interface().(BytesMarshaler); ok {
+		doc, err := marshaler.MarshalYAML()
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed to MarshalYAML")
+		}
+		node, err := e.encodeDocument(doc)
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed to encode document")
+		}
+		return node, nil
+	} else if marshaler, ok := v.Interface().(InterfaceMarshaler); ok {
 		marshalV, err := marshaler.MarshalYAML()
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to MarshalYAML")
@@ -99,7 +128,7 @@ func (e *Encoder) encodeValue(v reflect.Value, column int) (ast.Node, error) {
 		if mapSlice, ok := v.Interface().(MapSlice); ok {
 			return e.encodeMapSlice(mapSlice, column)
 		}
-		return e.encodeSlice(v), nil
+		return e.encodeSlice(v)
 	case reflect.Struct:
 		if mapItem, ok := v.Interface().(MapItem); ok {
 			return e.encodeMapItem(mapItem, column)
@@ -177,7 +206,7 @@ func (e *Encoder) encodeBool(v bool) ast.Node {
 	return ast.Bool(token.New(value, value, e.pos(e.column)))
 }
 
-func (e *Encoder) encodeSlice(value reflect.Value) ast.Node {
+func (e *Encoder) encodeSlice(value reflect.Value) (ast.Node, error) {
 	sequence := &ast.SequenceNode{
 		Start:  token.New("-", "-", e.pos(e.column)),
 		Values: []ast.Node{},
@@ -185,12 +214,11 @@ func (e *Encoder) encodeSlice(value reflect.Value) ast.Node {
 	for i := 0; i < value.Len(); i++ {
 		node, err := e.encodeValue(value.Index(i), e.column)
 		if err != nil {
-			panic(err)
-			return nil
+			return nil, errors.Wrapf(err, "failed to encode value for slice")
 		}
 		sequence.Values = append(sequence.Values, node)
 	}
-	return sequence
+	return sequence, nil
 }
 
 func (e *Encoder) encodeMapItem(item MapItem, column int) (*ast.MappingValueNode, error) {
@@ -337,6 +365,7 @@ func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error
 			for _, value := range c.Values {
 				if mvnode, ok := value.(*ast.MappingValueNode); ok {
 					mvnode.Key.GetToken().Position.Column += e.indent
+					mvnode.Value.GetToken().Position.Column += e.indent
 				}
 			}
 		}
diff --git a/scanner/scanner.go b/scanner/scanner.go
index 2454ad6..429fd5d 100644
--- a/scanner/scanner.go
+++ b/scanner/scanner.go
@@ -286,8 +286,12 @@ func (s *Scanner) scan(ctx *Context) (pos int) {
 		} else if ctx.isLiteral || ctx.isFolded || ctx.isRawFolded {
 			s.scanLiteral(ctx, c)
 			continue
-		} else if s.isMapContext && s.isChangedToIndentStateEqual() {
-			s.addBufferedTokenIfExists(ctx)
+		} else if s.isChangedToIndentStateEqual() {
+			// if first character is \n, buffer expect to raw folded literal
+			if len(ctx.obuf) > 0 && ctx.obuf[0] != '\n' {
+				// doesn't raw folded literal
+				s.addBufferedTokenIfExists(ctx)
+			}
 			s.isMapContext = false
 		} else if c != '-' && s.indentState == IndentStateUp {
 			s.isMapContext = false
diff --git a/yaml.go b/yaml.go
index a58a9fa..b45d920 100644
--- a/yaml.go
+++ b/yaml.go
@@ -7,19 +7,29 @@ import (
 	"golang.org/x/xerrors"
 )
 
-// Marshaler interface may be implemented by types to customize their
+// BytesMarshaler interface may be implemented by types to customize their
 // behavior when being marshaled into a YAML document. The returned value
 // is marshaled in place of the original value implementing Marshaler.
 //
 // If an error is returned by MarshalYAML, the marshaling procedure stops
 // and returns with the provided error.
-type Marshaler interface {
+type BytesMarshaler interface {
+	MarshalYAML() ([]byte, error)
+}
+
+// InterfaceMarshaler interface has MarshalYAML compatible with github.com/go-yaml/yaml package.
+type InterfaceMarshaler interface {
 	MarshalYAML() (interface{}, error)
 }
 
-// Unmarshaler interface may be implemented by types to customize their
+// BytesUnmarshaler interface may be implemented by types to customize their
 // behavior when being unmarshaled from a YAML document.
-type Unmarshaler interface {
+type BytesUnmarshaler interface {
+	UnmarshalYAML([]byte) error
+}
+
+// InterfaceUnmarshaler interface has UnmarshalYAML compatible with github.com/go-yaml/yaml package.
+type InterfaceUnmarshaler interface {
 	UnmarshalYAML(func(interface{}) error) error
 }
 
@@ -86,7 +96,7 @@ func Marshal(v interface{}) ([]byte, error) {
 	var buf bytes.Buffer
 	enc := NewEncoder(&buf)
 	if err := enc.Encode(v); err != nil {
-		return nil, errors.Wrapf(err, "failed to marshal", err)
+		return nil, errors.Wrapf(err, "failed to marshal")
 	}
 	return buf.Bytes(), nil
 }


Test patch:
diff --git a/yaml_test.go b/yaml_test.go
index c5449be..53b68ef 100644
--- a/yaml_test.go
+++ b/yaml_test.go
@@ -41,8 +41,8 @@ b: c
 
 type marshalTest struct{}
 
-func (t *marshalTest) MarshalYAML() (interface{}, error) {
-	return yaml.MapSlice{
+func (t *marshalTest) MarshalYAML() ([]byte, error) {
+	return yaml.Marshal(yaml.MapSlice{
 		{
 			"a", 1,
 		},
@@ -52,14 +52,32 @@ func (t *marshalTest) MarshalYAML() (interface{}, error) {
 		{
 			"c", true,
 		},
+	})
+}
+
+type marshalTest2 struct{}
+
+func (t *marshalTest2) MarshalYAML() (interface{}, error) {
+	return yaml.MapSlice{
+		{
+			"a", 2,
+		},
+		{
+			"b", "world",
+		},
+		{
+			"c", true,
+		},
 	}, nil
 }
 
 func TestMarshalYAML(t *testing.T) {
 	var v struct {
 		A *marshalTest
+		B *marshalTest2
 	}
 	v.A = &marshalTest{}
+	v.B = &marshalTest2{}
 	bytes, err := yaml.Marshal(v)
 	if err != nil {
 		t.Fatalf("failed to Marshal: %+v", err)
@@ -69,6 +87,10 @@ a:
   a: 1
   b: hello
   c: true
+b:
+  a: 2
+  b: world
+  c: true
 `
 	actual := "\n" + string(bytes)
 	if expect != actual {
@@ -82,7 +104,37 @@ type unmarshalTest struct {
 	c bool
 }
 
-func (t *unmarshalTest) UnmarshalYAML(unmarshal func(interface{}) error) error {
+func (t *unmarshalTest) UnmarshalYAML(b []byte) error {
+	if t.a != 0 {
+		return xerrors.New("unexpected field value to a")
+	}
+	if t.b != "" {
+		return xerrors.New("unexpected field value to b")
+	}
+	if t.c {
+		return xerrors.New("unexpected field value to c")
+	}
+	var v struct {
+		A int
+		B string
+		C bool
+	}
+	if err := yaml.Unmarshal(b, &v); err != nil {
+		return err
+	}
+	t.a = v.A
+	t.b = v.B
+	t.c = v.C
+	return nil
+}
+
+type unmarshalTest2 struct {
+	a int
+	b string
+	c bool
+}
+
+func (t *unmarshalTest2) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	var v struct {
 		A int
 		B string
@@ -112,9 +164,14 @@ a:
   a: 1
   b: hello
   c: true
+b:
+  a: 2
+  b: world
+  c: true
 `
 	var v struct {
 		A *unmarshalTest
+		B *unmarshalTest2
 	}
 	if err := yaml.Unmarshal([]byte(yml), &v); err != nil {
 		t.Fatalf("failed to Unmarshal: %+v", err)
@@ -131,4 +188,16 @@ a:
 	if !v.A.c {
 		t.Fatal("failed to UnmarshalYAML")
 	}
+	if v.B == nil {
+		t.Fatal("failed to UnmarshalYAML")
+	}
+	if v.B.a != 2 {
+		t.Fatal("failed to UnmarshalYAML")
+	}
+	if v.B.b != "world" {
+		t.Fatal("failed to UnmarshalYAML")
+	}
+	if !v.B.c {
+		t.Fatal("failed to UnmarshalYAML")
+	}
 }