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	// Prevent returning 0 for empty trees
230	line = int(^uint(0) >> 1)
231	// get lowest line number >= 0
232	for _, tv := range trees {
233		if tv.position.Line < line || line == 0 {
234			line = tv.position.Line
235		}
236	}
237	return
238}
239
240func sortByLines(t *Tree) (vals []sortNode) {
241	var (
242		line  int
243		lines []int
244		tv    *Tree
245		tom   *tomlValue
246		node  sortNode
247	)
248	vals = make([]sortNode, 0)
249	m := make(map[int]sortNode)
250
251	for k := range t.values {
252		v := t.values[k]
253		switch v.(type) {
254		case *Tree:
255			tv = v.(*Tree)
256			line = tv.position.Line
257			node = sortNode{key: k, complexity: valueComplex}
258		case []*Tree:
259			line = getTreeArrayLine(v.([]*Tree))
260			node = sortNode{key: k, complexity: valueComplex}
261		default:
262			tom = v.(*tomlValue)
263			line = tom.position.Line
264			node = sortNode{key: k, complexity: valueSimple}
265		}
266		lines = append(lines, line)
267		vals = append(vals, node)
268		m[line] = node
269	}
270	sort.Ints(lines)
271
272	for i, line := range lines {
273		vals[i] = m[line]
274	}
275
276	return vals
277}
278
279func sortAlphabetical(t *Tree) (vals []sortNode) {
280	var (
281		node     sortNode
282		simpVals []string
283		compVals []string
284	)
285	vals = make([]sortNode, 0)
286	m := make(map[string]sortNode)
287
288	for k := range t.values {
289		v := t.values[k]
290		switch v.(type) {
291		case *Tree, []*Tree:
292			node = sortNode{key: k, complexity: valueComplex}
293			compVals = append(compVals, node.key)
294		default:
295			node = sortNode{key: k, complexity: valueSimple}
296			simpVals = append(simpVals, node.key)
297		}
298		vals = append(vals, node)
299		m[node.key] = node
300	}
301
302	// Simples first to match previous implementation
303	sort.Strings(simpVals)
304	i := 0
305	for _, key := range simpVals {
306		vals[i] = m[key]
307		i++
308	}
309
310	sort.Strings(compVals)
311	for _, key := range compVals {
312		vals[i] = m[key]
313		i++
314	}
315
316	return vals
317}
318
319func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool) (int64, error) {
320	return t.writeToOrdered(w, indent, keyspace, bytesCount, arraysOneElementPerLine, OrderAlphabetical, "  ", false, false)
321}
322
323func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord MarshalOrder, indentString string, compactComments, parentCommented bool) (int64, error) {
324	var orderedVals []sortNode
325
326	switch ord {
327	case OrderPreserve:
328		orderedVals = sortByLines(t)
329	default:
330		orderedVals = sortAlphabetical(t)
331	}
332
333	for _, node := range orderedVals {
334		switch node.complexity {
335		case valueComplex:
336			k := node.key
337			v := t.values[k]
338
339			combinedKey := quoteKeyIfNeeded(k)
340			if keyspace != "" {
341				combinedKey = keyspace + "." + combinedKey
342			}
343
344			switch node := v.(type) {
345			// node has to be of those two types given how keys are sorted above
346			case *Tree:
347				tv, ok := t.values[k].(*Tree)
348				if !ok {
349					return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
350				}
351				if tv.comment != "" {
352					comment := strings.Replace(tv.comment, "\n", "\n"+indent+"#", -1)
353					start := "# "
354					if strings.HasPrefix(comment, "#") {
355						start = ""
356					}
357					writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment)
358					bytesCount += int64(writtenBytesCountComment)
359					if errc != nil {
360						return bytesCount, errc
361					}
362				}
363
364				var commented string
365				if parentCommented || t.commented || tv.commented {
366					commented = "# "
367				}
368				writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[", combinedKey, "]\n")
369				bytesCount += int64(writtenBytesCount)
370				if err != nil {
371					return bytesCount, err
372				}
373				bytesCount, err = node.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, compactComments, parentCommented || t.commented || tv.commented)
374				if err != nil {
375					return bytesCount, err
376				}
377			case []*Tree:
378				for _, subTree := range node {
379					var commented string
380					if parentCommented || t.commented || subTree.commented {
381						commented = "# "
382					}
383					writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[[", combinedKey, "]]\n")
384					bytesCount += int64(writtenBytesCount)
385					if err != nil {
386						return bytesCount, err
387					}
388
389					bytesCount, err = subTree.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, compactComments, parentCommented || t.commented || subTree.commented)
390					if err != nil {
391						return bytesCount, err
392					}
393				}
394			}
395		default: // Simple
396			k := node.key
397			v, ok := t.values[k].(*tomlValue)
398			if !ok {
399				return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
400			}
401
402			var commented string
403			if parentCommented || t.commented || v.commented {
404				commented = "# "
405			}
406			repr, err := tomlValueStringRepresentation(v, commented, indent, ord, arraysOneElementPerLine)
407			if err != nil {
408				return bytesCount, err
409			}
410
411			if v.comment != "" {
412				comment := strings.Replace(v.comment, "\n", "\n"+indent+"#", -1)
413				start := "# "
414				if strings.HasPrefix(comment, "#") {
415					start = ""
416				}
417				if !compactComments {
418					writtenBytesCountComment, errc := writeStrings(w, "\n")
419					bytesCount += int64(writtenBytesCountComment)
420					if errc != nil {
421						return bytesCount, errc
422					}
423				}
424				writtenBytesCountComment, errc := writeStrings(w, indent, start, comment, "\n")
425				bytesCount += int64(writtenBytesCountComment)
426				if errc != nil {
427					return bytesCount, errc
428				}
429			}
430
431			quotedKey := quoteKeyIfNeeded(k)
432			writtenBytesCount, err := writeStrings(w, indent, commented, quotedKey, " = ", repr, "\n")
433			bytesCount += int64(writtenBytesCount)
434			if err != nil {
435				return bytesCount, err
436			}
437		}
438	}
439
440	return bytesCount, nil
441}
442
443// quote a key if it does not fit the bare key format (A-Za-z0-9_-)
444// quoted keys use the same rules as strings
445func quoteKeyIfNeeded(k string) string {
446	// when encoding a map with the 'quoteMapKeys' option enabled, the tree will contain
447	// keys that have already been quoted.
448	// not an ideal situation, but good enough of a stop gap.
449	if len(k) >= 2 && k[0] == '"' && k[len(k)-1] == '"' {
450		return k
451	}
452	isBare := true
453	for _, r := range k {
454		if !isValidBareChar(r) {
455			isBare = false
456			break
457		}
458	}
459	if isBare {
460		return k
461	}
462	return quoteKey(k)
463}
464
465func quoteKey(k string) string {
466	return "\"" + encodeTomlString(k) + "\""
467}
468
469func writeStrings(w io.Writer, s ...string) (int, error) {
470	var n int
471	for i := range s {
472		b, err := io.WriteString(w, s[i])
473		n += b
474		if err != nil {
475			return n, err
476		}
477	}
478	return n, nil
479}
480
481// WriteTo encode the Tree as Toml and writes it to the writer w.
482// Returns the number of bytes written in case of success, or an error if anything happened.
483func (t *Tree) WriteTo(w io.Writer) (int64, error) {
484	return t.writeTo(w, "", "", 0, false)
485}
486
487// ToTomlString generates a human-readable representation of the current tree.
488// Output spans multiple lines, and is suitable for ingest by a TOML parser.
489// If the conversion cannot be performed, ToString returns a non-nil error.
490func (t *Tree) ToTomlString() (string, error) {
491	b, err := t.Marshal()
492	if err != nil {
493		return "", err
494	}
495	return string(b), nil
496}
497
498// String generates a human-readable representation of the current tree.
499// Alias of ToString. Present to implement the fmt.Stringer interface.
500func (t *Tree) String() string {
501	result, _ := t.ToTomlString()
502	return result
503}
504
505// ToMap recursively generates a representation of the tree using Go built-in structures.
506// The following types are used:
507//
508//	* bool
509//	* float64
510//	* int64
511//	* string
512//	* uint64
513//	* time.Time
514//	* map[string]interface{} (where interface{} is any of this list)
515//	* []interface{} (where interface{} is any of this list)
516func (t *Tree) ToMap() map[string]interface{} {
517	result := map[string]interface{}{}
518
519	for k, v := range t.values {
520		switch node := v.(type) {
521		case []*Tree:
522			var array []interface{}
523			for _, item := range node {
524				array = append(array, item.ToMap())
525			}
526			result[k] = array
527		case *Tree:
528			result[k] = node.ToMap()
529		case *tomlValue:
530			result[k] = tomlValueToGo(node.value)
531		}
532	}
533	return result
534}
535
536func tomlValueToGo(v interface{}) interface{} {
537	if tree, ok := v.(*Tree); ok {
538		return tree.ToMap()
539	}
540
541	rv := reflect.ValueOf(v)
542
543	if rv.Kind() != reflect.Slice {
544		return v
545	}
546	values := make([]interface{}, rv.Len())
547	for i := 0; i < rv.Len(); i++ {
548		item := rv.Index(i).Interface()
549		values[i] = tomlValueToGo(item)
550	}
551	return values
552}
553