From f3614c79b68ac531348a1c011ad6890e29efd04d Mon Sep 17 00:00:00 2001
From: Ryo Kitagawa <ryo-kitagawa@m3.com>
Date: Sat, 26 Jun 2021 20:38:11 +0900
Subject: [PATCH 1/2] Support to encode map which has defined type key

---
 encode.go      | 12 +++++++-----
 encode_test.go | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/encode.go b/encode.go
index 116999c5..2e075f32 100644
--- a/encode.go
+++ b/encode.go
@@ -423,11 +423,13 @@ func (e *Encoder) encodeMapSlice(ctx context.Context, value MapSlice, column int
 
 func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int) ast.Node {
 	node := ast.Mapping(token.New("", "", e.pos(column)), e.isFlowStyle)
-	keys := []string{}
-	for _, k := range value.MapKeys() {
-		keys = append(keys, k.Interface().(string))
+	keys := make([]interface{}, len(value.MapKeys()))
+	for i, k := range value.MapKeys() {
+		keys[i] = k.Interface()
 	}
-	sort.Strings(keys)
+	sort.Slice(keys, func(i, j int) bool {
+		return fmt.Sprintf("%s", keys[i]) < fmt.Sprintf("%s", keys[j])
+	})
 	for _, key := range keys {
 		k := reflect.ValueOf(key)
 		v := value.MapIndex(k)
@@ -440,7 +442,7 @@ func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int
 		}
 		node.Values = append(node.Values, ast.MappingValue(
 			nil,
-			e.encodeString(k.Interface().(string), column),
+			e.encodeString(fmt.Sprintf("%s", key), column),
 			value,
 		))
 	}
diff --git a/encode_test.go b/encode_test.go
index 503b95ed..eec8c45b 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -572,6 +572,24 @@ func TestEncodeStructIncludeMap(t *testing.T) {
 	}
 }
 
+func TestEncodeDefinedTypeKeyMap(t *testing.T) {
+	type K string
+	type U struct {
+		M map[K]string
+	}
+	bytes, err := yaml.Marshal(U{
+		M: map[K]string{K("x"): "y"},
+	})
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+	expect := "m:\n  x: y\n"
+	actual := string(bytes)
+	if actual != expect {
+		t.Fatalf("unexpected output. expect:[%s] actual:[%s]", expect, actual)
+	}
+}
+
 func TestEncodeWithAnchorAndAlias(t *testing.T) {
 	var buf bytes.Buffer
 	enc := yaml.NewEncoder(&buf)

From bae5ef26e03089fcce6b807ff3f643a8472cee8b Mon Sep 17 00:00:00 2001
From: Ryo Kitagawa <ryo-kitagawa@m3.com>
Date: Wed, 30 Jun 2021 14:32:52 +0900
Subject: [PATCH 2/2] Use fmt.Sprint

---
 encode.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/encode.go b/encode.go
index 2e075f32..4e2445f8 100644
--- a/encode.go
+++ b/encode.go
@@ -428,7 +428,7 @@ func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int
 		keys[i] = k.Interface()
 	}
 	sort.Slice(keys, func(i, j int) bool {
-		return fmt.Sprintf("%s", keys[i]) < fmt.Sprintf("%s", keys[j])
+		return fmt.Sprint(keys[i]) < fmt.Sprint(keys[j])
 	})
 	for _, key := range keys {
 		k := reflect.ValueOf(key)
@@ -442,7 +442,7 @@ func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int
 		}
 		node.Values = append(node.Values, ast.MappingValue(
 			nil,
-			e.encodeString(fmt.Sprintf("%s", key), column),
+			e.encodeString(fmt.Sprint(key), column),
 			value,
 		))
 	}