1package toml
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"math"
8	"math/big"
9	"reflect"
10	"sort"
11	"strconv"
12	"strings"
13	"time"
14)
15
16type valueComplexity int
17
18const (
19	valueSimple valueComplexity = iota + 1
20	valueComplex
21)
22
23type sortNode struct {
24	key        string
25	complexity valueComplexity
26}
27
28// Encodes a string to a TOML-compliant multi-line string value
29// This function is a clone of the existing encodeTomlString function, except that whitespace characters
30// are preserved. Quotation marks and backslashes are also not escaped.
31func encodeMultilineTomlString(value string, commented string) string {
32	var b bytes.Buffer
33	adjacentQuoteCount := 0
34
35	b.WriteString(commented)
36	for i, rr := range value {
37		if rr != '"' {
38			adjacentQuoteCount = 0
39		} else {
40			adjacentQuoteCount++
41		}
42		switch rr {
43		case '\b':
44			b.WriteString(`\b`)
45		case '\t':
46			b.WriteString("\t")
47		case '\n':
48			b.WriteString("\n" + commented)
49		case '\f':
50			b.WriteString(`\f`)
51		case '\r':
52			b.WriteString("\r")
53		case '"':
54			if adjacentQuoteCount >= 3 || i == len(value)-1 {
55				adjacentQuoteCount = 0
56				b.WriteString(`\"`)
57			} else {
58				b.WriteString(`"`)
59			}
60		case '\\':
61			b.WriteString(`\`)
62		default:
63			intRr := uint16(rr)
64			if intRr < 0x001F {
65				b.WriteString(fmt.Sprintf("\\u%0.4X", intRr))
66			} else {
67				b.WriteRune(rr)
68			}
69		}
70	}
71	return b.String()
72}
73
74// Encodes a string to a TOML-compliant string value
75func encodeTomlString(value string) string {
76	var b bytes.Buffer
77
78	for _, rr := range value {
79		switch rr {
80		case '\b':
81			b.WriteString(`\b`)
82		case '\t':
83			b.WriteString(`\t`)
84		case '\n':
85			b.WriteString(`\n`)
86		case '\f':
87			b.WriteString(`\f`)
88		case '\r':
89			b.WriteString(`\r`)
90		case '"':
91			b.WriteString(`\"`)
92		case '\\':
93			b.WriteString(`\\`)
94		default:
95			intRr := uint16(rr)
96			if intRr < 0x001F {
97				b.WriteString(fmt.Sprintf("\\u%0.4X", intRr))
98			} else {
99				b.WriteRune(rr)
100			}
101		}
102	}
103	return b.String()
104}
105
106func tomlTreeStringRepresentation(t *Tree, ord MarshalOrder) (string, error) {
107	var orderedVals []sortNode
108	switch ord {
109	case OrderPreserve:
110		orderedVals = sortByLines(t)
111	default:
112		orderedVals = sortAlphabetical(t)
113	}
114
115	var values []string
116	for _, node := range orderedVals {
117		k := node.key
118		v := t.values[k]
119
120		repr, err := tomlValueStringRepresentation(v, "", "", ord, false)
121		if err != nil {
122			return "", err
123		}
124		values = append(values, quoteKeyIfNeeded(k)+" = "+repr)
125	}
126	return "{ " + strings.Join(values, ", ") + " }", nil
127}
128
129func tomlValueStringRepresentation(v interface{}, commented string, indent string, ord MarshalOrder, arraysOneElementPerLine bool) (string, error) {
130	// this interface check is added to dereference the change made in the writeTo function.
131	// That change was made to allow this function to see formatting options.
132	tv, ok := v.(*tomlValue)
133	if ok {
134		v = tv.value
135	} else {
136		tv = &tomlValue{}
137	}
138
139	switch value := v.(type) {
140	case uint64:
141		return strconv.FormatUint(value, 10), nil
142	case int64:
143		return strconv.FormatInt(value, 10), nil
144	case float64:
145		// Default bit length is full 64
146		bits := 64
147		// Float panics if nan is used
148		if !math.IsNaN(value) {
149			// if 32 bit accuracy is enough to exactly show, use 32
150			_, acc := big.NewFloat(value).Float32()
151			if acc == big.Exact {
152				bits = 32
153			}
154		}
155		if math.Trunc(value) == value {
156			return strings.ToLower(strconv.FormatFloat(value, 'f', 1, bits)), nil
157		}
158		return strings.ToLower(strconv.FormatFloat(value, 'f', -1, bits)), nil
159	case string:
160		if tv.multiline {
161			if tv.literal {
162				b := strings.Builder{}
163				b.WriteString("'''\n")
164				b.Write([]byte(value))
165				b.WriteString("\n'''")
166				return b.String(), nil
167			} else {
168				return "\"\"\"\n" + encodeMultilineTomlString(value, commented) + "\"\"\"", nil
169			}
170		}
171		return "\"" + encodeTomlString(value) + "\"", nil
172	case []byte:
173		b, _ := v.([]byte)
174		return string(b), nil
175	case bool:
176		if value {
177			return "true", nil
178		}
179		return "false", nil
180	case time.Time:
181		return value.Format(time.RFC3339), nil
182	case LocalDate:
183		return value.String(), nil
184	case LocalDateTime:
185		return value.String(), nil
186	case LocalTime:
187		return value.String(), nil
188	case *Tree:
189		return tomlTreeStringRepresentation(value, ord)
190	case nil:
191		return "", nil
192	}
193
194	rv := reflect.ValueOf(v)
195
196	if rv.Kind() == reflect.Slice {
197		var values []string
198		for i := 0; i < rv.Len(); i++ {
199			item := rv.Index(i).Interface()
200			itemRepr, err := tomlValueStringRepresentation(item, commented, indent, ord, arraysOneElementPerLine)
201			if err != nil {
202				return "", err
203			}
204			values = append(values, itemRepr)
205		}
206		if arraysOneElementPerLine && len(values) > 1 {
207			stringBuffer := bytes.Buffer{}
208			valueIndent := indent + `  ` // TODO: move that to a shared encoder state
209
210			stringBuffer.WriteString("[\n")
211
212			for _, value := range values {
213				stringBuffer.WriteString(valueIndent)
214				stringBuffer.WriteString(commented + value)
215				stringBuffer.WriteString(`,`)
216				stringBuffer.WriteString("\n")
217			}
218
219			stringBuffer.WriteString(indent + commented + "]")
220
221			return stringBuffer.String(), nil
222		}
223		return "[" + strings.Join(values, ", ") + "]", nil
224	}
225	return "", fmt.Errorf("unsupported value type %T: %v", v, v)
226}
227
228func getTreeArrayLine(trees []*Tree) (line int) {
229	// get lowest line number that is not 0
230	for _, tv := range trees {
231		if tv.position.Line < line || line == 0 {
232			line = tv.position.Line
233		}
234	}
235	return
236}
237
238func sortByLines(t *Tree) (vals []sortNode) {
239	var (
240		line  int
241		lines []int
242		tv    *Tree
243		tom   *tomlValue
244		node  sortNode
245	)
246	vals = make([]sortNode, 0)
247	m := make(map[int]sortNode)
248
249	for k := range t.values {
250		v := t.values[k]
251		switch v.(type) {
252		case *Tree:
253			tv = v.(*Tree)
254			line = tv.position.Line
255			node = sortNode{key: k, complexity: valueComplex}
256		case []*Tree:
257			line = getTreeArrayLine(v.([]*Tree))
258			node = sortNode{key: k, complexity: valueComplex}
259		default:
260			tom = v.(*tomlValue)
261			line = tom.position.Line
262			node = sortNode{key: k, complexity: valueSimple}
263		}
264		lines = append(lines, line)
265		vals = append(vals, node)
266		m[line] = node
267	}
268	sort.Ints(lines)
269
270	for i, line := range lines {
271		vals[i] = m[line]
272	}
273
274	return vals
275}
276
277func sortAlphabetical(t *Tree) (vals []sortNode) {
278	var (
279		node     sortNode
280		simpVals []string
281		compVals []string
282	)
283	vals = make([]sortNode, 0)
284	m := make(map[string]sortNode)
285
286	for k := range t.values {
287		v := t.values[k]
288		switch v.(type) {
289		case *Tree, []*Tree:
290			node = sortNode{key: k, complexity: valueComplex}
291			compVals = append(compVals, node.key)
292		default:
293			node = sortNode{key: k, complexity: valueSimple}
294			simpVals = append(simpVals, node.key)
295		}
296		vals = append(vals, node)
297		m[node.key] = node
298	}
299
300	// Simples first to match previous implementation
301	sort.Strings(simpVals)
302	i := 0
303	for _, key := range simpVals {
304		vals[i] = m[key]
305		i++
306	}
307
308	sort.Strings(compVals)
309	for _, key := range compVals {
310		vals[i] = m[key]
311		i++
312	}
313
314	return vals
315}
316
317func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool) (int64, error) {
318	return t.writeToOrdered(w, indent, keyspace, bytesCount, arraysOneElementPerLine, OrderAlphabetical, "  ", false)
319}
320
321func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord MarshalOrder, indentString string, parentCommented bool) (int64, error) {
322	var orderedVals []sortNode
323
324	switch ord {
325	case OrderPreserve:
326		orderedVals = sortByLines(t)
327	default:
328		orderedVals = sortAlphabetical(t)
329	}
330
331	for _, node := range orderedVals {
332		switch node.complexity {
333		case valueComplex:
334			k := node.key
335			v := t.values[k]
336
337			combinedKey := quoteKeyIfNeeded(k)
338			if keyspace != "" {
339				combinedKey = keyspace + "." + combinedKey
340			}
341
342			switch node := v.(type) {
343			// node has to be of those two types given how keys are sorted above
344			case *Tree:
345				tv, ok := t.values[k].(*Tree)
346				if !ok {
347					return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
348				}
349				if tv.comment != "" {
350					comment := strings.Replace(tv.comment, "\n", "\n"+indent+"#", -1)
351					start := "# "
352					if strings.HasPrefix(comment, "#") {
353						start = ""
354					}
355					writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment)
356					bytesCount += int64(writtenBytesCountComment)
357					if errc != nil {
358						return bytesCount, errc
359					}
360				}
361
362				var commented string
363				if parentCommented || t.commented || tv.commented {
364					commented = "# "
365				}
366				writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[", combinedKey, "]\n")
367				bytesCount += int64(writtenBytesCount)
368				if err != nil {
369					return bytesCount, err
370				}
371				bytesCount, err = node.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, parentCommented || t.commented || tv.commented)
372				if err != nil {
373					return bytesCount, err
374				}
375			case []*Tree:
376				for _, subTree := range node {
377					var commented string
378					if parentCommented || t.commented || subTree.commented {
379						commented = "# "
380					}
381					writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[[", combinedKey, "]]\n")
382					bytesCount += int64(writtenBytesCount)
383					if err != nil {
384						return bytesCount, err
385					}
386
387					bytesCount, err = subTree.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, parentCommented || t.commented || subTree.commented)
388					if err != nil {
389						return bytesCount, err
390					}
391				}
392			}
393		default: // Simple
394			k := node.key
395			v, ok := t.values[k].(*tomlValue)
396			if !ok {
397				return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
398			}
399
400			var commented string
401			if parentCommented || t.commented || v.commented {
402				commented = "# "
403			}
404			repr, err := tomlValueStringRepresentation(v, commented, indent, ord, arraysOneElementPerLine)
405			if err != nil {
406				return bytesCount, err
407			}
408
409			if v.comment != "" {
410				comment := strings.Replace(v.comment, "\n", "\n"+indent+"#", -1)
411				start := "# "
412				if strings.HasPrefix(comment, "#") {
413					start = ""
414				}
415				writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment, "\n")
416				bytesCount += int64(writtenBytesCountComment)
417				if errc != nil {
418					return bytesCount, errc
419				}
420			}
421
422			quotedKey := quoteKeyIfNeeded(k)
423			writtenBytesCount, err := writeStrings(w, indent, commented, quotedKey, " = ", repr, "\n")
424			bytesCount += int64(writtenBytesCount)
425			if err != nil {
426				return bytesCount, err
427			}
428		}
429	}
430
431	return bytesCount, nil
432}
433
434// quote a key if it does not fit the bare key format (A-Za-z0-9_-)
435// quoted keys use the same rules as strings
436func quoteKeyIfNeeded(k string) string {
437	// when encoding a map with the 'quoteMapKeys' option enabled, the tree will contain
438	// keys that have already been quoted.
439	// not an ideal situation, but good enough of a stop gap.
440	if len(k) >= 2 && k[0] == '"' && k[len(k)-1] == '"' {
441		return k
442	}
443	isBare := true
444	for _, r := range k {
445		if !isValidBareChar(r) {
446			isBare = false
447			break
448		}
449	}
450	if isBare {
451		return k
452	}
453	return quoteKey(k)
454}
455
456func quoteKey(k string) string {
457	return "\"" + encodeTomlString(k) + "\""
458}
459
460func writeStrings(w io.Writer, s ...string) (int, error) {
461	var n int
462	for i := range s {
463		b, err := io.WriteString(w, s[i])
464		n += b
465		if err != nil {
466			return n, err
467		}
468	}
469	return n, nil
470}
471
472// WriteTo encode the Tree as Toml and writes it to the writer w.
473// Returns the number of bytes written in case of success, or an error if anything happened.
474func (t *Tree) WriteTo(w io.Writer) (int64, error) {
475	return t.writeTo(w, "", "", 0, false)
476}
477
478// ToTomlString generates a human-readable representation of the current tree.
479// Output spans multiple lines, and is suitable for ingest by a TOML parser.
480// If the conversion cannot be performed, ToString returns a non-nil error.
481func (t *Tree) ToTomlString() (string, error) {
482	b, err := t.Marshal()
483	if err != nil {
484		return "", err
485	}
486	return string(b), nil
487}
488
489// String generates a human-readable representation of the current tree.
490// Alias of ToString. Present to implement the fmt.Stringer interface.
491func (t *Tree) String() string {
492	result, _ := t.ToTomlString()
493	return result
494}
495
496// ToMap recursively generates a representation of the tree using Go built-in structures.
497// The following types are used:
498//
499//	* bool
500//	* float64
501//	* int64
502//	* string
503//	* uint64
504//	* time.Time
505//	* map[string]interface{} (where interface{} is any of this list)
506//	* []interface{} (where interface{} is any of this list)
507func (t *Tree) ToMap() map[string]interface{} {
508	result := map[string]interface{}{}
509
510	for k, v := range t.values {
511		switch node := v.(type) {
512		case []*Tree:
513			var array []interface{}
514			for _, item := range node {
515				array = append(array, item.ToMap())
516			}
517			result[k] = array
518		case *Tree:
519			result[k] = node.ToMap()
520		case *tomlValue:
521			result[k] = tomlValueToGo(node.value)
522		}
523	}
524	return result
525}
526
527func tomlValueToGo(v interface{}) interface{} {
528	if tree, ok := v.(*Tree); ok {
529		return tree.ToMap()
530	}
531
532	rv := reflect.ValueOf(v)
533
534	if rv.Kind() != reflect.Slice {
535		return v
536	}
537	values := make([]interface{}, rv.Len())
538	for i := 0; i < rv.Len(); i++ {
539		item := rv.Index(i).Interface()
540		values[i] = tomlValueToGo(item)
541	}
542	return values
543}
544