Merge pull request #169 from nsd20463/custom_map_key_type

Custom marshaler map key type support

Makefile

@@ -24,6 +24,7 @@ generate: root build
 		.root/src/$(PKG)/tests/omitempty.go \
 		.root/src/$(PKG)/tests/nothing.go \
 		.root/src/$(PKG)/tests/named_type.go \
+		.root/src/$(PKG)/tests/custom_map_key_type.go \
 		.root/src/$(PKG)/tests/embedded_type.go
 
 	.root/bin/easyjson -all .root/src/$(PKG)/tests/data.go 
@@ -34,6 +35,7 @@ generate: root build
 	.root/bin/easyjson -build_tags=use_easyjson .root/src/$(PKG)/benchmark/data.go
 	.root/bin/easyjson .root/src/$(PKG)/tests/nested_easy.go
 	.root/bin/easyjson .root/src/$(PKG)/tests/named_type.go
+	.root/bin/easyjson .root/src/$(PKG)/tests/custom_map_key_type.go
 	.root/bin/easyjson .root/src/$(PKG)/tests/embedded_type.go
 
 test: generate root

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 {

tests/basic_test.go

@@ -38,6 +38,7 @@ var testCases = []struct {
 	{&IntsValue, IntsString},
 	{&mapStringStringValue, mapStringStringString},
 	{&namedTypeValue, namedTypeValueString},
+	{&customMapKeyTypeValue, customMapKeyTypeValueString},
 	{&embeddedTypeValue, embeddedTypeValueString},
 	{&mapMyIntStringValue, mapMyIntStringValueString},
 	{&mapIntStringValue, mapIntStringValueString},

tests/custom_map_key_type.go

@@ -0,0 +1,30 @@
+package tests
+
+import fmt "fmt"
+
+//easyjson:json
+type CustomMapKeyType struct {
+	Map map[customKeyType]int
+}
+
+type customKeyType [2]byte
+
+func (k customKeyType) MarshalJSON() ([]byte, error) {
+	return []byte(fmt.Sprintf(`"%02x"`, k)), nil
+}
+
+func (k *customKeyType) UnmarshalJSON(b []byte) error {
+	_, err := fmt.Sscanf(string(b), `"%02x%02x"`, &k[0], &k[1])
+	return err
+}
+
+var customMapKeyTypeValue CustomMapKeyType
+
+func init() {
+	customMapKeyTypeValue.Map = map[customKeyType]int{
+		customKeyType{0x01, 0x01}: 1,
+		customKeyType{0x02, 0x02}: 2,
+	}
+}
+
+var customMapKeyTypeValueString = `{"Map":{"0101":1,"0202":2}}`