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

@@ -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))
 

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 {