From 82d245f54ea4167731f3e7030b3d9d1b028a115d Mon Sep 17 00:00:00 2001
From: Timur <timur.makarchuk@gmail.com>
Date: Sat, 23 Mar 2019 12:19:38 +0300
Subject: [PATCH] Better support for custom interfaces. Bring all of the
 original checks back

---
 gen/decoder.go | 18 ++++++++++++------
 gen/encoder.go | 24 ++++++++++++++++--------
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/gen/decoder.go b/gen/decoder.go
index 859390b..213664f 100644
--- a/gen/decoder.go
+++ b/gen/decoder.go
@@ -86,7 +86,7 @@ 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
+// returns true if 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()) ||
@@ -258,6 +258,7 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 		fmt.Fprintln(g.out, ws+"}")
 
 	case reflect.Interface:
+		fmt.Printf("//%v: %v", out, g.interfaceIsEasyjsonUnmarshaller(t))
 		if t.NumMethod() != 0 {
 			if g.interfaceIsEasyjsonUnmarshaller(t) {
 				fmt.Fprintln(g.out, ws+out+".UnmarshalEasyJSON(in)")
@@ -266,8 +267,15 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 			} else {
 				return fmt.Errorf("interface type %v not supported: only interface{} and easyjson/json Unmarshaler are allowed", t)
 			}
+		} else {
+			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)
 	}
@@ -276,13 +284,11 @@ 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)
+	return t.Implements(reflect.TypeOf((*easyjson.Unmarshaler)(nil)).Elem())
 }
 
 func (g *Generator) interfaceIsJsonUnmarshaller(t reflect.Type) bool {
-	unmarshalerType := reflect.TypeOf((*json.Unmarshaler)(nil))
-	return t.Implements(unmarshalerType)
+	return t.Implements(reflect.TypeOf((*json.Unmarshaler)(nil)).Elem())
 }
 
 func (g *Generator) genStructFieldDecoder(t reflect.Type, f reflect.StructField) error {
diff --git a/gen/encoder.go b/gen/encoder.go
index 9c0ff3b..465e053 100644
--- a/gen/encoder.go
+++ b/gen/encoder.go
@@ -110,7 +110,7 @@ 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
+// returns true if 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()) ||
@@ -245,13 +245,23 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 		if t.NumMethod() != 0 {
 			if g.interfaceIsEasyjsonMarshaller(t) {
 				fmt.Fprintln(g.out, ws+in+".MarshalEasyJSON(out)")
-			} else if g.interfaceIsJsonMarshaller(t) {
+			} else if g.interfaceIsJSONMarshaller(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 {")
 				fmt.Fprintln(g.out, ws+in+".MarshalJSON(out)")
+				fmt.Fprintln(g.out, ws+"}")
 			} else {
-				return fmt.Errorf("interface type %v not supported: only interface{} was allowed", t)
+				return fmt.Errorf("interface type %v not supported: only interface{} and interfaces that implement json or easyjson Marshaling are 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)
 	}
@@ -259,13 +269,11 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 }
 
 func (g *Generator) interfaceIsEasyjsonMarshaller(t reflect.Type) bool {
-	marshalerType := reflect.TypeOf((*easyjson.Marshaler)(nil))
-	return t.Implements(marshalerType)
+	return t.Implements(reflect.TypeOf((*easyjson.Marshaler)(nil)).Elem())
 }
 
-func (g *Generator) interfaceIsJsonMarshaller(t reflect.Type) bool {
-	marshalerType := reflect.TypeOf((*json.Marshaler)(nil))
-	return t.Implements(marshalerType)
+func (g *Generator) interfaceIsJSONMarshaller(t reflect.Type) bool {
+	return t.Implements(reflect.TypeOf((*json.Marshaler)(nil)).Elem())
 }
 
 func (g *Generator) notEmptyCheck(t reflect.Type, v string) string {