Allow MarshalTOML and MarshalText to be used on the document type itself

Fixes #383

Author: arp242

encode.go

@@ -136,6 +136,9 @@ func NewEncoder(w io.Writer) *Encoder {
 // document.
 func (enc *Encoder) Encode(v interface{}) error {
 	rv := eindirect(reflect.ValueOf(v))
+
+	// XXX
+
 	if err := enc.safeEncode(Key([]string{}), rv); err != nil {
 		return err
 	}
@@ -693,8 +696,11 @@ func (enc *Encoder) newline() {
 //	v      v   v  v    vv
 //	key = {k = 1, k2 = 2}
 func (enc *Encoder) writeKeyValue(key Key, val reflect.Value, inline bool) {
+	/// Marshaler used on top-level document; call eElement() to just call
+	/// Marshal{TOML,Text}.
 	if len(key) == 0 {
-		encPanic(errNoKey)
+		enc.eElement(val)
+		return
 	}
 	enc.wf("%s%s = ", enc.indentStr(key), key.maybeQuoted(len(key)-1))
 	enc.eElement(val)

encode_test.go

@@ -1261,6 +1261,43 @@ c = 3
 	}
 }
 
+type (
+	Doc1 struct{ N string }
+	Doc2 struct{ N string }
+)
+
+func (d Doc1) MarshalTOML() ([]byte, error) { return []byte(`marshal_toml = "` + d.N + `"`), nil }
+func (d Doc2) MarshalText() ([]byte, error) { return []byte(`marshal_text = "` + d.N + `"`), nil }
+
+// MarshalTOML and MarshalText on the top level type, rather than a field.
+func TestMarshalDoc(t *testing.T) {
+	t.Run("toml", func(t *testing.T) {
+		var buf bytes.Buffer
+		err := NewEncoder(&buf).Encode(Doc1{"asd"})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		want := `marshal_toml = "asd"`
+		if want != buf.String() {
+			t.Errorf("\nhave: %s\nwant: %s\n", buf.String(), want)
+		}
+	})
+
+	t.Run("text", func(t *testing.T) {
+		var buf bytes.Buffer
+		err := NewEncoder(&buf).Encode(Doc2{"asd"})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		want := `"marshal_text = \"asd\""`
+		if want != buf.String() {
+			t.Errorf("\nhave: %s\nwant: %s\n", buf.String(), want)
+		}
+	})
+}
+
 func encodeExpected(t *testing.T, label string, val interface{}, want string, wantErr error) {
 	t.Helper()
 	t.Run(label, func(t *testing.T) {