Option to keep fields ordered when marshal struct (#266)

Adds a new `Order()` option to preserve order of struct fields when
marshaling.

marshal.go

@@ -54,10 +54,21 @@ var annotationDefault = annotation{
 	defaultValue: tagDefault,
 }
 
+type marshalOrder int
+
+// Orders the Encoder can write the fields to the output stream.
+const (
+	// Sort fields alphabetically.
+	OrderAlphabetical marshalOrder = iota + 1
+	// Preserve the order the fields are encountered. For example, the order of fields in
+	// a struct.
+	OrderPreserve
+)
+
 var timeType = reflect.TypeOf(time.Time{})
 var marshalerType = reflect.TypeOf(new(Marshaler)).Elem()
 
-// Check if the given marshall type maps to a Tree primitive
+// Check if the given marshal type maps to a Tree primitive
 func isPrimitive(mtype reflect.Type) bool {
 	switch mtype.Kind() {
 	case reflect.Ptr:
@@ -79,7 +90,7 @@ func isPrimitive(mtype reflect.Type) bool {
 	}
 }
 
-// Check if the given marshall type maps to a Tree slice
+// Check if the given marshal type maps to a Tree slice
 func isTreeSlice(mtype reflect.Type) bool {
 	switch mtype.Kind() {
 	case reflect.Slice:
@@ -89,7 +100,7 @@ func isTreeSlice(mtype reflect.Type) bool {
 	}
 }
 
-// Check if the given marshall type maps to a non-Tree slice
+// Check if the given marshal type maps to a non-Tree slice
 func isOtherSlice(mtype reflect.Type) bool {
 	switch mtype.Kind() {
 	case reflect.Ptr:
@@ -101,7 +112,7 @@ func isOtherSlice(mtype reflect.Type) bool {
 	}
 }
 
-// Check if the given marshall type maps to a Tree
+// Check if the given marshal type maps to a Tree
 func isTree(mtype reflect.Type) bool {
 	switch mtype.Kind() {
 	case reflect.Map:
@@ -159,6 +170,8 @@ Tree primitive types and corresponding marshal types:
   string     string, pointers to same
   bool       bool, pointers to same
   time.Time  time.Time{}, pointers to same
+
+For additional flexibility, use the Encoder API.
 */
 func Marshal(v interface{}) ([]byte, error) {
 	return NewEncoder(nil).marshal(v)
@@ -169,6 +182,9 @@ type Encoder struct {
 	w io.Writer
 	encOpts
 	annotation
+	line  int
+	col   int
+	order marshalOrder
 }
 
 // NewEncoder returns a new encoder that writes to w.
@@ -177,6 +193,9 @@ func NewEncoder(w io.Writer) *Encoder {
 		w:          w,
 		encOpts:    encOptsDefaults,
 		annotation: annotationDefault,
+		line:       0,
+		col:        1,
+		order:      OrderAlphabetical,
 	}
 }
 
@@ -222,6 +241,12 @@ func (e *Encoder) ArraysWithOneElementPerLine(v bool) *Encoder {
 	return e
 }
 
+// Order allows to change in which order fields will be written to the output stream.
+func (e *Encoder) Order(ord marshalOrder) *Encoder {
+	e.order = ord
+	return e
+}
+
 // SetTagName allows changing default tag "toml"
 func (e *Encoder) SetTagName(v string) *Encoder {
 	e.tag = v
@@ -269,17 +294,22 @@ func (e *Encoder) marshal(v interface{}) ([]byte, error) {
 	}
 
 	var buf bytes.Buffer
-	_, err = t.writeTo(&buf, "", "", 0, e.arraysOneElementPerLine)
+	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order)
 
 	return buf.Bytes(), err
 }
 
+// Create next tree with a position based on Encoder.line
+func (e *Encoder) nextTree() *Tree {
+	return newTreeWithPosition(Position{Line: e.line, Col: 1})
+}
+
 // Convert given marshal struct or map value to toml tree
 func (e *Encoder) valueToTree(mtype reflect.Type, mval reflect.Value) (*Tree, error) {
 	if mtype.Kind() == reflect.Ptr {
 		return e.valueToTree(mtype.Elem(), mval.Elem())
 	}
-	tval := newTree()
+	tval := e.nextTree()
 	switch mtype.Kind() {
 	case reflect.Struct:
 		for i := 0; i < mtype.NumField(); i++ {
@@ -347,6 +377,7 @@ func (e *Encoder) valueToOtherSlice(mtype reflect.Type, mval reflect.Value) (int
 
 // Convert given marshal value to toml value
 func (e *Encoder) valueToToml(mtype reflect.Type, mval reflect.Value) (interface{}, error) {
+	e.line++
 	if mtype.Kind() == reflect.Ptr {
 		return e.valueToToml(mtype.Elem(), mval.Elem())
 	}

marshal_OrderPreserve_test.toml

@@ -0,0 +1,38 @@
+title = "TOML Marshal Testing"
+
+[basic_lists]
+  floats = [12.3,45.6,78.9]
+  bools = [true,false,true]
+  dates = [1979-05-27T07:32:00Z,1980-05-27T07:32:00Z]
+  ints = [8001,8001,8002]
+  uints = [5002,5003]
+  strings = ["One","Two","Three"]
+
+[[subdocptrs]]
+  name = "Second"
+
+[basic_map]
+  one = "one"
+  two = "two"
+
+[subdoc]
+
+  [subdoc.second]
+    name = "Second"
+
+  [subdoc.first]
+    name = "First"
+
+[basic]
+  uint = 5001
+  bool = true
+  float = 123.4
+  int = 5000
+  string = "Bite me"
+  date = 1979-05-27T07:32:00Z
+
+[[subdoclist]]
+  name = "List.First"
+
+[[subdoclist]]
+  name = "List.Second"

marshal_test.go

@@ -12,10 +12,10 @@ import (
 )
 
 type basicMarshalTestStruct struct {
-	String     string                      `toml:"string"`
+	String     string                      `toml:"Zstring"`
-	StringList []string                    `toml:"strlist"`
+	StringList []string                    `toml:"Ystrlist"`
-	Sub        basicMarshalTestSubStruct   `toml:"subdoc"`
+	Sub        basicMarshalTestSubStruct   `toml:"Xsubdoc"`
-	SubList    []basicMarshalTestSubStruct `toml:"sublist"`
+	SubList    []basicMarshalTestSubStruct `toml:"Wsublist"`
 }
 
 type basicMarshalTestSubStruct struct {
@@ -29,16 +29,29 @@ var basicTestData = basicMarshalTestStruct{
 	SubList:    []basicMarshalTestSubStruct{{"Two"}, {"Three"}},
 }
 
-var basicTestToml = []byte(`string = "Hello"
+var basicTestToml = []byte(`Ystrlist = ["Howdy","Hey There"]
-strlist = ["Howdy","Hey There"]
+Zstring = "Hello"
 
-[subdoc]
+[[Wsublist]]
-  String2 = "One"
-
-[[sublist]]
   String2 = "Two"
 
-[[sublist]]
+[[Wsublist]]
+  String2 = "Three"
+
+[Xsubdoc]
+  String2 = "One"
+`)
+
+var basicTestTomlOrdered = []byte(`Zstring = "Hello"
+Ystrlist = ["Howdy","Hey There"]
+
+[Xsubdoc]
+  String2 = "One"
+
+[[Wsublist]]
+  String2 = "Two"
+
+[[Wsublist]]
   String2 = "Three"
 `)
 
@@ -53,6 +66,18 @@ func TestBasicMarshal(t *testing.T) {
 	}
 }
 
+func TestBasicMarshalOrdered(t *testing.T) {
+	var result bytes.Buffer
+	err := NewEncoder(&result).Order(OrderPreserve).Encode(basicTestData)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := basicTestTomlOrdered
+	if !bytes.Equal(result.Bytes(), expected) {
+		t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, result.Bytes())
+	}
+}
+
 func TestBasicMarshalWithPointer(t *testing.T) {
 	result, err := Marshal(&basicTestData)
 	if err != nil {
@@ -64,6 +89,18 @@ func TestBasicMarshalWithPointer(t *testing.T) {
 	}
 }
 
+func TestBasicMarshalOrderedWithPointer(t *testing.T) {
+	var result bytes.Buffer
+	err := NewEncoder(&result).Order(OrderPreserve).Encode(&basicTestData)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := basicTestTomlOrdered
+	if !bytes.Equal(result.Bytes(), expected) {
+		t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, result.Bytes())
+	}
+}
+
 func TestBasicUnmarshal(t *testing.T) {
 	result := basicMarshalTestStruct{}
 	err := Unmarshal(basicTestToml, &result)
@@ -78,39 +115,39 @@ func TestBasicUnmarshal(t *testing.T) {
 
 type testDoc struct {
 	Title       string            `toml:"title"`
-	Basics      testDocBasics     `toml:"basic"`
 	BasicLists  testDocBasicLists `toml:"basic_lists"`
+	SubDocPtrs  []*testSubDoc     `toml:"subdocptrs"`
 	BasicMap    map[string]string `toml:"basic_map"`
 	Subdocs     testDocSubs       `toml:"subdoc"`
+	Basics      testDocBasics     `toml:"basic"`
 	SubDocList  []testSubDoc      `toml:"subdoclist"`
-	SubDocPtrs  []*testSubDoc     `toml:"subdocptrs"`
 	err         int               `toml:"shouldntBeHere"`
 	unexported  int               `toml:"shouldntBeHere"`
 	Unexported2 int               `toml:"-"`
 }
 
 type testDocBasics struct {
+	Uint       uint      `toml:"uint"`
 	Bool       bool      `toml:"bool"`
-	Date       time.Time `toml:"date"`
 	Float      float32   `toml:"float"`
 	Int        int       `toml:"int"`
-	Uint       uint      `toml:"uint"`
 	String     *string   `toml:"string"`
+	Date       time.Time `toml:"date"`
 	unexported int       `toml:"shouldntBeHere"`
 }
 
 type testDocBasicLists struct {
+	Floats  []*float32  `toml:"floats"`
 	Bools   []bool      `toml:"bools"`
 	Dates   []time.Time `toml:"dates"`
-	Floats  []*float32  `toml:"floats"`
 	Ints    []int       `toml:"ints"`
-	Strings []string    `toml:"strings"`
 	UInts   []uint      `toml:"uints"`
+	Strings []string    `toml:"strings"`
 }
 
 type testDocSubs struct {
-	First  testSubDoc  `toml:"first"`
 	Second *testSubDoc `toml:"second"`
+	First  testSubDoc  `toml:"first"`
 }
 
 type testSubDoc struct {
@@ -174,6 +211,18 @@ func TestDocMarshal(t *testing.T) {
 	}
 }
 
+func TestDocMarshalOrdered(t *testing.T) {
+	var result bytes.Buffer
+	err := NewEncoder(&result).Order(OrderPreserve).Encode(docData)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected, _ := ioutil.ReadFile("marshal_OrderPreserve_test.toml")
+	if !bytes.Equal(result.Bytes(), expected) {
+		t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, result.Bytes())
+	}
+}
+
 func TestDocMarshalPointer(t *testing.T) {
 	result, err := Marshal(&docData)
 	if err != nil {

toml.go

@@ -27,9 +27,13 @@ type Tree struct {
 }
 
 func newTree() *Tree {
+	return newTreeWithPosition(Position{})
+}
+
+func newTreeWithPosition(pos Position) *Tree {
 	return &Tree{
 		values:   make(map[string]interface{}),
-		position: Position{},
+		position: pos,
 	}
 }
 
@@ -194,10 +198,10 @@ func (t *Tree) SetWithOptions(key string, opts SetOptions, value interface{}) {
 // formatting instructions to the key, that will be reused by Marshal().
 func (t *Tree) SetPathWithOptions(keys []string, opts SetOptions, value interface{}) {
 	subtree := t
-	for _, intermediateKey := range keys[:len(keys)-1] {
+	for i, intermediateKey := range keys[:len(keys)-1] {
 		nextTree, exists := subtree.values[intermediateKey]
 		if !exists {
-			nextTree = newTree()
+			nextTree = newTreeWithPosition(Position{Line: t.position.Line + i, Col: t.position.Col})
 			subtree.values[intermediateKey] = nextTree // add new element here
 		}
 		switch node := nextTree.(type) {
@@ -207,7 +211,7 @@ func (t *Tree) SetPathWithOptions(keys []string, opts SetOptions, value interfac
 			// go to most recent element
 			if len(node) == 0 {
 				// create element if it does not exist
-				subtree.values[intermediateKey] = append(node, newTree())
+				subtree.values[intermediateKey] = append(node, newTreeWithPosition(Position{Line: t.position.Line + i, Col: t.position.Col}))
 			}
 			subtree = node[len(node)-1]
 		}
@@ -225,7 +229,11 @@ func (t *Tree) SetPathWithOptions(keys []string, opts SetOptions, value interfac
 		v.comment = opts.Comment
 		toInsert = v
 	default:
-		toInsert = &tomlValue{value: value, comment: opts.Comment, commented: opts.Commented, multiline: opts.Multiline}
+		toInsert = &tomlValue{value: value,
+			comment:   opts.Comment,
+			commented: opts.Commented,
+			multiline: opts.Multiline,
+			position:  Position{Line: subtree.position.Line + len(subtree.values) + 1, Col: subtree.position.Col}}
 	}
 
 	subtree.values[keys[len(keys)-1]] = toInsert
@@ -254,42 +262,7 @@ func (t *Tree) SetPath(keys []string, value interface{}) {
 // SetPathWithComment is the same as SetPath, but allows you to provide comment
 // information to the key, that will be reused by Marshal().
 func (t *Tree) SetPathWithComment(keys []string, comment string, commented bool, value interface{}) {
-	subtree := t
+	t.SetPathWithOptions(keys, SetOptions{Comment: comment, Commented: commented}, value)
-	for _, intermediateKey := range keys[:len(keys)-1] {
-		nextTree, exists := subtree.values[intermediateKey]
-		if !exists {
-			nextTree = newTree()
-			subtree.values[intermediateKey] = nextTree // add new element here
-		}
-		switch node := nextTree.(type) {
-		case *Tree:
-			subtree = node
-		case []*Tree:
-			// go to most recent element
-			if len(node) == 0 {
-				// create element if it does not exist
-				subtree.values[intermediateKey] = append(node, newTree())
-			}
-			subtree = node[len(node)-1]
-		}
-	}
-
-	var toInsert interface{}
-
-	switch v := value.(type) {
-	case *Tree:
-		v.comment = comment
-		toInsert = value
-	case []*Tree:
-		toInsert = value
-	case *tomlValue:
-		v.comment = comment
-		toInsert = v
-	default:
-		toInsert = &tomlValue{value: value, comment: comment, commented: commented}
-	}
-
-	subtree.values[keys[len(keys)-1]] = toInsert
 }
 
 // Delete removes a key from the tree.
@@ -329,10 +302,10 @@ func (t *Tree) DeletePath(keys []string) error {
 // Returns nil on success, error object on failure
 func (t *Tree) createSubTree(keys []string, pos Position) error {
 	subtree := t
-	for _, intermediateKey := range keys {
+	for i, intermediateKey := range keys {
 		nextTree, exists := subtree.values[intermediateKey]
 		if !exists {
-			tree := newTree()
+			tree := newTreeWithPosition(Position{Line: t.position.Line + i, Col: t.position.Col})
 			tree.position = pos
 			subtree.values[intermediateKey] = tree
 			nextTree = tree

tomltree_write.go

@@ -12,6 +12,18 @@ import (
 	"time"
 )
 
+type valueComplexity int
+
+const (
+	valueSimple valueComplexity = iota + 1
+	valueComplex
+)
+
+type sortNode struct {
+	key        string
+	complexity valueComplexity
+}
+
 // Encodes a string to a TOML-compliant multi-line string value
 // This function is a clone of the existing encodeTomlString function, except that whitespace characters
 // are preserved. Quotation marks and backslashes are also not escaped.
@@ -153,59 +165,113 @@ func tomlValueStringRepresentation(v interface{}, indent string, arraysOneElemen
 	return "", fmt.Errorf("unsupported value type %T: %v", v, v)
 }
 
-func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool) (int64, error) {
+func getTreeArrayLine(trees []*Tree) (line int) {
-	simpleValuesKeys := make([]string, 0)
+	// get lowest line number that is not 0
-	complexValuesKeys := make([]string, 0)
+	for _, tv := range trees {
+		if tv.position.Line < line || line == 0 {
+			line = tv.position.Line
+		}
+	}
+	return
+}
+
+func sortByLines(t *Tree) (vals []sortNode) {
+	var (
+		line  int
+		lines []int
+		tv    *Tree
+		tom   *tomlValue
+		node  sortNode
+	)
+	vals = make([]sortNode, 0)
+	m := make(map[int]sortNode)
+
+	for k := range t.values {
+		v := t.values[k]
+		switch v.(type) {
+		case *Tree:
+			tv = v.(*Tree)
+			line = tv.position.Line
+			node = sortNode{key: k, complexity: valueComplex}
+		case []*Tree:
+			line = getTreeArrayLine(v.([]*Tree))
+			node = sortNode{key: k, complexity: valueComplex}
+		default:
+			tom = v.(*tomlValue)
+			line = tom.position.Line
+			node = sortNode{key: k, complexity: valueSimple}
+		}
+		lines = append(lines, line)
+		vals = append(vals, node)
+		m[line] = node
+	}
+	sort.Ints(lines)
+
+	for i, line := range lines {
+		vals[i] = m[line]
+	}
+
+	return vals
+}
+
+func sortAlphabetical(t *Tree) (vals []sortNode) {
+	var (
+		node     sortNode
+		simpVals []string
+		compVals []string
+	)
+	vals = make([]sortNode, 0)
+	m := make(map[string]sortNode)
 
 	for k := range t.values {
 		v := t.values[k]
 		switch v.(type) {
 		case *Tree, []*Tree:
-			complexValuesKeys = append(complexValuesKeys, k)
+			node = sortNode{key: k, complexity: valueComplex}
+			compVals = append(compVals, node.key)
 		default:
-			simpleValuesKeys = append(simpleValuesKeys, k)
+			node = sortNode{key: k, complexity: valueSimple}
+			simpVals = append(simpVals, node.key)
 		}
+		vals = append(vals, node)
+		m[node.key] = node
 	}
 
-	sort.Strings(simpleValuesKeys)
+	// Simples first to match previous implementation
-	sort.Strings(complexValuesKeys)
+	sort.Strings(simpVals)
-
+	i := 0
-	for _, k := range simpleValuesKeys {
+	for _, key := range simpVals {
-		v, ok := t.values[k].(*tomlValue)
+		vals[i] = m[key]
-		if !ok {
+		i++
-			return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
 	}
 
-		repr, err := tomlValueStringRepresentation(v, indent, arraysOneElementPerLine)
+	sort.Strings(compVals)
-		if err != nil {
+	for _, key := range compVals {
-			return bytesCount, err
+		vals[i] = m[key]
+		i++
 	}
 
-		if v.comment != "" {
+	return vals
-			comment := strings.Replace(v.comment, "\n", "\n"+indent+"#", -1)
+}
-			start := "# "
+
-			if strings.HasPrefix(comment, "#") {
+func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool) (int64, error) {
-				start = ""
+	return t.writeToOrdered(w, indent, keyspace, bytesCount, arraysOneElementPerLine, OrderAlphabetical)
-			}
+}
-			writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment, "\n")
+
-			bytesCount += int64(writtenBytesCountComment)
+func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord marshalOrder) (int64, error) {
-			if errc != nil {
+	orderedVals := make([]sortNode, 0)
-				return bytesCount, errc
+
-			}
+	switch ord {
+	case OrderPreserve:
+		orderedVals = sortByLines(t)
+	default:
+		orderedVals = sortAlphabetical(t)
 	}
 
-		var commented string
+	for _, node := range orderedVals {
-		if v.commented {
+		switch node.complexity {
-			commented = "# "
+		case valueComplex:
-		}
+			k := node.key
-		writtenBytesCount, err := writeStrings(w, indent, commented, k, " = ", repr, "\n")
-		bytesCount += int64(writtenBytesCount)
-		if err != nil {
-			return bytesCount, err
-		}
-	}
-
-	for _, k := range complexValuesKeys {
 			v := t.values[k]
 
 			combinedKey := k
@@ -241,7 +307,7 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, a
 				if err != nil {
 					return bytesCount, err
 				}
-			bytesCount, err = node.writeTo(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine)
+				bytesCount, err = node.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord)
 				if err != nil {
 					return bytesCount, err
 				}
@@ -253,12 +319,47 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, a
 						return bytesCount, err
 					}
 
-				bytesCount, err = subTree.writeTo(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine)
+					bytesCount, err = subTree.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord)
 					if err != nil {
 						return bytesCount, err
 					}
 				}
 			}
+		default: // Simple
+			k := node.key
+			v, ok := t.values[k].(*tomlValue)
+			if !ok {
+				return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
+			}
+
+			repr, err := tomlValueStringRepresentation(v, indent, arraysOneElementPerLine)
+			if err != nil {
+				return bytesCount, err
+			}
+
+			if v.comment != "" {
+				comment := strings.Replace(v.comment, "\n", "\n"+indent+"#", -1)
+				start := "# "
+				if strings.HasPrefix(comment, "#") {
+					start = ""
+				}
+				writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment, "\n")
+				bytesCount += int64(writtenBytesCountComment)
+				if errc != nil {
+					return bytesCount, errc
+				}
+			}
+
+			var commented string
+			if v.commented {
+				commented = "# "
+			}
+			writtenBytesCount, err := writeStrings(w, indent, commented, k, " = ", repr, "\n")
+			bytesCount += int64(writtenBytesCount)
+			if err != nil {
+				return bytesCount, err
+			}
+		}
 	}
 
 	return bytesCount, nil