From fe875b37056b8971629f1f5f9aff45df314a5e65 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Fri, 29 May 2020 19:47:47 +0900
Subject: [PATCH] Add UseOrderedMap option

---
 decode.go      | 36 ++++++++++++++++++++++++++++++++++++
 decode_test.go | 25 +++++++++++++++++++++++++
 option.go      |  9 +++++++++
 3 files changed, 70 insertions(+)

diff --git a/decode.go b/decode.go
index 4a10e89d..c01d404c 100644
--- a/decode.go
+++ b/decode.go
@@ -34,6 +34,7 @@ type Decoder struct {
 	isResolvedReference  bool
 	validator            StructValidator
 	disallowUnknownField bool
+	useOrderedMap        bool
 }
 
 // NewDecoder returns a new decoder that reads from r.
@@ -49,6 +50,7 @@ func NewDecoder(r io.Reader, opts ...DecodeOption) *Decoder {
 		isRecursiveDir:       false,
 		isResolvedReference:  false,
 		disallowUnknownField: false,
+		useOrderedMap:        false,
 	}
 }
 
@@ -105,6 +107,25 @@ func (d *Decoder) setToMapValue(node ast.Node, m map[string]interface{}) {
 	}
 }
 
+func (d *Decoder) setToOrderedMapValue(node ast.Node, m *MapSlice) {
+	switch n := node.(type) {
+	case *ast.MappingValueNode:
+		if n.Key.Type() == ast.MergeKeyType && n.Value.Type() == ast.AliasType {
+			aliasNode := n.Value.(*ast.AliasNode)
+			aliasName := aliasNode.Value.GetToken().Value
+			node := d.anchorNodeMap[aliasName]
+			d.setToOrderedMapValue(node, m)
+		} else {
+			key := n.Key.GetToken().Value
+			*m = append(*m, MapItem{Key: key, Value: d.nodeToValue(n.Value)})
+		}
+	case *ast.MappingNode:
+		for _, value := range n.Values {
+			d.setToOrderedMapValue(value, m)
+		}
+	}
+}
+
 func (d *Decoder) nodeToValue(node ast.Node) interface{} {
 	switch n := node.(type) {
 	case *ast.NullNode:
@@ -150,15 +171,30 @@ func (d *Decoder) nodeToValue(node ast.Node) interface{} {
 			aliasNode := n.Value.(*ast.AliasNode)
 			aliasName := aliasNode.Value.GetToken().Value
 			node := d.anchorNodeMap[aliasName]
+			if d.useOrderedMap {
+				m := MapSlice{}
+				d.setToOrderedMapValue(node, &m)
+				return m
+			}
 			m := map[string]interface{}{}
 			d.setToMapValue(node, m)
 			return m
 		}
 		key := n.Key.GetToken().Value
+		if d.useOrderedMap {
+			return MapSlice{{Key: key, Value: d.nodeToValue(n.Value)}}
+		}
 		return map[string]interface{}{
 			key: d.nodeToValue(n.Value),
 		}
 	case *ast.MappingNode:
+		if d.useOrderedMap {
+			m := make(MapSlice, 0, len(n.Values))
+			for _, value := range n.Values {
+				d.setToOrderedMapValue(value, &m)
+			}
+			return m
+		}
 		m := make(map[string]interface{}, len(n.Values))
 		for _, value := range n.Values {
 			d.setToMapValue(value, m)
diff --git a/decode_test.go b/decode_test.go
index 7dea2bc9..7cad20b3 100644
--- a/decode_test.go
+++ b/decode_test.go
@@ -1703,3 +1703,28 @@ value: |
 		t.Fatal("failed to unmarshal literal with bytes unmarshaler")
 	}
 }
+
+func TestDecoder_UseOrderedMap(t *testing.T) {
+	yml := `
+a: b
+c: d
+e:
+  f: g
+  h: i
+j: k
+`
+	var v interface{}
+	if err := yaml.NewDecoder(strings.NewReader(yml), yaml.UseOrderedMap()).Decode(&v); err != nil {
+		t.Fatalf("%+v", err)
+	}
+	if _, ok := v.(yaml.MapSlice); !ok {
+		t.Fatalf("failed to convert to ordered map: %T", v)
+	}
+	bytes, err := yaml.Marshal(v)
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+	if string(yml) != "\n"+string(bytes) {
+		t.Fatalf("expected:[%s] actual:[%s]", string(yml), "\n"+string(bytes))
+	}
+}
diff --git a/option.go b/option.go
index 2a98628a..b81ed116 100644
--- a/option.go
+++ b/option.go
@@ -59,6 +59,15 @@ func DisallowUnknownField() DecodeOption {
 	}
 }
 
+// UseOrderedMap can be interpreted as a map,
+// and uses MapSlice ( ordered map ) aggressively if there is no type specification
+func UseOrderedMap() DecodeOption {
+	return func(d *Decoder) error {
+		d.useOrderedMap = true
+		return nil
+	}
+}
+
 // EncodeOption functional option type for Encoder
 type EncodeOption func(e *Encoder) error