From b7aa2cbb250709566f25c1d8c7157400d66ff721 Mon Sep 17 00:00:00 2001
From: Egor Kovetskiy <e.kovetskiy@office.ngs.ru>
Date: Mon, 23 May 2016 23:49:25 +0600
Subject: [PATCH] add Marshaler interface

---
 encode.go      | 29 +++++++++++++++++++
 encode_test.go | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/encode.go b/encode.go
index f538261a..cee51c98 100644
--- a/encode.go
+++ b/encode.go
@@ -38,6 +38,12 @@ var quotedReplacer = strings.NewReplacer(
 	"\\", "\\\\",
 )
 
+// Marshaler is the interface implemented by objects that can marshal
+// themselves to a TOML representation.
+type Marshaler interface {
+	MarshalTOML() ([]byte, error)
+}
+
 // Encoder controls the encoding of Go values to a TOML document to some
 // io.Writer.
 //
@@ -106,6 +112,10 @@ func (enc *Encoder) safeEncode(key Key, rv reflect.Value) (err error) {
 }
 
 func (enc *Encoder) encode(key Key, rv reflect.Value) {
+	if _, ok := rv.Interface().(Marshaler); ok {
+		rv = marshalDecode(rv)
+	}
+
 	// Special case. Time needs to be in ISO8601 format.
 	// Special case. If we can marshal the type to text, then we used that.
 	// Basically, this prevents the encoder for handling these types as
@@ -253,6 +263,10 @@ func (enc *Encoder) eTable(key Key, rv reflect.Value) {
 }
 
 func (enc *Encoder) eMapOrStruct(key Key, rv reflect.Value) {
+	if _, ok := rv.Interface().(Marshaler); ok {
+		rv = marshalDecode(rv)
+	}
+
 	switch rv := eindirect(rv); rv.Kind() {
 	case reflect.Map:
 		enc.eMap(key, rv)
@@ -537,6 +551,21 @@ func encPanic(err error) {
 	panic(tomlEncodeError{err})
 }
 
+func marshalDecode(rv reflect.Value) reflect.Value {
+	buf, err := rv.Interface().(Marshaler).MarshalTOML()
+	if err != nil {
+		encPanic(err)
+	}
+
+	v := map[string]interface{}{}
+	_, err = Decode(string(buf), &v)
+	if err != nil {
+		encPanic(err)
+	}
+
+	return reflect.ValueOf(v)
+}
+
 func eindirect(v reflect.Value) reflect.Value {
 	switch v.Kind() {
 	case reflect.Ptr, reflect.Interface:
diff --git a/encode_test.go b/encode_test.go
index 673b7b00..87987278 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -531,6 +531,83 @@ func TestEncodeAnonymousStruct(t *testing.T) {
 	encodeExpected(t, "embedded anonymous tagged struct", v1, expected, nil)
 }
 
+type withMarshalTOML struct{}
+
+func (withMarshalTOML) MarshalTOML() ([]byte, error) {
+	return []byte("x = 1\ny = 1\n"), nil
+}
+
+func TestEncodeMarshalerAnonymousStruct(t *testing.T) {
+	type Inner struct {
+		N int
+		withMarshalTOML
+	}
+	type Outer0 struct{ Inner }
+	type Outer1 struct {
+		Inner `toml:"inner"`
+		X     int
+	}
+
+	v0 := Inner{}
+	expected := "x = 1\ny = 1\n"
+	encodeExpected(t, "embedded struct with MarshalTOML", v0, expected, nil)
+
+	v1 := Outer0{Inner{}}
+	expected = "x = 1\ny = 1\n"
+	encodeExpected(t, "embedded anonymous untagged struct with MarshalTOML", v1, expected, nil)
+
+	v2 := Outer1{Inner{}, 3}
+	expected = "x = 1\ny = 1\n"
+	encodeExpected(t, "embedded anonymous tagged struct with MarshalTOML", v2, expected, nil)
+}
+
+func TestEncodeMarshalerField(t *testing.T) {
+	type Inner struct {
+		withMarshalTOML
+		N int
+		Z int
+	}
+	type Outer struct {
+		A     int
+		Field Inner
+	}
+	type Outer1 struct {
+		B   int
+		Ins []Inner `toml:"slice"`
+	}
+
+	v0 := Outer{
+		A:     1,
+		Field: Inner{},
+	}
+	expected := `A = 1
+
+[Field]
+  x = 1
+  y = 1
+`
+	encodeExpected(t, "struct field with MarshalTOML", v0, expected, nil)
+
+	v1 := Outer1{
+		B: 1,
+		Ins: []Inner{
+			Inner{},
+			Inner{},
+		},
+	}
+	expected = `B = 1
+
+[[slice]]
+  x = 1
+  y = 1
+
+[[slice]]
+  x = 1
+  y = 1
+`
+	encodeExpected(t, "struct slice field with MarshalTOML", v1, expected, nil)
+}
+
 func TestEncodeAnonymousStructPointerField(t *testing.T) {
 	type Inner struct{ N int }
 	type Outer0 struct{ *Inner }