From a06183da62cb3e02e05eb265809164c91ca7b16b Mon Sep 17 00:00:00 2001
From: "Nicolas S. Dade" <nic.dade@gmail.com>
Date: Tue, 6 Mar 2018 13:29:46 -0800
Subject: [PATCH] support maps with key types which have custom
 marshler/unmarshalers

by assuming the caller knows what they are doing and that the custom
marshler will generate JSON appropriate for a key.

The standard library's encoding/json supports these.
---
 gen/decoder.go | 24 ++++++++++++++++++++----
 gen/encoder.go | 23 +++++++++++++++++++----
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/gen/decoder.go b/gen/decoder.go
index 021933a..f740c68 100644
--- a/gen/decoder.go
+++ b/gen/decoder.go
@@ -84,6 +84,14 @@ func (g *Generator) genTypeDecoder(t reflect.Type, out string, tags fieldTags, i
 	return err
 }
 
+// returns true of the type t implements one of the custom unmarshaler interfaces
+func hasCustomUnmarshaler(t reflect.Type) bool {
+	t = reflect.PtrTo(t)
+	return t.Implements(reflect.TypeOf((*easyjson.Unmarshaler)(nil)).Elem()) ||
+		t.Implements(reflect.TypeOf((*json.Unmarshaler)(nil)).Elem()) ||
+		t.Implements(reflect.TypeOf((*encoding.TextUnmarshaler)(nil)).Elem())
+}
+
 // genTypeDecoderNoCheck generates decoding code for the type t.
 func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags fieldTags, indent int) error {
 	ws := strings.Repeat("  ", indent)
@@ -208,9 +216,9 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 	case reflect.Map:
 		key := t.Key()
 		keyDec, ok := primitiveStringDecoders[key.Kind()]
-		if !ok {
-			return fmt.Errorf("map type %v not supported: only string and integer keys are allowed", key)
-		}
+		if !ok && !hasCustomUnmarshaler(key) {
+			return fmt.Errorf("map type %v not supported: only string and integer keys and types implementing json.Unmarshaler are allowed", key)
+		} // else assume the caller knows what they are doing and that the custom unmarshaler performs the translation from string or integer keys to the key type
 		elem := t.Elem()
 		tmpVar := g.uniqueVarName()
 
@@ -225,7 +233,15 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 		fmt.Fprintln(g.out, ws+"  }")
 
 		fmt.Fprintln(g.out, ws+"  for !in.IsDelim('}') {")
-		fmt.Fprintln(g.out, ws+"    key := "+g.getType(key)+"("+keyDec+")")
+		if keyDec != "" {
+			fmt.Fprintln(g.out, ws+"    key := "+g.getType(key)+"("+keyDec+")")
+		} else {
+			fmt.Fprintln(g.out, ws+"    var key "+g.getType(key))
+			if err := g.genTypeDecoder(key, "key", tags, indent+2); err != nil {
+				return err
+			}
+		}
+
 		fmt.Fprintln(g.out, ws+"    in.WantColon()")
 		fmt.Fprintln(g.out, ws+"    var "+tmpVar+" "+g.getType(elem))
 
diff --git a/gen/encoder.go b/gen/encoder.go
index 48cba15..4eca160 100644
--- a/gen/encoder.go
+++ b/gen/encoder.go
@@ -108,6 +108,14 @@ func (g *Generator) genTypeEncoder(t reflect.Type, in string, tags fieldTags, in
 	return err
 }
 
+// returns true of the type t implements one of the custom marshaler interfaces
+func hasCustomMarshaler(t reflect.Type) bool {
+	t = reflect.PtrTo(t)
+	return t.Implements(reflect.TypeOf((*easyjson.Marshaler)(nil)).Elem()) ||
+		t.Implements(reflect.TypeOf((*json.Marshaler)(nil)).Elem()) ||
+		t.Implements(reflect.TypeOf((*encoding.TextMarshaler)(nil)).Elem())
+}
+
 // genTypeEncoderNoCheck generates code that encodes in of type t into the writer.
 func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldTags, indent int, assumeNonEmpty bool) error {
 	ws := strings.Repeat("  ", indent)
@@ -197,9 +205,9 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 	case reflect.Map:
 		key := t.Key()
 		keyEnc, ok := primitiveStringEncoders[key.Kind()]
-		if !ok {
-			return fmt.Errorf("map key type %v not supported: only string and integer keys are allowed", key)
-		}
+		if !ok && !hasCustomMarshaler(key) {
+			return fmt.Errorf("map key type %v not supported: only string and integer keys and types implementing Marshaler interfaces are allowed", key)
+		} // else assume the caller knows what they are doing and that the custom marshaler performs the translation from the key type to a string or integer
 		tmpVar := g.uniqueVarName()
 
 		if !assumeNonEmpty {
@@ -213,7 +221,14 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 		fmt.Fprintln(g.out, ws+"  "+tmpVar+"First := true")
 		fmt.Fprintln(g.out, ws+"  for "+tmpVar+"Name, "+tmpVar+"Value := range "+in+" {")
 		fmt.Fprintln(g.out, ws+"    if "+tmpVar+"First { "+tmpVar+"First = false } else { out.RawByte(',') }")
-		fmt.Fprintln(g.out, ws+"    "+fmt.Sprintf(keyEnc, tmpVar+"Name"))
+		if keyEnc != "" {
+			fmt.Fprintln(g.out, ws+"    "+fmt.Sprintf(keyEnc, tmpVar+"Name"))
+		} else {
+			if err := g.genTypeEncoder(key, tmpVar+"Name", tags, indent+2, false); err != nil {
+				return err
+			}
+		}
+
 		fmt.Fprintln(g.out, ws+"    out.RawByte(':')")
 
 		if err := g.genTypeEncoder(t.Elem(), tmpVar+"Value", tags, indent+2, false); err != nil {