From 728039f679cbcd4f6a54e080d2219a4c4928c546 Mon Sep 17 00:00:00 2001
From: Gregory Oschwald <goschwald@maxmind.com>
Date: Mon, 29 Apr 2019 20:50:10 -0700
Subject: [PATCH] Handle other key types in Unmarshal (#276)

Previously, this would fail with:

```
panic: reflect.Value.SetMapIndex: value of type string is not assignable to type toml.letter [recovered]
panic: reflect.Value.SetMapIndex: value of type string is not assignable to type toml.letter
```

Now this only panics when the key type cannot be converted from a
string.
---
 marshal.go      |  2 +-
 marshal_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/marshal.go b/marshal.go
index 2e16a38..0e1c57e 100644
--- a/marshal.go
+++ b/marshal.go
@@ -615,7 +615,7 @@ func (d *Decoder) valueFromTree(mtype reflect.Type, tval *Tree) (reflect.Value,
 			if err != nil {
 				return mval, formatError(err, tval.GetPosition(key))
 			}
-			mval.SetMapIndex(reflect.ValueOf(key), mvalf)
+			mval.SetMapIndex(reflect.ValueOf(key).Convert(mtype.Key()), mvalf)
 		}
 	}
 	return mval, nil
diff --git a/marshal_test.go b/marshal_test.go
index 999235f..cc23dca 100644
--- a/marshal_test.go
+++ b/marshal_test.go
@@ -1126,6 +1126,32 @@ func TestUnmarshalMap(t *testing.T) {
 	}
 }
 
+func TestUnmarshalMapWithTypedKey(t *testing.T) {
+	testToml := []byte(`
+		a = 1
+		b = 2
+		c = 3
+		`)
+
+	type letter string
+	var result map[letter]int
+	err := Unmarshal(testToml, &result)
+	if err != nil {
+		t.Errorf("Received unexpected error: %s", err)
+		return
+	}
+
+	expected := map[letter]int{
+		"a": 1,
+		"b": 2,
+		"c": 3,
+	}
+
+	if !reflect.DeepEqual(result, expected) {
+		t.Errorf("Bad unmarshal: expected %v, got %v", expected, result)
+	}
+}
+
 func TestUnmarshalNonPointer(t *testing.T) {
 	a := 1
 	err := Unmarshal([]byte{}, a)