1package toml
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"math"
8	"reflect"
9	"sort"
10	"strconv"
11	"strings"
12	"time"
13)
14
15// Encodes a string to a TOML-compliant multi-line string value
16// This function is a clone of the existing encodeTomlString function, except that whitespace characters
17// are preserved. Quotation marks and backslashes are also not escaped.
18func encodeMultilineTomlString(value string) string {
19	var b bytes.Buffer
20
21	for _, rr := range value {
22		switch rr {
23		case '\b':
24			b.WriteString(`\b`)
25		case '\t':
26			b.WriteString("\t")
27		case '\n':
28			b.WriteString("\n")
29		case '\f':
30			b.WriteString(`\f`)
31		case '\r':
32			b.WriteString("\r")
33		case '"':
34			b.WriteString(`"`)
35		case '\\':
36			b.WriteString(`\`)
37		default:
38			intRr := uint16(rr)
39			if intRr < 0x001F {
40				b.WriteString(fmt.Sprintf("\\u%0.4X", intRr))
41			} else {
42				b.WriteRune(rr)
43			}
44		}
45	}
46	return b.String()
47}
48
49// Encodes a string to a TOML-compliant string value
50func encodeTomlString(value string) string {
51	var b bytes.Buffer
52
53	for _, rr := range value {
54		switch rr {
55		case '\b':
56			b.WriteString(`\b`)
57		case '\t':
58			b.WriteString(`\t`)
59		case '\n':
60			b.WriteString(`\n`)
61		case '\f':
62			b.WriteString(`\f`)
63		case '\r':
64			b.WriteString(`\r`)
65		case '"':
66			b.WriteString(`\"`)
67		case '\\':
68			b.WriteString(`\\`)
69		default:
70			intRr := uint16(rr)
71			if intRr < 0x001F {
72				b.WriteString(fmt.Sprintf("\\u%0.4X", intRr))
73			} else {
74				b.WriteRune(rr)
75			}
76		}
77	}
78	return b.String()
79}
80
81func tomlValueStringRepresentation(v interface{}, indent string, arraysOneElementPerLine bool) (string, error) {
82	// this interface check is added to dereference the change made in the writeTo function.
83	// That change was made to allow this function to see formatting options.
84	tv, ok := v.(*tomlValue)
85	if ok {
86		v = tv.value
87	} else {
88		tv = &tomlValue{}
89	}
90
91	switch value := v.(type) {
92	case uint64:
93		return strconv.FormatUint(value, 10), nil
94	case int64:
95		return strconv.FormatInt(value, 10), nil
96	case float64:
97		// Ensure a round float does contain a decimal point. Otherwise feeding
98		// the output back to the parser would convert to an integer.
99		if math.Trunc(value) == value {
100			return strings.ToLower(strconv.FormatFloat(value, 'f', 1, 32)), nil
101		}
102		return strings.ToLower(strconv.FormatFloat(value, 'f', -1, 32)), nil
103	case string:
104		if tv.multiline {
105			return "\"\"\"\n" + encodeMultilineTomlString(value) + "\"\"\"", nil
106		}
107		return "\"" + encodeTomlString(value) + "\"", nil
108	case []byte:
109		b, _ := v.([]byte)
110		return tomlValueStringRepresentation(string(b), indent, arraysOneElementPerLine)
111	case bool:
112		if value {
113			return "true", nil
114		}
115		return "false", nil
116	case time.Time:
117		return value.Format(time.RFC3339), nil
118	case nil:
119		return "", nil
120	}
121
122	rv := reflect.ValueOf(v)
123
124	if rv.Kind() == reflect.Slice {
125		var values []string
126		for i := 0; i < rv.Len(); i++ {
127			item := rv.Index(i).Interface()
128			itemRepr, err := tomlValueStringRepresentation(item, indent, arraysOneElementPerLine)
129			if err != nil {
130				return "", err
131			}
132			values = append(values, itemRepr)
133		}
134		if arraysOneElementPerLine && len(values) > 1 {
135			stringBuffer := bytes.Buffer{}
136			valueIndent := indent + `  ` // TODO: move that to a shared encoder state
137
138			stringBuffer.WriteString("[\n")
139
140			for _, value := range values {
141				stringBuffer.WriteString(valueIndent)
142				stringBuffer.WriteString(value)
143				stringBuffer.WriteString(`,`)
144				stringBuffer.WriteString("\n")
145			}
146
147			stringBuffer.WriteString(indent + "]")
148
149			return stringBuffer.String(), nil
150		}
151		return "[" + strings.Join(values, ",") + "]", nil
152	}
153	return "", fmt.Errorf("unsupported value type %T: %v", v, v)
154}
155
156func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool) (int64, error) {
157	simpleValuesKeys := make([]string, 0)
158	complexValuesKeys := make([]string, 0)
159
160	for k := range t.values {
161		v := t.values[k]
162		switch v.(type) {
163		case *Tree, []*Tree:
164			complexValuesKeys = append(complexValuesKeys, k)
165		default:
166			simpleValuesKeys = append(simpleValuesKeys, k)
167		}
168	}
169
170	sort.Strings(simpleValuesKeys)
171	sort.Strings(complexValuesKeys)
172
173	for _, k := range simpleValuesKeys {
174		v, ok := t.values[k].(*tomlValue)
175		if !ok {
176			return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
177		}
178
179		repr, err := tomlValueStringRepresentation(v, indent, arraysOneElementPerLine)
180		if err != nil {
181			return bytesCount, err
182		}
183
184		if v.comment != "" {
185			comment := strings.Replace(v.comment, "\n", "\n"+indent+"#", -1)
186			start := "# "
187			if strings.HasPrefix(comment, "#") {
188				start = ""
189			}
190			writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment, "\n")
191			bytesCount += int64(writtenBytesCountComment)
192			if errc != nil {
193				return bytesCount, errc
194			}
195		}
196
197		var commented string
198		if v.commented {
199			commented = "# "
200		}
201		writtenBytesCount, err := writeStrings(w, indent, commented, k, " = ", repr, "\n")
202		bytesCount += int64(writtenBytesCount)
203		if err != nil {
204			return bytesCount, err
205		}
206	}
207
208	for _, k := range complexValuesKeys {
209		v := t.values[k]
210
211		combinedKey := k
212		if keyspace != "" {
213			combinedKey = keyspace + "." + combinedKey
214		}
215		var commented string
216		if t.commented {
217			commented = "# "
218		}
219
220		switch node := v.(type) {
221		// node has to be of those two types given how keys are sorted above
222		case *Tree:
223			tv, ok := t.values[k].(*Tree)
224			if !ok {
225				return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
226			}
227			if tv.comment != "" {
228				comment := strings.Replace(tv.comment, "\n", "\n"+indent+"#", -1)
229				start := "# "
230				if strings.HasPrefix(comment, "#") {
231					start = ""
232				}
233				writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment)
234				bytesCount += int64(writtenBytesCountComment)
235				if errc != nil {
236					return bytesCount, errc
237				}
238			}
239			writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[", combinedKey, "]\n")
240			bytesCount += int64(writtenBytesCount)
241			if err != nil {
242				return bytesCount, err
243			}
244			bytesCount, err = node.writeTo(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine)
245			if err != nil {
246				return bytesCount, err
247			}
248		case []*Tree:
249			for _, subTree := range node {
250				writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[[", combinedKey, "]]\n")
251				bytesCount += int64(writtenBytesCount)
252				if err != nil {
253					return bytesCount, err
254				}
255
256				bytesCount, err = subTree.writeTo(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine)
257				if err != nil {
258					return bytesCount, err
259				}
260			}
261		}
262	}
263
264	return bytesCount, nil
265}
266
267func writeStrings(w io.Writer, s ...string) (int, error) {
268	var n int
269	for i := range s {
270		b, err := io.WriteString(w, s[i])
271		n += b
272		if err != nil {
273			return n, err
274		}
275	}
276	return n, nil
277}
278
279// WriteTo encode the Tree as Toml and writes it to the writer w.
280// Returns the number of bytes written in case of success, or an error if anything happened.
281func (t *Tree) WriteTo(w io.Writer) (int64, error) {
282	return t.writeTo(w, "", "", 0, false)
283}
284
285// ToTomlString generates a human-readable representation of the current tree.
286// Output spans multiple lines, and is suitable for ingest by a TOML parser.
287// If the conversion cannot be performed, ToString returns a non-nil error.
288func (t *Tree) ToTomlString() (string, error) {
289	var buf bytes.Buffer
290	_, err := t.WriteTo(&buf)
291	if err != nil {
292		return "", err
293	}
294	return buf.String(), nil
295}
296
297// String generates a human-readable representation of the current tree.
298// Alias of ToString. Present to implement the fmt.Stringer interface.
299func (t *Tree) String() string {
300	result, _ := t.ToTomlString()
301	return result
302}
303
304// ToMap recursively generates a representation of the tree using Go built-in structures.
305// The following types are used:
306//
307//	* bool
308//	* float64
309//	* int64
310//	* string
311//	* uint64
312//	* time.Time
313//	* map[string]interface{} (where interface{} is any of this list)
314//	* []interface{} (where interface{} is any of this list)
315func (t *Tree) ToMap() map[string]interface{} {
316	result := map[string]interface{}{}
317
318	for k, v := range t.values {
319		switch node := v.(type) {
320		case []*Tree:
321			var array []interface{}
322			for _, item := range node {
323				array = append(array, item.ToMap())
324			}
325			result[k] = array
326		case *Tree:
327			result[k] = node.ToMap()
328		case *tomlValue:
329			result[k] = node.value
330		}
331	}
332	return result
333}
334