From 63896f4315d5daee7fbf7e5c5ade6ad72a08ad17 Mon Sep 17 00:00:00 2001
From: Timur Makarchuk <makarchuk@group-ib.com>
Date: Wed, 23 Jan 2019 19:53:03 +0300
Subject: [PATCH] Add checks for json/easyjson marshaler/unmarshaler interfaces

---
 gen/decoder.go | 26 ++++++++++++++++++--------
 gen/encoder.go | 25 +++++++++++++++++--------
 2 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/gen/decoder.go b/gen/decoder.go
index 606602f..859390b 100644
--- a/gen/decoder.go
+++ b/gen/decoder.go
@@ -259,15 +259,15 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 
 	case reflect.Interface:
 		if t.NumMethod() != 0 {
-			return fmt.Errorf("interface type %v not supported: only interface{} is allowed", t)
+			if g.interfaceIsEasyjsonUnmarshaller(t) {
+				fmt.Fprintln(g.out, ws+out+".UnmarshalEasyJSON(in)")
+			} else if g.interfaceIsJsonUnmarshaller(t) {
+				fmt.Fprintln(g.out, ws+out+".UnmarshalJSON(in.Raw())")
+			} else {
+				return fmt.Errorf("interface type %v not supported: only interface{} and easyjson/json Unmarshaler are allowed", t)
+			}
 		}
-		fmt.Fprintln(g.out, ws+"if m, ok := "+out+".(easyjson.Unmarshaler); ok {")
-		fmt.Fprintln(g.out, ws+"m.UnmarshalEasyJSON(in)")
-		fmt.Fprintln(g.out, ws+"} else if m, ok := "+out+".(json.Unmarshaler); ok {")
-		fmt.Fprintln(g.out, ws+"_ = m.UnmarshalJSON(in.Raw())")
-		fmt.Fprintln(g.out, ws+"} else {")
-		fmt.Fprintln(g.out, ws+"  "+out+" = in.Interface()")
-		fmt.Fprintln(g.out, ws+"}")
+		fmt.Fprintln(g.out, ws+out+" = in.Interface()")
 	default:
 		return fmt.Errorf("don't know how to decode %v", t)
 	}
@@ -275,6 +275,16 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 
 }
 
+func (g *Generator) interfaceIsEasyjsonUnmarshaller(t reflect.Type) bool {
+	unmarshalerType := reflect.TypeOf((*easyjson.Unmarshaler)(nil))
+	return t.Implements(unmarshalerType)
+}
+
+func (g *Generator) interfaceIsJsonUnmarshaller(t reflect.Type) bool {
+	unmarshalerType := reflect.TypeOf((*json.Unmarshaler)(nil))
+	return t.Implements(unmarshalerType)
+}
+
 func (g *Generator) genStructFieldDecoder(t reflect.Type, f reflect.StructField) error {
 	jsonName := g.fieldNamer.GetJSONFieldName(t, f)
 	tags := parseFieldTags(f)
diff --git a/gen/encoder.go b/gen/encoder.go
index b2be743..9c0ff3b 100644
--- a/gen/encoder.go
+++ b/gen/encoder.go
@@ -243,22 +243,31 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 
 	case reflect.Interface:
 		if t.NumMethod() != 0 {
-			return fmt.Errorf("interface type %v not supported: only interface{} is allowed", t)
+			if g.interfaceIsEasyjsonMarshaller(t) {
+				fmt.Fprintln(g.out, ws+in+".MarshalEasyJSON(out)")
+			} else if g.interfaceIsJsonMarshaller(t) {
+				fmt.Fprintln(g.out, ws+in+".MarshalJSON(out)")
+			} else {
+				return fmt.Errorf("interface type %v not supported: only interface{} was allowed", t)
+			}
 		}
-		fmt.Fprintln(g.out, ws+"if m, ok := "+in+".(easyjson.Marshaler); ok {")
-		fmt.Fprintln(g.out, ws+"  m.MarshalEasyJSON(out)")
-		fmt.Fprintln(g.out, ws+"} else if m, ok := "+in+".(json.Marshaler); ok {")
-		fmt.Fprintln(g.out, ws+"  out.Raw(m.MarshalJSON())")
-		fmt.Fprintln(g.out, ws+"} else {")
 		fmt.Fprintln(g.out, ws+"  out.Raw(json.Marshal("+in+"))")
-		fmt.Fprintln(g.out, ws+"}")
-
 	default:
 		return fmt.Errorf("don't know how to encode %v", t)
 	}
 	return nil
 }
 
+func (g *Generator) interfaceIsEasyjsonMarshaller(t reflect.Type) bool {
+	marshalerType := reflect.TypeOf((*easyjson.Marshaler)(nil))
+	return t.Implements(marshalerType)
+}
+
+func (g *Generator) interfaceIsJsonMarshaller(t reflect.Type) bool {
+	marshalerType := reflect.TypeOf((*json.Marshaler)(nil))
+	return t.Implements(marshalerType)
+}
+
 func (g *Generator) notEmptyCheck(t reflect.Type, v string) string {
 	optionalIface := reflect.TypeOf((*easyjson.Optional)(nil)).Elem()
 	if reflect.PtrTo(t).Implements(optionalIface) {