From 5c5e85117e45bc20ba67974ae27021b457716257 Mon Sep 17 00:00:00 2001
From: Marat Khasanov <marat.khasanov@lazada.com>
Date: Tue, 11 Oct 2016 12:16:52 +0300
Subject: [PATCH] Allow interface{} fields to be marshaled using
 easyjson.Marshaler methods.

---
 Makefile             |  1 +
 gen/encoder.go       |  4 +++-
 gen/generator.go     |  9 ++++++---
 tests/basic_test.go  | 23 +++++++++++++++++++++++
 tests/nested_easy.go | 23 +++++++++++++++++++++++
 5 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 tests/nested_easy.go

diff --git a/Makefile b/Makefile
index 5a95104..e2c9bd6 100644
--- a/Makefile
+++ b/Makefile
@@ -28,6 +28,7 @@ generate: root build
 	.root/bin/easyjson -snake_case .root/src/$(PKG)/tests/snake.go
 	.root/bin/easyjson -omit_empty .root/src/$(PKG)/tests/omitempty.go
 	.root/bin/easyjson -build_tags=use_easyjson .root/src/$(PKG)/benchmark/data.go
+	.root/bin/easyjson .root/src/$(PKG)/tests/nested_easy.go
 
 test: generate root
 	go test \
diff --git a/gen/encoder.go b/gen/encoder.go
index 06303e0..b67f1e9 100644
--- a/gen/encoder.go
+++ b/gen/encoder.go
@@ -199,7 +199,9 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 		if t.NumMethod() != 0 {
 			return fmt.Errorf("interface type %v not supported: only interface{} is allowed", t)
 		}
-		fmt.Fprintln(g.out, ws+"if m, ok := "+in+".(json.Marshaler); ok {")
+		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+"))")
diff --git a/gen/generator.go b/gen/generator.go
index 1c527d5..78edaff 100644
--- a/gen/generator.go
+++ b/gen/generator.go
@@ -15,6 +15,7 @@ import (
 
 const pkgWriter = "github.com/mailru/easyjson/jwriter"
 const pkgLexer = "github.com/mailru/easyjson/jlexer"
+const pkgEasyjson = "github.com/mailru/easyjson"
 
 // FieldNamer defines a policy for generating names for struct fields.
 type FieldNamer interface {
@@ -59,6 +60,7 @@ func NewGenerator(filename string) *Generator {
 		imports: map[string]string{
 			pkgWriter:       "jwriter",
 			pkgLexer:        "jlexer",
+			pkgEasyjson:      "easyjson",
 			"encoding/json": "json",
 		},
 		fieldNamer:    DefaultFieldNamer{},
@@ -159,9 +161,10 @@ func (g *Generator) printHeader() {
 	fmt.Println("")
 	fmt.Println("// suppress unused package warning")
 	fmt.Println("var (")
-	fmt.Println("   _ = json.RawMessage{}")
-	fmt.Println("   _ = jlexer.Lexer{}")
-	fmt.Println("   _ = jwriter.Writer{}")
+	fmt.Println("   _ *json.RawMessage")
+	fmt.Println("   _ *jlexer.Lexer")
+	fmt.Println("   _ *jwriter.Writer")
+	fmt.Println("   _ easyjson.Marshaler")
 	fmt.Println(")")
 
 	fmt.Println()
diff --git a/tests/basic_test.go b/tests/basic_test.go
index 3e994ea..8292788 100644
--- a/tests/basic_test.go
+++ b/tests/basic_test.go
@@ -156,3 +156,26 @@ func TestUnderflowArray(t *testing.T) {
 		t.Errorf("Unmarshal(%v) = %+v; want %+v", arrayUnderflowString, a, arrayUnderflowValue)
 	}
 }
+
+func TestNestedEasyJsonMarshal(t *testing.T) {
+	n := map[string]*NestedEasyMarshaler{
+		"Value":  {},
+		"Slice1": {},
+		"Slice2": {},
+		"Map1":   {},
+		"Map2":   {},
+	}
+
+	ni := NestedInterfaces{
+		Value: n["Value"],
+		Slice: []interface{}{n["Slice1"], n["Slice2"]},
+		Map:   map[string]interface{}{"1": n["Map1"], "2": n["Map2"]},
+	}
+	easyjson.Marshal(ni)
+
+	for k, v := range n {
+		if !v.EasilyMarshaled {
+			t.Errorf("Nested interface %s wasn't easily marshaled", k)
+		}
+	}
+}
diff --git a/tests/nested_easy.go b/tests/nested_easy.go
new file mode 100644
index 0000000..c7f0f4b
--- /dev/null
+++ b/tests/nested_easy.go
@@ -0,0 +1,23 @@
+package tests
+
+import (
+	"github.com/mailru/easyjson"
+	"github.com/mailru/easyjson/jwriter"
+)
+
+//easyjson:json
+type NestedInterfaces struct {
+	Value interface{}
+	Slice []interface{}
+	Map   map[string]interface{}
+}
+
+type NestedEasyMarshaler struct {
+	EasilyMarshaled bool
+}
+
+var _ easyjson.Marshaler = &NestedEasyMarshaler{}
+
+func (i *NestedEasyMarshaler) MarshalEasyJSON(w *jwriter.Writer) {
+	i.EasilyMarshaled = true
+}
\ No newline at end of file