feat: parse into a Document that keeps order and comments

Assisted-by: DeepSeek V4.1 Flash

parser.go

@@ -44,6 +44,27 @@ type parser struct {
 	arrays  map[string]bool
 
 	currentPath []string
+
+	// wantDoc asks for the node tree the Document is built from; doc is that
+	// tree, and it stays nil when only the value tree is wanted. currentNode
+	// is the node of p.current; pending collects the comment lines since the
+	// last statement and trailing the comment on the statement's own line;
+	// lastEntry and lastTable name the statement those comments belong to;
+	// lastInline and lastArrayElems carry the nodes of the value parseValue has
+	// just produced.
+	wantDoc        bool
+	doc            *Table
+	currentNode    *Table
+	pending        []string
+	trailing       string
+	lastEntry      *Entry
+	lastTable      *Table
+	lastInline     *Table
+	lastArrayElems []*Table
+
+	// footer holds the comment lines that follow the last statement, which
+	// belong to the document rather than to any table or key.
+	footer []string
 }
 
 // maxNestingDepth bounds how deeply arrays and inline tables may nest when no
@@ -71,6 +92,10 @@ func (p *parser) parse() (map[string]any, error) {
 	p.dotted = map[string]bool{}
 	p.arrays = map[string]bool{}
 	p.currentPath = nil
+	if p.wantDoc {
+		p.doc = newTable(p.root)
+		p.currentNode = p.doc
+	}
 
 	for i := 0; ; i++ {
 		if i%ctxCheckInterval == 0 {
@@ -84,6 +109,7 @@ func (p *parser) parse() (map[string]any, error) {
 		if p.eof() {
 			break
 		}
+		p.lastEntry, p.lastTable = nil, nil
 		c := p.peek()
 		switch {
 		case c == '[':
@@ -98,10 +124,38 @@ func (p *parser) parse() (map[string]any, error) {
 		if err := p.expectLineEnd(); err != nil {
 			return nil, err
 		}
+		p.attachComments()
 	}
+	p.attachFooter()
 	return p.root, nil
 }
 
+// attachComments hands the collected comments to the statement just parsed:
+// the lines above it to its entry or table, the comment on its own line as the
+// trailing one.
+func (p *parser) attachComments() {
+	if p.doc != nil {
+		switch {
+		case p.lastEntry != nil:
+			p.lastEntry.comments = p.pending
+			p.lastEntry.trailing = p.trailing
+		case p.lastTable != nil:
+			p.lastTable.comments = p.pending
+			p.lastTable.trailing = p.trailing
+		}
+	}
+	p.pending, p.trailing = nil, ""
+}
+
+// attachFooter hands the comment lines that follow the last statement to the
+// document, which is where a comment block at the end of a file belongs.
+func (p *parser) attachFooter() {
+	if p.doc != nil && len(p.pending) > 0 {
+		p.footer = p.pending
+	}
+	p.pending = nil
+}
+
 // checkCtx returns ctx.Err() when the context has been cancelled, nil
 // otherwise. The call is a no-op when ctx is nil or the zero Background
 // context, both of which never cancel.
@@ -140,7 +194,7 @@ func (p *parser) parseTableHeader() error {
 	}
 
 	if array {
-		tbl, err := p.appendArrayTable(key)
+		tbl, elem, err := p.appendArrayTable(key)
 		if err != nil {
 			return err
 		}
@@ -150,6 +204,8 @@ func (p *parser) parseTableHeader() error {
 		p.arrays[pathKey(key)] = true
 		p.current = tbl
 		p.currentPath = key
+		p.currentNode = elem
+		p.lastTable = elem
 		return nil
 	}
 
@@ -159,69 +215,91 @@ func (p *parser) parseTableHeader() error {
 	}
 	p.headers[pk] = true
 
-	tbl, err := p.tableAt(key)
+	tbl, node, err := p.tableAt(key)
 	if err != nil {
 		return err
 	}
 	p.current = tbl
 	p.currentPath = key
+	p.currentNode = node
+	p.lastTable = node
 	return nil
 }
 
 // tableAt walks (creating intermediate tables) to the table named by key,
 // relative to the document root, rejecting any step into a frozen inline table.
-func (p *parser) tableAt(key []string) (map[string]any, error) {
+func (p *parser) tableAt(key []string) (map[string]any, *Table, error) {
 	cur := p.root
+	node := p.doc
 	path := make([]string, 0, len(key))
 	for _, k := range key {
 		path = append(path, k)
 		if p.frozen[pathKey(path)] {
-			return nil, p.errf("cannot extend inline table %q", strings.Join(path, "."))
+			return nil, nil, p.errf("cannot extend inline table %q", strings.Join(path, "."))
 		}
 		existing, ok := cur[k]
 		if !ok {
 			next := map[string]any{}
 			cur[k] = next
 			cur = next
+			if node != nil {
+				node = node.addTable(k, next)
+			}
 			continue
 		}
 		switch v := existing.(type) {
 		case map[string]any:
 			cur = v
+			if node != nil {
+				node = node.addTable(k, v)
+			}
 		case []map[string]any:
 			if len(v) == 0 {
-				return nil, p.errf("key %q is an empty array of tables", k)
+				return nil, nil, p.errf("key %q is an empty array of tables", k)
 			}
 			cur = v[len(v)-1]
+			if node != nil {
+				node = node.lastElement(k)
+			}
 		default:
-			return nil, p.errf("key %q is not a table", k)
+			return nil, nil, p.errf("key %q is not a table", k)
 		}
 	}
-	return cur, nil
+	return cur, node, nil
 }
 
-func (p *parser) appendArrayTable(key []string) (map[string]any, error) {
+func (p *parser) appendArrayTable(key []string) (map[string]any, *Table, error) {
 	parent := p.root
+	node := p.doc
 	path := make([]string, 0, len(key))
 	for _, k := range key[:len(key)-1] {
 		path = append(path, k)
 		if p.frozen[pathKey(path)] {
-			return nil, p.errf("cannot extend inline table %q", strings.Join(path, "."))
+			return nil, nil, p.errf("cannot extend inline table %q", strings.Join(path, "."))
 		}
 		existing, ok := parent[k]
 		if !ok {
 			next := map[string]any{}
 			parent[k] = next
 			parent = next
+			if node != nil {
+				node = node.addTable(k, next)
+			}
 			continue
 		}
 		switch v := existing.(type) {
 		case map[string]any:
 			parent = v
+			if node != nil {
+				node = node.addTable(k, v)
+			}
 		case []map[string]any:
 			parent = v[len(v)-1]
+			if node != nil {
+				node = node.lastElement(k)
+			}
 		default:
-			return nil, p.errf("key %q is not a table", k)
+			return nil, nil, p.errf("key %q is not a table", k)
 		}
 	}
 
@@ -233,9 +311,13 @@ func (p *parser) appendArrayTable(key []string) (map[string]any, error) {
 	case []map[string]any:
 		parent[leaf] = append(existing, tbl)
 	default:
-		return nil, p.errf("key %q is not an array of tables", leaf)
+		return nil, nil, p.errf("key %q is not an array of tables", leaf)
 	}
-	return tbl, nil
+	var elem *Table
+	if node != nil {
+		elem = node.addElement(leaf, tbl)
+	}
+	return tbl, elem, nil
 }
 
 // --- key/value -------------------------------------------------------------
@@ -262,6 +344,9 @@ func (p *parser) parseKeyValue() error {
 	// top-level statement reuses it for the leaf.
 	abs := make([]string, 0, len(p.currentPath)+len(key))
 	abs = append(abs, p.currentPath...)
+	// dests collects the map each dotted key descended into, which the node
+	// tree needs to build the matching tables around the value.
+	var dests []map[string]any
 	for _, k := range key[:len(key)-1] {
 		abs = append(abs, k)
 		if p.frozen[pathKey(abs)] {
@@ -276,6 +361,7 @@ func (p *parser) parseKeyValue() error {
 			next := map[string]any{}
 			dest[k] = next
 			dest = next
+			dests = append(dests, next)
 			continue
 		}
 		m, ok := existing.(map[string]any)
@@ -283,6 +369,7 @@ func (p *parser) parseKeyValue() error {
 			return p.errf("key %q is not a table", k)
 		}
 		dest = m
+		dests = append(dests, m)
 	}
 	leaf := key[len(key)-1]
 	abs = append(abs, leaf)
@@ -290,10 +377,47 @@ func (p *parser) parseKeyValue() error {
 		return p.errf("duplicate key %q", leaf)
 	}
 	dest[leaf] = val
+	if p.doc != nil {
+		node := p.currentNode
+		for i, k := range key[:len(key)-1] {
+			node = node.addTable(k, dests[i])
+		}
+		_, inline := val.(map[string]any)
+		entry := node.addValue(leaf, val, inline)
+		if inline {
+			entry.child = p.takeInline(val)
+		}
+		if nodes := p.takeArrayElems(val); nodes != nil {
+			entry.elements = nodes
+		}
+		p.lastEntry = entry
+	}
 	p.freezeInline(abs, val)
 	return nil
 }
 
+// takeInline returns the node of the inline table just parsed, when v is that
+// table's value, and clears it so a later value cannot pick it up.
+func (p *parser) takeInline(v any) *Table {
+	node := p.lastInline
+	p.lastInline = nil
+	if _, ok := v.(map[string]any); !ok {
+		return nil
+	}
+	return node
+}
+
+// takeArrayElems returns the element nodes of the array just parsed, when v is
+// that array's value, and clears them.
+func (p *parser) takeArrayElems(v any) []*Table {
+	nodes := p.lastArrayElems
+	p.lastArrayElems = nil
+	if _, ok := v.([]any); !ok {
+		return nil
+	}
+	return nodes
+}
+
 // freezeInline marks the path of an inline table (and any nested inline tables)
 // as immutable, so a later header or dotted key cannot extend it.
 func (p *parser) freezeInline(path []string, val any) {
@@ -383,6 +507,9 @@ func (p *parser) parseValue() (any, error) {
 	if p.eof() {
 		return nil, p.errf("expected a value")
 	}
+	// A container value leaves its node behind for the caller to pick up; a
+	// value that follows must not find the previous one.
+	p.lastInline, p.lastArrayElems = nil, nil
 	switch c := p.peek(); {
 	case c == '"':
 		return p.parseBasicString()
@@ -698,13 +825,23 @@ func (p *parser) readUnicode(n int) (rune, error) {
 
 // --- arrays and inline tables ---------------------------------------------
 
-func (p *parser) parseArray() (any, error) {
+func (p *parser) parseArray() (val any, err error) {
 	if err := p.enterNesting(); err != nil {
 		return nil, err
 	}
 	defer p.leaveNesting()
 	p.pos++ // '['
 	arr := []any{}
+	// elems carries the node of each element that is an inline table, so the
+	// caller can keep its key order; the entries are nil for other values.
+	var elems []*Table
+	if p.doc != nil {
+		defer func() {
+			if err == nil {
+				p.lastArrayElems = elems
+			}
+		}()
+	}
 	for {
 		if err := p.skipNestedSpace(); err != nil {
 			return nil, err
@@ -720,6 +857,9 @@ func (p *parser) parseArray() (any, error) {
 		if err != nil {
 			return nil, err
 		}
+		if p.doc != nil {
+			elems = append(elems, p.takeInline(v))
+		}
 		arr = append(arr, v)
 		if err := p.skipNestedSpace(); err != nil {
 			return nil, err
@@ -739,7 +879,7 @@ func (p *parser) parseArray() (any, error) {
 	}
 }
 
-func (p *parser) parseInlineTable() (any, error) {
+func (p *parser) parseInlineTable() (val any, err error) {
 	if err := p.enterNesting(); err != nil {
 		return nil, err
 	}
@@ -747,6 +887,18 @@ func (p *parser) parseInlineTable() (any, error) {
 	p.pos++ // '{'
 	tbl := map[string]any{}
 	assigned := map[string]bool{}
+	// The inline table is a node of its own, so the keys keep their order; the
+	// caller picks the node up when the table parses.
+	var node *Table
+	if p.doc != nil {
+		node = newTable(tbl)
+		node.inline = true
+		defer func() {
+			if err == nil {
+				p.lastInline = node
+			}
+		}()
+	}
 	// TOML 1.1 lets an inline table span lines: interior whitespace includes
 	// newlines and comments, and a trailing comma is allowed before the
 	// closing brace.
@@ -778,6 +930,7 @@ func (p *parser) parseInlineTable() (any, error) {
 
 		dest := tbl
 		path := make([]string, 0, len(key))
+		var dests []map[string]any
 		for _, k := range key[:len(key)-1] {
 			path = append(path, k)
 			if assigned[pathKey(path)] {
@@ -788,6 +941,7 @@ func (p *parser) parseInlineTable() (any, error) {
 				m := map[string]any{}
 				dest[k] = m
 				dest = m
+				dests = append(dests, m)
 				continue
 			}
 			m, isMap := existing.(map[string]any)
@@ -795,6 +949,7 @@ func (p *parser) parseInlineTable() (any, error) {
 				return nil, p.errf("key %q is already defined", k)
 			}
 			dest = m
+			dests = append(dests, m)
 		}
 		leaf := key[len(key)-1]
 		path = append(path, leaf)
@@ -803,6 +958,20 @@ func (p *parser) parseInlineTable() (any, error) {
 		}
 		dest[leaf] = val
 		assigned[pathKey(path)] = true
+		if node != nil {
+			child := node
+			for i, k := range key[:len(key)-1] {
+				child = child.addTable(k, dests[i])
+			}
+			_, inline := val.(map[string]any)
+			entry := child.addValue(leaf, val, inline)
+			if inline {
+				entry.child = p.takeInline(val)
+			}
+			if nodes := p.takeArrayElems(val); nodes != nil {
+				entry.elements = nodes
+			}
+		}
 
 		if err := p.skipNestedSpace(); err != nil {
 			return nil, err
@@ -897,7 +1066,9 @@ func (p *parser) skipNestedSpace() error {
 			p.line++
 			p.pos++
 		case '#':
-			if err := p.skipComment(); err != nil {
+			// A comment between values inside an array or an inline table is
+			// skipped; the Document does not carry those yet.
+			if _, err := p.skipComment(); err != nil {
 				return err
 			}
 		default:
@@ -921,9 +1092,11 @@ func (p *parser) skipBlank() error {
 			p.line++
 			p.pos++
 		case '#':
-			if err := p.skipComment(); err != nil {
+			line, err := p.skipComment()
+			if err != nil {
 				return err
 			}
+			p.pending = append(p.pending, line)
 		default:
 			return nil
 		}
@@ -931,27 +1104,34 @@ func (p *parser) skipBlank() error {
 	return nil
 }
 
-func (p *parser) skipComment() error {
+func (p *parser) skipComment() (string, error) {
 	p.pos++ // consume '#'
+	start := p.pos
 	for !p.eof() {
 		c := p.peek()
 		switch {
 		case c == '\n':
-			return nil
+			return commentText(string(p.src[start:p.pos])), nil
 		case c == '\r':
 			if p.pos+1 < len(p.src) && p.src[p.pos+1] == '\n' {
-				return nil
+				return commentText(string(p.src[start:p.pos])), nil
 			}
-			return p.errf("bare carriage return is not allowed")
+			return "", p.errf("bare carriage return is not allowed")
 		case c == '\t':
 			p.pos++
 		case c < 0x20 || c == 0x7f:
-			return p.errf("control character U+%04X is not allowed in a comment", c)
+			return "", p.errf("control character U+%04X is not allowed in a comment", c)
 		default:
 			p.pos++
 		}
 	}
-	return nil
+	return commentText(string(p.src[start:p.pos])), nil
+}
+
+// commentText drops the one space that usually follows the '#', so a line
+// stored in a Document reads as the comment itself.
+func commentText(s string) string {
+	return strings.TrimPrefix(s, " ")
 }
 
 // expectCRLF consumes a carriage return that must be immediately followed by a
@@ -972,9 +1152,13 @@ func (p *parser) expectLineEnd() error {
 		return nil
 	}
 	if p.peek() == '#' {
-		if err := p.skipComment(); err != nil {
+		line, err := p.skipComment()
+		if err != nil {
 			return err
 		}
+		if p.doc != nil {
+			p.trailing = line
+		}
 	}
 	if p.eof() {
 		return nil