fix(document): rebuild set nodes and write documents back round-trip

Assisted-by: GLM 5.3

docwrite.go

@@ -43,8 +43,8 @@ func (e *encoder) writeDocument(doc *Document) error {
 }
 
 // writeDocumentFooter writes the comment lines that follow the last
-// statement, each separated from it by a blank line, the shape the parser
-// reads them back from.
+// statement. The parser collects them wherever they sit after it, so the
+// writer needs no blank line of its own to have them read back.
 func (e *encoder) writeDocumentFooter(footer []string) {
 	for _, line := range footer {
 		e.buf.WriteString("# ")
@@ -53,8 +53,9 @@ func (e *encoder) writeDocumentFooter(footer []string) {
 	}
 }
 
-// writeTableEntries writes one table's entries in written order at the given
-// header path, nil for the document root, whose keys need no header.
+// writeTableEntries writes one table at the given header path, nil for the
+// document root, whose keys need no header: the blank line, the comments,
+// the header line with its trailing comment, then the body.
 func (e *encoder) writeTableEntries(t *Table, path []string) error {
 	if t == nil {
 		return nil
@@ -66,77 +67,222 @@ func (e *encoder) writeTableEntries(t *Table, path []string) error {
 		if err := e.writeKeyPath(path); err != nil {
 			return err
 		}
-		e.buf.WriteString("]\n")
+		e.buf.WriteString("]")
 		if tr := t.Trailing(); tr != "" {
 			e.buf.WriteString(" # ")
 			e.buf.WriteString(tr)
-			e.buf.WriteByte('\n')
+		}
+		e.buf.WriteByte('\n')
+	}
+	return e.writeTableBody(t, path)
+}
+
+// writeTableBody writes one table's entries: the value lines first, in
+// written order, then the header sections. In a valid document every line at
+// one level precedes the headers below it, so the split reorders nothing;
+// what it prevents is a table a dotted key introduced, which the parse nests
+// as a sub-table at the position of a line, from swallowing the lines that
+// follow it into its header.
+func (e *encoder) writeTableBody(t *Table, path []string) error {
+	for _, entry := range t.Entries() {
+		if err := e.checkCtx(); err != nil {
+			return err
+		}
+		if !e.isLineEntry(entry) {
+			continue
+		}
+		if err := e.writeLineEntry(entry, path); err != nil {
+			return err
 		}
 	}
 	for _, entry := range t.Entries() {
 		if err := e.checkCtx(); err != nil {
 			return err
 		}
-		if _, isTables := entry.Value().([]map[string]any); isTables && len(entry.Elements()) > 0 {
-			for i, el := range entry.Elements() {
-				elemPath := append(append([]string{}, path...), entry.Key())
-				e.writeBlankLine()
-				if i == 0 {
-					e.writeComments(entry.Comments())
-				}
-				e.buf.WriteString("[[")
-				if err := e.writeKeyPath(elemPath); err != nil {
-					return err
-				}
-				e.buf.WriteString("]]\n")
-				if err := e.writeTableEntries(el, elemPath); err != nil {
-					return err
-				}
-			}
-			continue
-		}
-		if child := entry.Table(); child != nil && !entry.Inline() {
-			headerPath := append(append([]string{}, path...), entry.Key())
-			if err := e.writeTableEntries(child, headerPath); err != nil {
+		if child := entry.Table(); child != nil && child.dotted && !entry.Inline() {
+			// A dotted table writes as lines above; its own header-form
+			// sub-tables are sections the document placed after those lines,
+			// so the section pass reaches through the dotted entry.
+			if err := e.writeDottedSections(child, append(append([]string{}, path...), entry.Key())); err != nil {
 				return err
 			}
 			continue
 		}
-		if err := e.writeDocumentEntry(entry, path); err != nil {
+		if e.isLineEntry(entry) {
+			continue
+		}
+		if err := e.writeSectionEntry(entry, path); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
+// writeDottedSections writes the header-form sub-tables of a dotted table:
+// the sections the document placed after the dotted lines, reached through
+// the dotted entry itself.
+func (e *encoder) writeDottedSections(t *Table, path []string) error {
+	for _, entry := range t.Entries() {
+		if err := e.checkCtx(); err != nil {
+			return err
+		}
+		if child := entry.Table(); child != nil && child.dotted && !entry.Inline() {
+			if err := e.writeDottedSections(child, append(append([]string{}, path...), entry.Key())); err != nil {
+				return err
+			}
+			continue
+		}
+		if e.isLineEntry(entry) {
+			continue
+		}
+		if err := e.writeSectionEntry(entry, path); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// writeSectionEntry writes one entry the line pass left behind: a table or
+// an array of tables under its header, at the path this level carries.
+func (e *encoder) writeSectionEntry(entry *Entry, path []string) error {
+	if _, isTables := entry.Value().([]map[string]any); isTables {
+		// An array of tables keeps its header form, one element per header
+		// with the element's own comments above it; the body that follows is
+		// the element's, with no header of its own to repeat.
+		elemPath := append(append([]string{}, path...), entry.Key())
+		for i, el := range entry.Elements() {
+			e.writeBlankLine()
+			if i == 0 {
+				e.writeComments(entry.Comments())
+			}
+			e.writeComments(el.Comments())
+			e.buf.WriteString("[[")
+			if err := e.writeKeyPath(elemPath); err != nil {
+				return err
+			}
+			e.buf.WriteString("]]")
+			if tr := el.Trailing(); tr != "" {
+				e.buf.WriteString(" # ")
+				e.buf.WriteString(tr)
+			}
+			e.buf.WriteByte('\n')
+			if err := e.writeTableBody(el, elemPath); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+	headerPath := append(append([]string{}, path...), entry.Key())
+	return e.writeTableEntries(entry.Table(), headerPath)
+}
+
+// isLineEntry reports whether an entry writes as one or more "key = value"
+// lines at its own level: a value, an inline table, or a table a dotted key
+// introduced, which goes back as dotted keys. An emptied array of tables
+// counts as one only so the line pass can drop it, the omission the value
+// encoder applies to an empty array of tables too.
+func (e *encoder) isLineEntry(entry *Entry) bool {
+	if child := entry.Table(); child != nil {
+		return entry.Inline() || child.dotted
+	}
+	if _, isTables := entry.Value().([]map[string]any); isTables {
+		return len(entry.Elements()) == 0
+	}
+	return true
+}
+
+// writeLineEntry writes one entry as lines at this level, and drops an
+// emptied array of tables, which has no TOML form.
+func (e *encoder) writeLineEntry(entry *Entry, path []string) error {
+	if child := entry.Table(); child != nil && !entry.Inline() {
+		return e.writeDottedTable(child, append(append([]string{}, path...), entry.Key()))
+	}
+	if _, isTables := entry.Value().([]map[string]any); isTables {
+		return nil
+	}
+	return e.writeDocumentEntry(entry)
+}
+
+// writeDottedTable writes a table a dotted key introduced as one dotted line
+// per leaf, in written order: `a.b = 1`. A sub-table the document added
+// under a header stays a section and is left to the section pass.
+func (e *encoder) writeDottedTable(t *Table, path []string) error {
+	for _, entry := range t.Entries() {
+		if err := e.checkCtx(); err != nil {
+			return err
+		}
+		if child := entry.Table(); child != nil && !entry.Inline() && !child.dotted {
+			continue
+		}
+		leafPath := append(append([]string{}, path...), entry.Key())
+		if child := entry.Table(); child != nil && !entry.Inline() {
+			if err := e.writeDottedTable(child, leafPath); err != nil {
+				return err
+			}
+			continue
+		}
+		e.writeComments(entry.Comments())
+		if err := e.writeKeyPath(leafPath); err != nil {
+			return err
+		}
+		e.buf.WriteString(" = ")
+		if err := e.writeEntryValueNodes(entry); err != nil {
+			return err
+		}
+		e.buf.WriteByte('\n')
+	}
+	return nil
+}
+
 // writeDocumentEntry writes one "key = value" line of a document, with the
 // comments the key carried. A value that is itself an inline table renders
 // inline from its node, in the written order.
-func (e *encoder) writeDocumentEntry(entry *Entry, path []string) error {
+func (e *encoder) writeDocumentEntry(entry *Entry) error {
 	e.writeComments(entry.Comments())
 	if err := e.writeKey(entry.Key()); err != nil {
 		return err
 	}
 	e.buf.WriteString(" = ")
+	if err := e.writeEntryValueNodes(entry); err != nil {
+		return err
+	}
+	e.buf.WriteByte('\n')
+	return nil
+}
+
+// writeEntryValueNodes writes the value of a document entry. An inline table
+// node keeps the written key order even inside a value array, where the
+// ordinary value writer would sort the keys.
+func (e *encoder) writeEntryValueNodes(entry *Entry) error {
 	if child := entry.Table(); child != nil {
 		if err := e.writeInlineTableNode(child); err != nil {
 			return err
 		}
-		if tr := entry.Trailing(); tr != "" {
-			e.buf.WriteString(" # ")
-			e.buf.WriteString(tr)
+	} else if arr, ok := entry.Value().([]any); ok {
+		elems := entry.Elements()
+		e.buf.WriteByte('[')
+		for i, item := range arr {
+			if i > 0 {
+				e.buf.WriteString(", ")
+			}
+			if i < len(elems) && elems[i] != nil {
+				if err := e.writeInlineTableNode(elems[i]); err != nil {
+					return err
+				}
+				continue
+			}
+			if err := e.writeValue(item); err != nil {
+				return err
+			}
 		}
-		e.buf.WriteByte('\n')
-		return nil
-	}
-	if err := e.writeValue(entry.Value()); err != nil {
+		e.buf.WriteByte(']')
+	} else if err := e.writeValue(entry.Value()); err != nil {
 		return err
 	}
 	if tr := entry.Trailing(); tr != "" {
 		e.buf.WriteString(" # ")
 		e.buf.WriteString(tr)
 	}
-	e.buf.WriteByte('\n')
 	return nil
 }