v0.3.0 format a document, each comment with its owner; fewer allocations on the hot path

* Benchmarks for the parse, decode, encode and document paths

A command's payload from text to its struct, a settings file with comment
blocks, a scenario file of arrays of inline tables across lines, decoding
and encoding that scenario, and ParseDocument: the readings the changes
that follow are measured by.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ThctvSRwt5kV1pVMKtsdWc

* Fewer allocations on the parse, decode and encode paths

The parser's lexer drops what it never reads: a comment's text and a
number's literal, which the parser parses from the input bytes in place.
A string without escapes is one allocation, one with escapes is decoded
into a buffer sized to its source. The lexer scans ASCII bytewise, the
parser holds its lexer by value, keeps key paths on the stack, gathers an
array's elements in a reused buffer and allocates the array once at its
length, and holds the first few frozen inline tables in place; an inline
table that is an array element is not held, as no key can reach it.

Decode, Marshal and Unmapped read a struct's fields from one cache per
type instead of splitting every tag on every call. Decode reuses one key
and one element per map; Marshal sets a map's values into one array, and
writes keys, headers and numbers into its buffer without concatenating.

Tokens, positions, trees, error messages, encoded bytes and decoded
values are unchanged: a differential fuzz against the previous sources
ran 1.7M inputs over the lexer and parser and 9.5M over Marshal and
Decode without a difference.

benchstat, 8 runs each, Xeon 2.1 GHz:
  payload, parse and decode  3.31us -> 2.30us   31 -> 14 allocs  1.41 -> 1.09 KiB
  parse payload              2.40us -> 1.89us   24 -> 12 allocs
  parse settings            10.25us -> 6.98us   68 -> 32 allocs
  parse scenario             536us -> 391us   4959 -> 2970 allocs
  decode scenario            225us -> 110us   1863 -> 776 allocs
  marshal scenario           397us -> 252us   2326 -> 976 allocs
  ParseDocument settings    21.7us -> 16.5us   128 -> 85 allocs

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ThctvSRwt5kV1pVMKtsdWc

* Document: a comment goes with what fewer blank lines part it from

ParseDocument lays a document out as statements, runs of statements no
blank line parts, and the comment groups between them. A group belongs
to the side fewer blank lines part it from, the side below on a tie:
right above a statement it is that statement's own; right below a run, or
apart from it, it is the run's. Between two runs several groups split at
the widest gap. The file's start counts as no blank line and its end as
half of one, so a comment at the top with a blank line under it, or at
the end with one above it, is the file's, and one touching the last
statement stays with it.

Delete takes a statement with its own comments; a run going whole takes
its comments too and one blank run beside it, so no comment is left
orphaned and no gap doubles. Set puts a new table above the file's
closing comments and a first root key above the first table's comments,
a blank line apart. What v0.2.0 kept by position (the comment above a
dropped key, a dropped table's heading) now goes with its owner.

FuzzDocument checks that a delete leaves whole lines the document had,
in order, and that adding a key loses none; it ran 12.6M inputs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ThctvSRwt5kV1pVMKtsdWc

* Format: one layout for a document, every comment kept with its owner

Format, and Formatter for its options, writes a document spaced as
Marshal writes one: one space around = and after a comma, [1, 2] and
{ a = 1 }, LF line ends, a value's lines a level deeper than the line its
bracket opens on, and a trailing comma exactly when the closing bracket
has a line of its own. Keys, scalars and comments keep their spelling.

Comments stay with what they belong to by the rule Document reads them
with, and the result's blank lines encode it again: a run of blank lines
becomes one, or two below a comment group that belongs above it, so a
second Format reads every comment as the first did and writes the same
bytes. Each table is written once, at its first header: a header opening
it again gives its keys there, its comments above the key they touched or
else below what is there before them, and a dotted key reaching a table
that has a header moves into it with its own comments. Tables of an array
of tables never merge.

Options: Indent (up to 8 tabs or spaces, a tab by default), IndentTables
(a level per enclosing table whose header comes first), Align (= and
comments of consecutive keys).

Hardening: Format refuses what Parse refuses, a comment holding a
bidirectional embedding, override or isolate or a line break other than
LF (U+202A-202E, U+2066-2069, U+0085, U+2028, U+2029), which can show a
reader text other than what parses, and a result that would not parse to
the same values. Such a rune in a string or quoted key is written as its
\u escape, by Format and by the encoder alike. A line indents at most 16
levels and aligns within 32 runes of key and 80 of line, so the result
stays within a fixed multiple of its source; table paths sit in a trie,
so placing a key costs its length, not its square.

Marshal writes { a = 1 } and no blank first line, so Format(Marshal(v))
is Marshal(v), which FuzzRoundTrip now checks. FuzzFormat checks, with
the options off and on, that the result reads as the source, formats to
itself, keeps every comment once, holds no CR or trailing blank and stays
within the size bound; CI runs it for 10 s. All 29 TOML files in vif
format idempotently to the same values; vif.toml and keymap.toml are
unchanged byte for byte.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ThctvSRwt5kV1pVMKtsdWc

---------

Co-authored-by: Claude <noreply@anthropic.com>

.github/workflows/test.yml

@@ -17,3 +17,4 @@ jobs:
       - run: go test -run '^$' -fuzz '^FuzzRoundTrip$' -fuzztime=15s -parallel=2
       - run: go test -run '^$' -fuzz '^FuzzIntegerConversions$' -fuzztime=10s -parallel=2
       - run: go test -run '^$' -fuzz '^FuzzDocument$' -fuzztime=10s -parallel=2
+      - run: go test -run '^$' -fuzz '^FuzzFormat$' -fuzztime=10s -parallel=2

README.md

@@ -1,6 +1,6 @@
 # toml
 
-A deterministic TOML subset encoder/decoder using Go standard library.
+A deterministic TOML subset encoder, decoder and formatter using Go standard library.
 Requires Go 1.27.1.
 
 ```go
@@ -62,7 +62,8 @@ empty containers. Duplicate struct tag names error.
 
 Basic-string escapes emitted by the encoder (`\b`, `\f`, `\n`, `\r`, `\t`,
 quotes, backslashes, Unicode escapes) decode symmetrically; the TOML 1.1 `\e`
-and `\xHH` decode too. Invalid UTF-8 and unescaped control characters in
+and `\xHH` decode too. The encoder escapes the runes that can hide text, as
+`Format` does (below). Invalid UTF-8 and unescaped control characters in
 strings are rejected, as are control characters other than tab in comments and
 a carriage return outside CRLF. Encoding and typed decoding
 have recursion bounds (1000 levels), so cyclic traversals fail instead of
@@ -74,8 +75,11 @@ change it. Errors quote keys and cut long literals, and never repeat a string
 value. `Marshal` writes a table inline once its header path would pass 256
 bytes, so output stays linear in depth.
 
+`Marshal` writes an inline table as `{ a = 1 }` and a blank line above each
+header but one opening the document; its output is what `Format` writes for it.
 Comments and original formatting are not preserved by `Parse`/`Marshal`; a
-`Document` edits a file in place instead.
+`Document` edits a file in place, and `Format` rewrites one with its comments
+kept.
 
 ## Editing a document in place
 
@@ -89,19 +93,31 @@ data = doc.Bytes()
 ```
 
 `ParseDocument` parses as `Parse` does and keeps the bytes. An edit rewrites only
-the statements it names, so comments, blank lines, order and the spelling of
+the statements it names, with the comments that are theirs, so blank lines,
-every other value stay as written; the result is parsed again, and an edit it
+order, every other comment and the spelling of every other value stay as
-would not parse after is refused with the document unchanged.
+written; the result is parsed again, and an edit it would not parse after is
+refused with the document unchanged.
+
+A group of comment lines belongs to whichever side fewer blank lines part it
+from, the side below on a tie: the statement right under it, or the run of
+statements (lines no blank line parts) above or below it. The file's start and
+end count as sides, so a comment at the top with a blank line under it, or at
+the end with one above it, belongs to the file and stays, while one touching
+the last statement is that statement's run's. Between two runs, several groups
+split at the widest gap.
 
 - `Set` replaces the value of a key the document sets, keeping the rest of its
   line and its comment. A key it lacks goes after the last key of its table,
   spelled as that key is (`play.mouse` beside `play.hud`), else after the table's
-  header, else in a `[table]` appended at the end. A key under an array of
+  header, else in a `[table]` appended at the end, above the file's closing
-  tables, inside an inline table or array, or naming a table fails. The value is
+  comments. A root key with no other goes above the first table and its
-  encoded as `Marshal` writes one.
+  comments. A key under an array of tables, inside an inline table or array, or
+  naming a table fails. The value is encoded as `Marshal` writes one.
 - `Delete` removes the key and every key under it: each statement setting one and
-  each header naming one, in every table of an array of tables. It reports
+  each header naming one, in every table of an array of tables, each with its
-  whether it removed anything; a key inside an inline value fails.
+  comments. A run of statements going whole takes the comments that belong to
+  it and one blank run beside it. It reports whether it removed anything; a key
+  inside an inline value fails.
 - `Unmapped` names the keys no field of the destination takes, matched as `Decode`
   matches them (tags, then field names; unexported and `-` fields take nothing),
   each once and sorted. A map or interface destination takes every key under it;
@@ -111,6 +127,46 @@ would not parse after is refused with the document unchanged.
 
 A token's `Pos` and `End` are its byte offsets in the input.
 
+## Formatting
+
+```go
+out, err := toml.Format(data)
+out, err = toml.Formatter{Indent: "  ", IndentTables: true, Align: true}.Format(data)
+```
+
+`Format` writes a document spaced as `Marshal` writes one: one space around `=`
+and after a comma, `[1, 2]`, `{ a = 1 }`, `[a.b]`, LF line ends. A value's lines
+go a level deeper than the line its opening bracket is on, and an array or
+inline table has a trailing comma exactly when its closing bracket has a line of
+its own. Keys, scalars and comments keep their spelling; a run of blank lines
+becomes one.
+
+Every comment stays with what it belongs to, as a `Document` reads it, and the
+result's blank lines say so again: two blank lines follow a group of comments
+that belongs to what is above it, after a blank line. Each table is written
+once, at its first header. The keys of a header opening a table again go there,
+a blank line apart, the header's comments above them; a dotted key reaching a
+table that has a header moves into it, spelled from it, with its own comments.
+Tables of an array of tables stay as written.
+
+- `Indent` is one level, up to 8 tabs or spaces; a tab when empty.
+- `IndentTables` puts a table a level deeper for each table it is under whose
+  header comes before it: `[a.b]` after `[a]`.
+- `Align` lines up the `=` and the comments of keys on consecutive lines, with
+  spaces. A comment or blank line, a value of several lines, a key over 32 runes
+  or a line over 80 ends the run.
+
+A line indents no deeper than 16 levels, so a result stays within a fixed
+multiple of its source. `Format` refuses what `Parse` refuses, and a comment
+holding a bidirectional embedding, override or isolate (U+202A to U+202E, U+2066
+to U+2069) or a line break other than LF (U+0085, U+2028, U+2029): such a rune
+can show a reader text other than what parses, and a comment has no escape for
+it. In a string or quoted key it is written as its `\u` escape. A result that
+would not parse to the same values is refused.
+
+`Parse`, `Unmarshal` and `Decode` never format: they skip a comment without
+reading its text.
+
 ## Deliberate subset and deviations
 
 - Datetime syntax is unsupported. Use quoted strings and parse them in the

bench_test.go

@@ -0,0 +1,179 @@
+package toml
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+)
+
+// benchSettings is a settings file as people write one: comment blocks over
+// tables of scalars
+const benchSettings = `# settings: what a flag would otherwise say on every run. A flag given on the
+# command line wins; a key left out keeps the default written here.
+
+[paths]
+# A configuration root searched before the user root.
+root = ""
+# Where logs, journals and music go; empty is the platform default.
+log = ""
+journal = ""
+music = ""
+
+[audio]
+# Mixer period and output latency, in milliseconds.
+buffer_ms = 50
+
+[presentation]
+window = false
+maximized = false
+cell = 0
+
+# What the configuration menu changes in a run, written back so the next starts
+# as this one was left.
+[play]
+sfx = true
+music = true
+sfx_volume = 0.5
+auto_fire = "on"
+mouse = true
+hud = false
+`
+
+// benchPayload is an event payload a command or script carries
+const benchPayload = `_p = { entity = 7, x = 3, y = -2, pointer = true, name = "cursor" }`
+
+// scenarioDoc writes n states as a scenario file does: headers, arrays of inline
+// tables across lines, comments between elements
+func scenarioDoc(n int) []byte {
+	var b strings.Builder
+	b.WriteString("# A generated region: states and their transitions\n\n")
+	for i := range n {
+		fmt.Fprintf(&b, `# --- state %d ---
+
+[states.S%d]
+parent = "S%d"
+on_enter = [
+    { action = "EmitEvent", event = "EventSpawn%d" },
+]
+transitions = [
+    # the usual way on
+    { trigger = "EventDone", target = "S%d", guard = "StateTimeExceeds", guard_args = { ms = %d } },
+    { trigger = "Tick", target = "S%d", weight = 0.25 },
+]
+
+`, i, i, i/2, i, i+1, 100+i, i)
+	}
+	return []byte(b.String())
+}
+
+type benchTransition struct {
+	Trigger   string         `toml:"trigger"`
+	Target    string         `toml:"target"`
+	Guard     string         `toml:"guard,omitempty"`
+	GuardArgs map[string]any `toml:"guard_args,omitempty"`
+	Weight    float64        `toml:"weight,omitempty"`
+}
+
+type benchState struct {
+	Parent      string            `toml:"parent"`
+	OnEnter     []map[string]any  `toml:"on_enter,omitempty"`
+	Transitions []benchTransition `toml:"transitions"`
+}
+
+type benchScenario struct {
+	States map[string]benchState `toml:"states"`
+}
+
+type benchEvent struct {
+	Entity  int64  `toml:"entity"`
+	X       int    `toml:"x"`
+	Y       int    `toml:"y"`
+	Pointer bool   `toml:"pointer"`
+	Name    string `toml:"name"`
+}
+
+func BenchmarkParse(b *testing.B) {
+	for _, c := range []struct {
+		name string
+		data []byte
+	}{
+		{"payload", []byte(benchPayload)},
+		{"settings", []byte(benchSettings)},
+		{"scenario", scenarioDoc(64)},
+	} {
+		b.Run(c.name, func(b *testing.B) {
+			b.SetBytes(int64(len(c.data)))
+			b.ReportAllocs()
+			for b.Loop() {
+				if _, err := NewParser(c.data).Parse(); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+// BenchmarkPayload is a command's payload from text to its typed struct
+func BenchmarkPayload(b *testing.B) {
+	data := []byte(benchPayload)
+	b.ReportAllocs()
+	for b.Loop() {
+		tree, err := NewParser(data).Parse()
+		if err != nil {
+			b.Fatal(err)
+		}
+		var ev benchEvent
+		if err := Decode(tree["_p"], &ev); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkDecode(b *testing.B) {
+	tree, err := NewParser(scenarioDoc(64)).Parse()
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		var s benchScenario
+		if err := Decode(tree, &s); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkMarshal(b *testing.B) {
+	var s benchScenario
+	if err := Unmarshal(scenarioDoc(64), &s); err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Marshal(s); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkFormat(b *testing.B) {
+	data := scenarioDoc(64)
+	b.SetBytes(int64(len(data)))
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Format(data); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkParseDocument(b *testing.B) {
+	data := []byte(benchSettings)
+	b.SetBytes(int64(len(data)))
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := ParseDocument(data); err != nil {
+			b.Fatal(err)
+		}
+	}
+}

decode.go

@@ -1,10 +1,12 @@
 package toml
 
 import (
+	"cmp"
 	"fmt"
 	"math"
 	"reflect"
 	"strings"
+	"sync"
 )
 
 // Unmarshal parses TOML data and stores the result in the value pointed to by v.
@@ -132,13 +134,16 @@ func decodeAt(data any, val reflect.Value, depth int) error {
 			return fmt.Errorf("expected map, got %T", data)
 		}
 		newMap := reflect.MakeMapWithSize(val.Type(), len(dataMap))
-		elemType := val.Type().Elem()
+		// One key and one element, zeroed per entry, which SetMapIndex copies
+		key := reflect.New(val.Type().Key()).Elem()
+		elem := reflect.New(val.Type().Elem()).Elem()
 		for k, vData := range dataMap {
-			newVal := reflect.New(elemType).Elem()
+			elem.SetZero()
-			if err := decodeAt(vData, newVal, depth+1); err != nil {
+			if err := decodeAt(vData, elem, depth+1); err != nil {
 				return atPath("map key "+brief(k), err)
 			}
-			newMap.SetMapIndex(reflect.ValueOf(k).Convert(val.Type().Key()), newVal)
+			key.SetString(k)
+			newMap.SetMapIndex(key, elem)
 		}
 		val.Set(newMap)
 
@@ -207,39 +212,50 @@ func decodeAt(data any, val reflect.Value, depth int) error {
 }
 
 func decodeStruct(data map[string]any, val reflect.Value, depth int) error {
-	typ := val.Type()
+	for _, f := range fieldsOf(val.Type()) {
-
-	for i := 0; i < val.NumField(); i++ {
-		field := val.Field(i)
-		fieldType := typ.Field(i)
-
-		// Safety check: skip unexported fields that cannot be set
-		if !field.CanSet() {
-			continue
-		}
-
-		// Determine key name
-		key := fieldType.Name
-		if tag := fieldType.Tag.Get("toml"); tag != "" {
-			parts := strings.Split(tag, ",")
-			if parts[0] == "-" {
-				continue
-			}
-			if parts[0] != "" {
-				key = parts[0]
-			}
-		}
-
 		// Look up in data map (case sensitive)
-		if vData, ok := data[key]; ok {
+		if vData, ok := data[f.key]; ok {
-			if err := decodeAt(vData, field, depth+1); err != nil {
+			if err := decodeAt(vData, val.Field(f.index), depth+1); err != nil {
-				return atPath(typ.Name()+"."+fieldType.Name, err)
+				return atPath(val.Type().Name()+"."+f.name, err)
 			}
 		}
 	}
 	return nil
 }
 
+// field is a struct field a document can set: exported and not tagged "-"
+type field struct {
+	index     int
+	key       string // the tag's name, else the Go name
+	name      string // the Go name
+	omitempty bool
+}
+
+var fieldCache sync.Map // reflect.Type -> []field
+
+// fieldsOf lists the fields of struct type t a document names, in declaration
+// order; read once per type, as the tags never change
+func fieldsOf(t reflect.Type) []field {
+	if f, ok := fieldCache.Load(t); ok {
+		return f.([]field)
+	}
+	out := make([]field, 0, t.NumField())
+	for i := range t.NumField() {
+		f := t.Field(i)
+		name, opts, _ := strings.Cut(f.Tag.Get("toml"), ",")
+		if !f.IsExported() || name == "-" {
+			continue
+		}
+		x := field{index: i, key: cmp.Or(name, f.Name), name: f.Name}
+		for o := range strings.SplitSeq(opts, ",") {
+			x.omitempty = x.omitempty || o == "omitempty"
+		}
+		out = append(out, x)
+	}
+	f, _ := fieldCache.LoadOrStore(t, out)
+	return f.([]field)
+}
+
 // Numeric conversions check before conversion: rounded float64(MaxInt64) is
 // 2^63, so the upper bound must be exclusive. No float intermediary for integers.
 func toInt64(v any) (int64, bool) {

document.go

@@ -2,7 +2,6 @@ package toml
 
 import (
 	"bytes"
-	"cmp"
 	"fmt"
 	"reflect"
 	"slices"
@@ -31,23 +30,35 @@ func (k Key) String() string {
 func (k Key) under(p Key) bool { return len(k) >= len(p) && slices.Equal(k[:len(p)], p) }
 
 // Document is a parsed document that keeps its bytes. An edit rewrites only the
-// statements it names, so comments, blank lines, order and the spelling of every
+// statements it names, with the comments that belong to them, so blank lines,
-// other value stay as written, and the result is parsed again before it stands.
+// order, every other comment and the spelling of every other value stay as
+// written, and the result is parsed again before it stands.
 type Document struct {
 	data  []byte
 	tree  map[string]any
 	stmts []statement
+	units []unit
+	pre   int // where the comments heading the file end, 0 when none
+	foot  int // where those closing it start, -1 when none
 }
 
 // statement is a table header or a key and its value, as the lines it spans: from
 // the start of its line through its newline, a comment after it on that line
-// included. A key is its path relative to scope, the header it sits under.
+// included. lead is where the comment lines right above it, which are its own,
+// start. A key is its path relative to scope, the header it sits under.
 type statement struct {
-	start, end int
+	start, end, lead int
-	scope, key Key
+	scope, key       Key
-	header     bool
+	header           bool
-	array      bool   // an [[array]] header
+	array            bool   // an [[array]] header
-	value      [2]int // a key's value bytes
+	value            [2]int // a key's value bytes
+}
+
+// unit is a run of statements no blank line parts, from the first comment line
+// that belongs to the run through the last, whole lines: stmts[first:last]
+type unit struct {
+	start, end  int
+	first, last int
 }
 
 func (s statement) path() Key {
@@ -64,12 +75,94 @@ func ParseDocument(data []byte) (*Document, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &Document{data: data, tree: tree, stmts: statements(data)}, nil
+	d := layout(data)
+	d.tree = tree
+	return &d, nil
 }
 
-// statements splits a document the parser accepted into its statements
+// item is a statement, or a group of comment lines no blank line parts, with the
-func statements(data []byte) []statement {
+// blank lines between it and the item before it or the file's start
-	lx := NewLexer(data)
+type item struct {
+	start, end int
+	stmt       int // its index in stmts, -1 for comments
+	gap        int
+}
+
+// layout splits a document the parser accepted into statements and units and
+// gives each comment group to the side fewer blank lines part it from, below on
+// a tie
+func layout(data []byte) Document {
+	d := Document{data: data, foot: -1}
+	items := d.scan()
+	para := make([]int, len(items)) // a statement's unit
+	open := false                   // no blank line since the last unit's last item
+	for k, it := range items {
+		open = open && it.gap == 0
+		if it.stmt < 0 {
+			continue
+		}
+		if !open {
+			d.units = append(d.units, unit{start: it.start, first: it.stmt})
+		}
+		para[k], open = len(d.units)-1, true
+		d.units[para[k]].end, d.units[para[k]].last = it.end, it.stmt+1
+	}
+	for i := 0; i < len(items); {
+		if items[i].stmt >= 0 {
+			i++
+			continue
+		}
+		j := i
+		for j < len(items) && items[j].stmt < 0 {
+			j++
+		}
+		split := placeComments(items, i, j)
+		for _, c := range items[i:split] {
+			if i == 0 {
+				d.pre = c.end
+			} else {
+				d.units[para[i-1]].end = max(d.units[para[i-1]].end, c.end)
+			}
+		}
+		for _, c := range items[split:j] {
+			if j < len(items) {
+				d.units[para[j]].start = min(d.units[para[j]].start, c.start)
+			} else if d.foot < 0 {
+				d.foot = c.start
+			}
+		}
+		if split < j && j < len(items) && items[j].gap == 0 {
+			d.stmts[items[j].stmt].lead = items[j-1].start
+		}
+		i = j
+	}
+	return d
+}
+
+// placeComments splits the comment groups items[i:j] at the widest run of blank
+// lines around them, the first of equals: those above it go with what is above,
+// the rest with what is below. The file's start counts as no blank line and its
+// end as half of one, so a comment touching a statement stays with it.
+func placeComments(items []item, i, j int) int {
+	split, widest := i, -1
+	for k := i; k <= j; k++ {
+		w := 1 // the file's end
+		switch {
+		case k == 0:
+			w = 0
+		case k < len(items):
+			w = 2 * items[k].gap
+		}
+		if w > widest {
+			split, widest = k, w
+		}
+	}
+	return split
+}
+
+// scan lexes data into its items, appending d.stmts
+func (d *Document) scan() (items []item) {
+	lx := Lexer{input: d.data, line: 1, lean: true}
 	t := lx.NextToken()
 	key := func() Key {
 		var k Key
@@ -93,12 +186,26 @@ func statements(data []byte) []statement {
 		return end
 	}
 	var (
-		out   []statement
 		scope Key
+		blank int
 	)
 	for t.Type != TokenEOF {
-		s := statement{start: bytes.LastIndexByte(data[:t.Pos], '\n') + 1, scope: scope}
+		start := bytes.LastIndexByte(d.data[:t.Pos], '\n') + 1
+		s := statement{start: start, lead: start, scope: scope}
 		switch t.Type {
+		case TokenNewline:
+			blank++
+			t = lx.NextToken()
+			continue
+		case TokenComment:
+			end := lineEnd()
+			if n := len(items); n > 0 && items[n-1].stmt < 0 && blank == 0 {
+				items[n-1].end = end
+			} else {
+				items = append(items, item{start: start, end: end, stmt: -1, gap: blank})
+			}
+			blank = 0
+			continue
 		case TokenLBracket:
 			s.header, t = true, lx.NextToken()
 			if t.Type == TokenLBracket {
@@ -109,7 +216,7 @@ func statements(data []byte) []statement {
 				t = lx.NextToken()
 			}
 			s.scope, scope = nil, s.key
-		case TokenIdent, TokenString:
+		default: // a key
 			s.key = key()
 			t = lx.NextToken() // past =
 			s.value[0] = t.Pos
@@ -127,14 +234,13 @@ func statements(data []byte) []statement {
 					break
 				}
 			}
-		default: // a blank or comment line
-			t = lx.NextToken()
-			continue
 		}
 		s.end = lineEnd()
-		out = append(out, s)
+		items = append(items, item{start: s.start, end: s.end, stmt: len(d.stmts), gap: blank})
+		d.stmts = append(d.stmts, s)
+		blank = 0
 	}
-	return out
+	return items
 }
 
 // Bytes is the document as it now reads; the caller must not change it
@@ -174,13 +280,9 @@ func unmapped(data any, t reflect.Type, at Key, seen map[string]Key, depth int)
 		if !ok {
 			return
 		}
-		fields := make(map[string]reflect.Type, t.NumField())
+		fields := make(map[string]reflect.Type)
-		for i := range t.NumField() {
+		for _, f := range fieldsOf(t) {
-			f := t.Field(i)
+			fields[f.key] = t.Field(f.index).Type
-			name, _, _ := strings.Cut(f.Tag.Get("toml"), ",")
-			if f.IsExported() && name != "-" {
-				fields[cmp.Or(name, f.Name)] = f.Type
-			}
 		}
 		for k, v := range m {
 			key := append(slices.Clip(at), k)
@@ -210,25 +312,55 @@ func unmapped(data any, t reflect.Type, at Key, seen map[string]Key, depth int)
 	}
 }
 
-// Delete removes key and every key under it: each statement setting one and each
+// Delete removes key and every key under it, each statement setting one and each
-// table header naming one, in every table of an array of tables, every other line
+// header naming one, in every table of an array of tables, with the comments that
-// kept. It reports whether anything was removed; a key inside an inline value, a
+// are theirs; a run of statements going whole takes its comments and one blank run
-// table or array written as a value, cannot be edited and fails.
+// beside it. It reports whether anything was removed; a key inside an inline value,
+// a table or array written as a value, cannot be edited and fails.
 func (d *Document) Delete(key Key) (bool, error) {
-	var out []byte
+	gone := make([]bool, len(d.stmts))
-	at := 0
+	some := false
-	for _, s := range d.stmts {
+	for i, s := range d.stmts {
 		p := s.path()
 		if !s.header && len(p) < len(key) && key.under(p) && d.inline(s) {
 			return false, fmt.Errorf("key %s is inside the value of %s", key, p)
 		}
-		if p.under(key) {
+		gone[i] = p.under(key)
-			out, at = append(out, d.data[at:s.start]...), s.end
+		some = some || gone[i]
-		}
 	}
-	if at == 0 {
+	if !some {
 		return false, nil
 	}
+	whole := func(u unit) bool { return !slices.Contains(gone[u.first:u.last], false) }
+	var out []byte
+	at := 0
+	cut := func(from, to int) { out, at = append(out, d.data[at:from]...), to }
+	for k := 0; k < len(d.units); k++ {
+		u := d.units[k]
+		if !whole(u) {
+			for i := u.first; i < u.last; i++ {
+				if gone[i] {
+					cut(d.stmts[i].lead, d.stmts[i].end)
+				}
+			}
+			continue
+		}
+		j := k
+		for j+1 < len(d.units) && whole(d.units[j+1]) {
+			j++
+		}
+		switch {
+		case j+1 < len(d.units):
+			cut(u.start, d.units[j+1].start)
+		case d.foot >= 0:
+			cut(u.start, d.foot)
+		case k > 0:
+			cut(d.units[k-1].end, len(d.data))
+		default:
+			cut(d.pre, len(d.data))
+		}
+		k = j
+	}
 	return true, d.replace(append(out, d.data[at:]...))
 }
 
@@ -238,11 +370,11 @@ func (d *Document) inline(s statement) bool {
 	return c == '{' || c == '['
 }
 
-// Set gives key the value v, encoded as Marshal writes a value. A key the document
+// Set gives key the value v, encoded as Marshal writes one. A key the document sets
-// sets has its value replaced, the rest of its line kept; a key it lacks goes after
+// has its value replaced, the rest of its line kept; one it lacks goes after the
-// the last key of its table, else after the table's header, else in a table
+// last key of its table, else after its header, else in a table appended above the
-// appended to the document. A key under an array of tables or inside an inline
+// file's closing comments, a first root key above the first table and its comments.
-// value fails, and so does a value the document then would not parse with.
+// A key under an array of tables or an inline value fails, as does an unparsable edit.
 func (d *Document) Set(key Key, v any) error {
 	if len(key) == 0 || v == nil {
 		return fmt.Errorf("a key and a value are needed")
@@ -273,21 +405,22 @@ func (d *Document) Set(key Key, v any) error {
 			}
 		}
 	}
-	head := ""
+	head, foot := "", "\n"
 	switch {
 	case after >= 0:
-	case len(parent) == 0: // a root key goes before the first header
+	case len(parent) == 0 && len(d.units) > 0: // a blank line keeps the table's comments its own
+		after, foot = d.units[0].start, "\n\n"
+	case len(parent) == 0:
 		after = len(d.data)
-		if i := slices.IndexFunc(d.stmts, func(s statement) bool { return s.header }); i >= 0 {
+	case d.foot >= 0:
-			after = d.stmts[i].start
+		after, scope, head, foot = d.foot, parent, "["+parent.String()+"]\n", "\n\n"
-		}
 	default:
 		after, scope, head = len(d.data), parent, "["+parent.String()+"]\n"
 		if after > 0 {
 			head = "\n" + head
 		}
 	}
-	line := slices.Concat([]byte(head+key[len(scope):].String()+" = "), value, []byte("\n"))
+	line := slices.Concat([]byte(head+key[len(scope):].String()+" = "), value, []byte(foot))
 	if after > 0 && d.data[after-1] != '\n' {
 		line = slices.Concat([]byte("\n"), line)
 	}

document_test.go

@@ -3,6 +3,7 @@ package toml
 import (
 	"fmt"
 	"slices"
+	"strings"
 	"testing"
 )
 
@@ -33,7 +34,7 @@ buffer_ms = 20
 // TestDocumentEditsOnlyWhatItNames: a set value keeps the rest of its line, a
 // missing key goes after its table's last key in that key's spelling, a missing
 // table is appended, a delete takes a key's lines or a table's header with every
-// key under it, and every other byte, comments and blank lines, stays.
+// key under it and the comments that are theirs, and every other byte stays.
 func TestDocumentEditsOnlyWhatItNames(t *testing.T) {
 	d, err := ParseDocument([]byte(settingsDoc))
 	if err != nil {
@@ -62,8 +63,6 @@ log = "logs" # beside the file
 mouse = false # keys
 sfx = false
 
-# about later
-
 [audio]
 buffer_ms = 20
 "buffer ms" = 30
@@ -110,6 +109,65 @@ func TestSetSpellsAKeyAsItsTableDoes(t *testing.T) {
 	}
 }
 
+// TestCommentsGoWithTheNearerSide: a comment group belongs to what fewer blank
+// lines part it from, below on a tie; the file's start and end count as sides, so
+// its opening and closing comments stay. A key takes its own comments; a run of
+// statements going whole takes its comments and one blank run beside it.
+func TestCommentsGoWithTheNearerSide(t *testing.T) {
+	for _, c := range []struct {
+		name, doc string
+		key       Key
+		want      string
+	}{
+		{"right above", "a = 1\n# about b\nb = 2\nc = 3\n", Key{"b"}, "a = 1\nc = 3\n"},
+		{"right below, run stays", "a = 1\nb = 2\n# after\n\n[t]\nk = 1\n", Key{"b"}, "a = 1\n# after\n\n[t]\nk = 1\n"},
+		{"right below, run goes", "[a]\nx = 1\n# end of a\n\n[b]\ny = 2\n", Key{"a"}, "[b]\ny = 2\n"},
+		{"tie goes below", "[a]\nx = 1\n\n# about b\n\n[b]\ny = 2\n\n[c]\nz = 3\n", Key{"b"}, "[a]\nx = 1\n\n[c]\nz = 3\n"},
+		{"nearer above", "[a]\nx = 1\n\n# about a\n\n\n[b]\ny = 2\n", Key{"b"}, "[a]\nx = 1\n\n# about a\n"},
+		{"widest gap splits", "[a]\nx = 1\n# a1\n\n# a2\n\n\n# b1\n\n# b2\n[b]\ny = 2\n", Key{"b"}, "[a]\nx = 1\n# a1\n\n# a2\n"},
+		{"file's opening", "# file\n\n[a]\nx = 1\n", Key{"a"}, "# file\n"},
+		{"file's closing", "[a]\nx = 1\n\n[b]\ny = 2\n\n# end\n", Key{"b"}, "[a]\nx = 1\n\n# end\n"},
+		{"touching the end", "[a]\nx = 1\n\n[b]\ny = 2\n# y's\n", Key{"b"}, "[a]\nx = 1\n"},
+		{"one blank run left", "a = 1\n\n[b]\nx = 1\n\n\n[c]\ny = 1\n", Key{"b"}, "a = 1\n\n[c]\ny = 1\n"},
+	} {
+		d, err := ParseDocument([]byte(c.doc))
+		if err != nil {
+			t.Fatalf("%s: %v", c.name, err)
+		}
+		if err := del(d, c.key); err != nil {
+			t.Fatalf("%s: %v", c.name, err)
+		}
+		if got := string(d.Bytes()); got != c.want {
+			t.Errorf("%s:\n%s\nwant:\n%s", c.name, got, c.want)
+		}
+	}
+}
+
+// TestSetLeavesCommentsTheirOwners: a root key added above the first table keeps
+// a blank line from that table's comments, and a new table goes above the file's
+// closing comments.
+func TestSetLeavesCommentsTheirOwners(t *testing.T) {
+	for _, c := range []struct {
+		doc  string
+		key  Key
+		want string
+	}{
+		{"# file\n\n# about a\n[a]\nx = 1\n", Key{"top"}, "# file\n\ntop = 1\n\n# about a\n[a]\nx = 1\n"},
+		{"[a]\nx = 1\n\n# end\n", Key{"b", "y"}, "[a]\nx = 1\n\n[b]\ny = 1\n\n# end\n"},
+	} {
+		d, err := ParseDocument([]byte(c.doc))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := d.Set(c.key, 1); err != nil {
+			t.Fatal(err)
+		}
+		if got := string(d.Bytes()); got != c.want {
+			t.Errorf("set %s:\n%s\nwant:\n%s", c.key, got, c.want)
+		}
+	}
+}
+
 // TestDocumentRefusesAnEditItCannotPlace: under an array of tables, inside an
 // inline value, or over a table, an edit fails and the document stays as it was.
 func TestDocumentRefusesAnEditItCannotPlace(t *testing.T) {
@@ -200,14 +258,17 @@ func TestKeyStringQuotesWhatABareKeyCannotHold(t *testing.T) {
 }
 
 // FuzzDocument: a document that parses keeps its bytes; deleting any key it sets
-// leaves a shorter document that parses without that key, and setting one or a
+// leaves a shorter document that parses without that key, of whole lines it had
-// key beside it, where the edit is taken, leaves it holding the value set.
+// in their order, and setting one or a key beside it, where the edit is taken,
+// leaves it holding the value set, a new key keeping every line it had.
 func FuzzDocument(f *testing.F) {
 	for _, s := range []string{
 		settingsDoc,
 		"a = 1\n[t]\nk = [1,\n2] # c\n[[r]]\nx = 1\n[[r]]\nx = 2\n",
 		"t.a.b = { c = 1 }\n# end",
 		"\"q.k\" = 1\r\n[\"x y\".z]\r\nv = \"#\"\r\n",
+		"# file\n\n# a\n[a]\nx = 1 # x\n# end a\n\n\n# b\n\n[b]\ny = [\n# in\n1,\n]\n\n# foot\n",
+		"A={}",
 	} {
 		f.Add([]byte(s))
 	}
@@ -234,13 +295,21 @@ func FuzzDocument(f *testing.F) {
 			if len(c.Bytes()) >= len(data) {
 				t.Fatalf("delete of %s did not shorten the document", s.path())
 			}
+			if !keptLines(c.Bytes(), data) {
+				t.Fatalf("delete of %s left lines the document did not have:\n%s", s.path(), c.Bytes())
+			}
 		}
 		for _, s := range d.stmts {
 			for _, k := range []Key{s.path(), append(s.path()[:len(s.path())-1:len(s.path())-1], "added")} {
 				c, _ := ParseDocument(data)
-				if err := c.Set(k, int64(7)); err == nil && value(c.Tree(), k) != int64(7) {
+				added := value(c.Tree(), k) == nil
+				err := c.Set(k, int64(7))
+				if err == nil && value(c.Tree(), k) != int64(7) {
 					t.Fatalf("set %s took but reads %v:\n%s", k, value(c.Tree(), k), c.Bytes())
 				}
+				if err == nil && added && !keptLines(data, c.Bytes()) {
+					t.Fatalf("set of new %s lost a line:\n%s", k, c.Bytes())
+				}
 			}
 		}
 	})
@@ -248,6 +317,22 @@ func FuzzDocument(f *testing.F) {
 
 func has(tree map[string]any, k Key) bool { return value(tree, k) != nil }
 
+// keptLines reports every line of part is a line of whole, in whole's order; a
+// last line may gain its newline
+func keptLines(part, whole []byte) bool {
+	lines := strings.Split(string(whole), "\n")
+	for _, l := range strings.Split(string(part), "\n") {
+		for len(lines) > 0 && lines[0] != l {
+			lines = lines[1:]
+		}
+		if len(lines) == 0 {
+			return l == ""
+		}
+		lines = lines[1:]
+	}
+	return true
+}
+
 // value is what k holds in tree, in the last table of an array of tables
 func value(tree map[string]any, k Key) any {
 	var cur any = tree

encode.go

@@ -8,15 +8,14 @@ import (
 	"reflect"
 	"slices"
 	"strconv"
-	"strings"
 	"unicode/utf8"
 )
 
-// Marshal encodes a struct or string-keyed map with deterministic key order.
+// Marshal encodes a struct or string-keyed map with deterministic key order, laid
-// Nil pointer/interface fields are omitted; nil pointer/interface array elements,
+// out as Format lays it. Nil pointer/interface fields are omitted; nil
-// cycles/excessive nesting,
+// pointer/interface array elements, cycles/excessive nesting, non-finite floats,
-// non-finite floats, and integers outside TOML's signed 64-bit range error.
+// and integers outside TOML's signed 64-bit range error. Datetimes are not
-// Comments and whitespace are not retained. Datetimes are not supported.
+// supported.
 func Marshal(v any) ([]byte, error) {
 	var e encoder
 	if err := e.table(reflect.ValueOf(v), "", 0); err != nil {
@@ -27,7 +26,8 @@ func Marshal(v any) ([]byte, error) {
 
 type encoder struct{ buf bytes.Buffer }
 type entry struct {
-	key   string
+	name  string // the key
+	key   string // as a document spells it
 	value reflect.Value
 }
 
@@ -58,52 +58,44 @@ func entries(v reflect.Value) ([]entry, error) {
 		if v.Type().Key().Kind() != reflect.String {
 			return nil, fmt.Errorf("map keys must be strings")
 		}
-		out = make([]entry, 0, v.Len())
+		// The values go into one array, each set in place, not a copy apiece
-		it := v.MapRange()
+		n := v.Len()
-		for it.Next() {
+		out = make([]entry, 0, n)
-			out = append(out, entry{it.Key().String(), it.Value()})
+		values := reflect.MakeSlice(reflect.SliceOf(v.Type().Elem()), n, n)
+		key := reflect.New(v.Type().Key()).Elem()
+		for it := v.MapRange(); it.Next() && len(out) < n; {
+			key.SetIterKey(it)
+			value := values.Index(len(out))
+			value.SetIterValue(it)
+			out = append(out, entry{name: key.String(), value: value})
 		}
 	case reflect.Struct:
-		t := v.Type()
+		fields := fieldsOf(v.Type())
-		out = make([]entry, 0, t.NumField())
+		out = make([]entry, 0, len(fields))
-		for i := 0; i < t.NumField(); i++ {
+		for _, f := range fields {
-			f := t.Field(i)
+			if value := v.Field(f.index); !f.omitempty || !isEmptyValue(value) {
-			if !f.IsExported() {
+				out = append(out, entry{name: f.key, value: value})
-				continue
 			}
-			parts := strings.Split(f.Tag.Get("toml"), ",")
-			if parts[0] == "-" {
-				continue
-			}
-			key := f.Name
-			if parts[0] != "" {
-				key = parts[0]
-			}
-			if slices.Contains(parts[1:], "omitempty") && isEmptyValue(v.Field(i)) {
-				continue
-			}
-			out = append(out, entry{key, v.Field(i)})
 		}
 	default:
 		return nil, fmt.Errorf("expected struct or map, got %s", v.Kind())
 	}
-	slices.SortFunc(out, func(a, b entry) int { return cmp.Compare(a.key, b.key) })
+	slices.SortFunc(out, func(a, b entry) int { return cmp.Compare(a.name, b.name) })
 	result := out[:0]
 	previous := ""
 	for i, x := range out {
-		if _, err := keyText(x.key); err != nil {
+		if x.key, err = keyText(x.name); err != nil {
 			return nil, err
 		}
-		if i > 0 && previous == x.key {
+		if i > 0 && previous == x.name {
-			return nil, fmt.Errorf("duplicate TOML key %q", x.key)
+			return nil, fmt.Errorf("duplicate TOML key %q", x.name)
 		}
-		previous = x.key
+		previous = x.name
-		value, err := indirect(x.value)
+		if x.value, err = indirect(x.value); err != nil {
-		if err != nil {
 			return nil, err
 		}
-		if value.IsValid() {
+		if x.value.IsValid() {
-			result = append(result, entry{x.key, value})
+			result = append(result, x)
 		}
 	}
 	return result, nil
@@ -121,32 +113,32 @@ func (e *encoder) table(v reflect.Value, prefix string, depth int) error {
 	if err != nil {
 		return err
 	}
-	var tables []entry
+	tables := fields[:0] // filtered in place: a table is never ahead of its entry
 	for _, x := range fields {
-		key, _ := keyText(x.key)
+		if (isTable(x.value) || isTableArray(x.value)) && len(prefix)+len(x.key) < maxHeaderBytes {
-		if (isTable(x.value) || isTableArray(x.value)) && len(prefix)+len(key) < maxHeaderBytes {
 			tables = append(tables, x)
 			continue
 		}
-		e.buf.WriteString(key + " = ")
+		e.buf.WriteString(x.key)
+		e.buf.WriteString(" = ")
 		if err := e.value(x.value, depth+1); err != nil {
-			return fmt.Errorf("key %q: %w", x.key, err)
+			return fmt.Errorf("key %q: %w", x.name, err)
 		}
 		e.buf.WriteByte('\n')
 	}
 	for _, x := range tables {
-		key, _ := keyText(x.key)
+		key := x.key
 		if prefix != "" {
 			key = prefix + "." + key
 		}
 		if isTable(x.value) {
-			e.buf.WriteString("\n[" + key + "]\n")
+			e.header(key, "[", "]")
 			if err := e.table(x.value, key, depth+1); err != nil {
 				return err
 			}
 		} else {
 			for i := 0; i < x.value.Len(); i++ {
-				e.buf.WriteString("\n[[" + key + "]]\n")
+				e.header(key, "[[", "]]")
 				if err := e.table(x.value.Index(i), key, depth+1); err != nil {
 					return err
 				}
@@ -156,6 +148,17 @@ func (e *encoder) table(v reflect.Value, prefix string, depth int) error {
 	return nil
 }
 
+// header writes a table's header a blank line below what is above it
+func (e *encoder) header(key, open, close string) {
+	if e.buf.Len() > 0 {
+		e.buf.WriteString("\n")
+	}
+	e.buf.WriteString(open)
+	e.buf.WriteString(key)
+	e.buf.WriteString(close)
+	e.buf.WriteString("\n")
+}
+
 func isTable(v reflect.Value) bool { return v.Kind() == reflect.Map || v.Kind() == reflect.Struct }
 func isTableArray(v reflect.Value) bool {
 	if v.Kind() != reflect.Slice && v.Kind() != reflect.Array || v.Len() == 0 {
@@ -190,22 +193,22 @@ func (e *encoder) value(v reflect.Value, depth int) error {
 		}
 		e.string(v.String())
 	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
-		e.buf.WriteString(strconv.FormatInt(v.Int(), 10))
+		e.buf.Write(strconv.AppendInt(e.buf.AvailableBuffer(), v.Int(), 10))
 	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
 		if v.Uint() > math.MaxInt64 {
 			return fmt.Errorf("unsigned integer %d exceeds TOML int64 range", v.Uint())
 		}
-		e.buf.WriteString(strconv.FormatUint(v.Uint(), 10))
+		e.buf.Write(strconv.AppendUint(e.buf.AvailableBuffer(), v.Uint(), 10))
 	case reflect.Float32, reflect.Float64:
 		f := v.Float()
 		if math.IsNaN(f) || math.IsInf(f, 0) {
 			return fmt.Errorf("non-finite float")
 		}
-		s := strconv.FormatFloat(f, 'g', -1, v.Type().Bits())
+		b := strconv.AppendFloat(e.buf.AvailableBuffer(), f, 'g', -1, v.Type().Bits())
-		if !strings.ContainsAny(s, ".eE") {
+		if !bytes.ContainsAny(b, ".eE") {
-			s += ".0"
+			b = append(b, ".0"...)
 		}
-		e.buf.WriteString(s)
+		e.buf.Write(b)
 	case reflect.Slice, reflect.Array:
 		e.buf.WriteByte('[')
 		for i := 0; i < v.Len(); i++ {
@@ -225,14 +228,18 @@ func (e *encoder) value(v reflect.Value, depth int) error {
 		e.buf.WriteByte('{')
 		for i, x := range fields {
 			if i > 0 {
-				e.buf.WriteString(", ")
+				e.buf.WriteByte(',')
 			}
-			key, _ := keyText(x.key)
+			e.buf.WriteByte(' ')
-			e.buf.WriteString(key + " = ")
+			e.buf.WriteString(x.key)
+			e.buf.WriteString(" = ")
 			if err := e.value(x.value, depth+1); err != nil {
 				return err
 			}
 		}
+		if len(fields) > 0 {
+			e.buf.WriteByte(' ')
+		}
 		e.buf.WriteByte('}')
 	default:
 		return fmt.Errorf("unsupported type %s", v.Type())
@@ -274,7 +281,7 @@ func (e *encoder) string(s string) {
 		case '\f':
 			e.buf.WriteString(`\f`)
 		default:
-			if r < 0x20 || r == 0x7f {
+			if r < 0x20 || r == 0x7f || hides(r) {
 				fmt.Fprintf(&e.buf, `\u%04X`, r)
 			} else {
 				e.buf.WriteRune(r)

format.go

@@ -0,0 +1,730 @@
+package toml
+
+import (
+	"bytes"
+	"cmp"
+	"fmt"
+	"reflect"
+	"slices"
+	"strings"
+	"unicode/utf8"
+)
+
+// Formatter writes a document in one layout, each comment kept with what it
+// belongs to as a Document reads it. The zero value indents with a tab.
+type Formatter struct {
+	Indent       string // one level of indentation, up to 8 tabs or spaces; a tab when empty
+	IndentTables bool   // indent a table a level per table it is under whose header is above it
+	Align        bool   // line up the = and the comments of keys on consecutive lines
+}
+
+// A line indents at most maxIndent levels and aligns to at most alignKey runes of
+// key and alignCode of line, so no line of the result grows past its source by
+// more than a fixed amount, however deep or long the source
+const (
+	maxIndent = 16
+	alignKey  = 32
+	alignCode = 80
+)
+
+// Format is Formatter{}.Format
+func Format(data []byte) ([]byte, error) { return Formatter{}.Format(data) }
+
+// Format writes data spaced as Marshal writes, each table once at its first header:
+// a reopening header's keys, and a dotted key reaching a table with a header, move
+// there with their comments. A blank line run becomes one, or two below a comment
+// that belongs above. It refuses what Parse refuses, a comment with a rune that can
+// hide text, and a result that would not read as data does.
+func (f Formatter) Format(data []byte) ([]byte, error) {
+	indent := cmp.Or(f.Indent, "\t")
+	if strings.Trim(indent, " \t") != "" || len(indent) > 8 {
+		return nil, fmt.Errorf("indent %q is not up to 8 tabs and spaces", f.Indent)
+	}
+	tree, err := NewParser(data).Parse()
+	if err != nil {
+		return nil, err
+	}
+	if err := hiddenInComments(data); err != nil {
+		return nil, err
+	}
+	d := layout(data)
+	p := printer{indent: indent, align: f.Align}
+	p.document(&d, regroup(&d, tree), f.IndentTables)
+	out := p.buf.Bytes()
+	got, err := NewParser(out).Parse()
+	if err != nil {
+		return nil, fmt.Errorf("format refused, its result would not parse: %w", err)
+	}
+	if !reflect.DeepEqual(got, tree) {
+		return nil, fmt.Errorf("format refused, its result would read otherwise")
+	}
+	return out, nil
+}
+
+// hides reports a rune that changes how the text after it displays: an
+// embedding, override or isolate of bidirectional text, or a line break but LF
+func hides(r rune) bool {
+	return r >= 0x202A && r <= 0x202E || r >= 0x2066 && r <= 0x2069 || r == 0x85 || r == 0x2028 || r == 0x2029
+}
+
+// hiddenInComments refuses a comment holding a rune hides reports; a comment has
+// no escape to write it as
+func hiddenInComments(data []byte) error {
+	lx := NewLexer(data)
+	for t := lx.NextToken(); t.Type != TokenEOF; t = lx.NextToken() {
+		if t.Type != TokenComment {
+			continue
+		}
+		if i := strings.IndexFunc(t.Literal, hides); i >= 0 {
+			r, _ := utf8.DecodeRuneInString(t.Literal[i:])
+			return fmt.Errorf("comment at line %d holds %U, which can hide text", t.Line, r)
+		}
+	}
+	return nil
+}
+
+// visible spells a string's source with each rune hides reports escaped, the
+// value the same: a lone backslash before one, an unknown escape kept as
+// written, becomes an escaped backslash
+func visible(src string) string {
+	if !strings.ContainsFunc(src, hides) {
+		return src
+	}
+	var b strings.Builder
+	escaped := false
+	for _, r := range src {
+		switch {
+		case hides(r):
+			if escaped {
+				b.WriteByte('\\')
+			}
+			fmt.Fprintf(&b, `\u%04X`, r)
+			escaped = false
+		case r == '\\':
+			b.WriteRune(r)
+			escaped = !escaped
+		default:
+			b.WriteRune(r)
+			escaped = false
+		}
+	}
+	return b.String()
+}
+
+// piece is a statement as Format writes it, with the comments that go where it
+// goes
+type piece struct {
+	key    []string // its parts as spelled, relative to its section
+	header bool
+	array  bool
+	value  *node
+	note   string     // the comment after it on its line
+	heads  [][]string // comment groups above it, each a blank line apart
+	lead   []string   // comment lines right above it
+	tail   []string   // comment lines right below it
+	trail  [][]string // comment groups below it, each a blank line apart
+	brk    bool       // a blank line above it
+}
+
+// node is a value: a scalar's source, or an array or inline table
+type node struct {
+	text     string
+	open     TokenType // TokenLBracket or TokenLBrace for a container
+	note     string    // a comment after the opening bracket on its line
+	items    []elem
+	tail     []line // comment lines after the last element
+	closeOwn bool   // the closing bracket on a line of its own
+}
+
+type elem struct {
+	before []line // comment lines above it
+	broken bool   // on a line after what precedes it
+	blank  bool   // a blank line right above it
+	key    []string
+	value  *node
+	note   string
+}
+
+type line struct {
+	text  string
+	blank bool // a blank line right above it
+}
+
+// section is a table and the pieces written under its header, the root's first
+type section struct {
+	path   Key
+	header *piece
+	pieces []*piece
+	into   *section // the first header of the table this one opens again
+	keep   bool     // under an array of tables, where no key moves in or out
+	depth  int      // the tables it is under whose headers come before it
+}
+
+// tables is a trie of the paths headers name, so finding a key's table costs
+// its length: home is the section of a table's first header keys can move to
+type tables struct {
+	sub    map[string]*tables
+	home   *section
+	headed bool
+}
+
+// child is t's node for k, made if missing
+func (t *tables) child(k string) *tables {
+	if t.sub == nil {
+		t.sub = map[string]*tables{}
+	}
+	c := t.sub[k]
+	if c == nil {
+		c = &tables{}
+		t.sub[k] = c
+	}
+	return c
+}
+
+// regroup gives each statement's piece to the section it is written in. A key
+// goes to the deepest table its path reaches that has a header, a reopened
+// header's keys to the first header; heads and lead go with the first piece of a
+// run, tail and trail with the last, and a run that loses its first piece keeps
+// its blank line above.
+func regroup(d *Document, tree map[string]any) []*section {
+	pieces := make([]*piece, len(d.stmts))
+	for i, s := range d.stmts {
+		pieces[i] = d.piece(s)
+	}
+	for k, u := range d.units {
+		first, last := pieces[u.first], pieces[u.last-1]
+		first.heads, first.brk = d.groups(u.start, d.stmts[u.first].lead), k > 0 || d.pre > 0
+		last.tail, last.trail = d.trailing(d.stmts[u.last-1].end, u.end)
+	}
+
+	root := &section{}
+	sections := []*section{root}
+	var trie tables
+	in := make([]*section, len(d.stmts))
+	cur := root
+	for i, s := range d.stmts {
+		if s.header {
+			cur = &section{path: s.key, header: pieces[i], keep: s.array || underArray(tree, s.key)}
+			t := &trie
+			for n, k := range s.key {
+				if t = t.child(k); t.headed && n < len(s.key)-1 {
+					cur.depth++
+				}
+			}
+			t.headed = true
+			switch {
+			case cur.keep:
+				sections = append(sections, cur)
+			case t.home == nil:
+				t.home = cur
+				sections = append(sections, cur)
+			default:
+				cur.into = t.home
+			}
+		}
+		in[i] = cur
+	}
+	dst := make([]*section, len(d.stmts))
+	for i, s := range d.stmts {
+		src := in[i]
+		dst[i] = cmp.Or(src.into, src)
+		if s.header || src.keep {
+			continue
+		}
+		at, t := 0, &trie
+		for _, k := range src.path {
+			t = t.sub[k]
+		}
+		for n, k := range s.key[:len(s.key)-1] {
+			if t = t.sub[k]; t == nil {
+				break
+			}
+			if t.home != nil {
+				dst[i], at = t.home, n+1
+			}
+		}
+		pieces[i].key = pieces[i].key[at:]
+	}
+	for _, u := range d.units {
+		if dst[u.first] == in[u.first] {
+			continue
+		}
+		for i := u.first + 1; i < u.last; i++ {
+			if dst[i] == in[i] {
+				pieces[i].brk = true
+				break
+			}
+		}
+	}
+
+	// A reopened header's comments go above the key that touched it, else below
+	// what is written before them where its keys go
+	type move struct {
+		p     *piece
+		to    *section
+		trail [][]string // with no piece, comments for the last piece of to
+	}
+	var moves []move
+	opened := map[*section]bool{} // a reopened section that has begun in its first
+	for i, s := range d.stmts {
+		p, src := pieces[i], in[i]
+		switch {
+		case s.header && src.into == nil:
+		case s.header:
+			at := slices.Concat(p.lead, notes(p.note), p.tail) // the lines around the header's
+			next := i + 1
+			if len(at) > 0 && next < len(d.stmts) && in[next] == src && dst[next] == src.into && !pieces[next].brk {
+				k := pieces[next]
+				k.heads, k.lead = p.heads, slices.Concat(at, k.lead)
+			} else {
+				groups := slices.Clone(p.heads)
+				if len(at) > 0 {
+					groups = append(groups, at)
+				}
+				if groups = append(groups, p.trail...); len(groups) > 0 {
+					moves = append(moves, move{to: src.into, trail: groups})
+				}
+			}
+		case dst[i] == src:
+			src.pieces = append(src.pieces, p)
+		default:
+			if dst[i] == src.into && !opened[src] {
+				opened[src], p.brk = true, true
+			}
+			moves = append(moves, move{p: p, to: dst[i]})
+		}
+	}
+	for _, m := range moves {
+		if m.p == nil {
+			last := m.to.header
+			if n := len(m.to.pieces); n > 0 {
+				last = m.to.pieces[n-1]
+			}
+			last.trail = append(last.trail, m.trail...)
+			continue
+		}
+		if len(m.to.pieces) == 0 {
+			m.p.brk = len(m.p.heads) > 0
+		}
+		m.to.pieces = append(m.to.pieces, m.p)
+	}
+	return sections
+}
+
+// underArray reports a path through an array of tables, whose headers name a
+// different table each time
+func underArray(tree map[string]any, path Key) bool {
+	var cur any = tree
+	for _, k := range path {
+		m, ok := cur.(map[string]any)
+		if !ok {
+			break
+		}
+		cur = m[k]
+	}
+	_, ok := cur.([]map[string]any)
+	return ok
+}
+
+// piece reads a statement's source into its piece
+func (d *Document) piece(s statement) *piece {
+	c := cursor{src: d.data[s.start:s.end]}
+	c.lx = Lexer{input: c.src, line: 1}
+	c.next()
+	p := &piece{header: s.header, array: s.array}
+	if s.header {
+		for c.t.Type == TokenLBracket {
+			c.next()
+		}
+		p.key = c.key()
+		for c.t.Type == TokenRBracket {
+			c.next()
+		}
+	} else {
+		p.key = c.key()
+		c.next() // past =
+		p.value = c.value()
+	}
+	if c.t.Type == TokenComment {
+		p.note = comment(c.t.Literal)
+	}
+	if g := d.groups(s.lead, s.start); len(g) > 0 {
+		p.lead = g[0]
+	}
+	return p
+}
+
+// comment is a comment's text as written, its trailing blanks dropped
+func comment(text string) string { return "#" + strings.TrimRight(text, " \t") }
+
+// notes is a header's comment on its line as a line of its own, if it has one
+func notes(note string) []string {
+	if note == "" {
+		return nil
+	}
+	return []string{note}
+}
+
+// groups reads the comment lines of data[from:to], which holds only comment and
+// blank lines, a group to each run no blank line parts
+func (d *Document) groups(from, to int) [][]string {
+	var out [][]string
+	fresh := true
+	for l := range strings.SplitSeq(string(d.data[from:to]), "\n") {
+		switch l = strings.Trim(l, " \t\r"); {
+		case l == "":
+			fresh = true
+		case fresh:
+			out, fresh = append(out, []string{l}), false
+		default:
+			out[len(out)-1] = append(out[len(out)-1], l)
+		}
+	}
+	return out
+}
+
+// trailing reads the comments below a run: the group touching it, then the rest
+func (d *Document) trailing(from, to int) ([]string, [][]string) {
+	g := d.groups(from, to)
+	first, _, _ := bytes.Cut(d.data[from:to], []byte("\n"))
+	if len(g) > 0 && len(bytes.Trim(first, " \t\r")) > 0 {
+		return g[0], g[1:]
+	}
+	return nil, g
+}
+
+// cursor walks a statement's tokens
+type cursor struct {
+	src []byte
+	lx  Lexer
+	t   Token
+}
+
+func (c *cursor) next() { c.t = c.lx.NextToken() }
+
+// spell is the token as written; a string's runes that can hide text escaped
+func (c *cursor) spell() string {
+	text := string(c.src[c.t.Pos:c.t.End])
+	if c.t.Type == TokenString {
+		return visible(text)
+	}
+	return text
+}
+
+// key reads a dotted key's parts as spelled, leaving the cursor past them
+func (c *cursor) key() []string {
+	var k []string
+	for {
+		k = append(k, c.spell())
+		if c.next(); c.t.Type != TokenDot {
+			return k
+		}
+		c.next()
+	}
+}
+
+// value reads a value, leaving the cursor past it
+func (c *cursor) value() *node {
+	if c.t.Type != TokenLBracket && c.t.Type != TokenLBrace {
+		n := &node{text: c.spell()}
+		c.next()
+		return n
+	}
+	n := &node{open: c.t.Type}
+	closer := TokenRBracket
+	if n.open == TokenLBrace {
+		closer = TokenRBrace
+	}
+	c.next()
+	var pending []line
+	nl := 0 // newlines since the last element or comment
+	for c.t.Type != closer {
+		inside := len(n.items) > 0 || len(pending) > 0
+		switch c.t.Type {
+		case TokenNewline:
+			nl++
+		case TokenComma:
+		case TokenComment:
+			text := comment(c.t.Literal)
+			switch {
+			case nl > 0:
+				pending = append(pending, line{text, nl > 1 && inside})
+			case len(n.items) > 0:
+				n.items[len(n.items)-1].note = text
+			default:
+				n.note = text
+			}
+			nl = 0
+		default:
+			e := elem{before: pending, broken: nl > 0 || len(pending) > 0, blank: nl > 1 && inside}
+			pending = nil
+			if n.open == TokenLBrace {
+				e.key = c.key()
+				c.next() // past =
+			}
+			e.value = c.value()
+			n.items = append(n.items, e)
+			nl = 0
+			continue
+		}
+		c.next()
+	}
+	n.tail, n.closeOwn = pending, nl > 0
+	c.next()
+	return n
+}
+
+// printer writes pieces, its depth the indentation of the line being written
+type printer struct {
+	buf    bytes.Buffer
+	indent string
+	align  bool
+	depth  int
+}
+
+func (p *printer) newline(depth int) {
+	p.buf.WriteByte('\n')
+	p.indentTo(depth)
+}
+
+func (p *printer) indentTo(depth int) {
+	for range min(depth, maxIndent) {
+		p.buf.WriteString(p.indent)
+	}
+	p.depth = depth
+}
+
+// document writes the file's opening comments, each section's pieces and the
+// file's closing comments. Two blank lines follow a group of comments that
+// belongs to what is above it, after a blank line, so it reads that way again.
+func (p *printer) document(d *Document, sections []*section, indentTables bool) {
+	pre := d.groups(0, d.pre)
+	p.groups(pre, 0)
+	gap := min(len(pre), 2) // blank lines after the opening comments
+	var foot [][]string
+	if d.foot >= 0 {
+		foot = d.groups(d.foot, len(d.data))
+	} else if s := lastSection(sections); s != nil {
+		// Comments below the last piece, a blank line apart, read as the file's
+		// closing ones; written so, they read the same again
+		last := s.header
+		if n := len(s.pieces); n > 0 {
+			last = s.pieces[n-1]
+		}
+		foot, last.trail = last.trail, nil
+	}
+	var last *piece
+	for _, s := range sections {
+		depth := 0
+		if indentTables {
+			depth = s.depth
+		}
+		all := s.pieces
+		if s.header != nil {
+			all = append([]*piece{s.header}, all...)
+		}
+		widths := p.widths(all, depth)
+		for i, pc := range all {
+			p.blanks(gap, last, pc.brk || len(pc.heads) > 0)
+			p.piece(pc, depth, widths[i])
+			last = pc
+		}
+	}
+	if len(foot) > 0 {
+		p.blanks(gap, last, true)
+		p.groups(foot, 0)
+	}
+}
+
+// lastSection is the last section with a piece to write
+func lastSection(sections []*section) *section {
+	for i := len(sections) - 1; i >= 0; i-- {
+		if s := sections[i]; s.header != nil || len(s.pieces) > 0 {
+			return s
+		}
+	}
+	return nil
+}
+
+// blanks writes the blank lines above what follows last: two below comment groups
+// that belong to it, else one where either asks for it
+func (p *printer) blanks(gap int, last *piece, brk bool) {
+	n := 0
+	switch {
+	case last == nil:
+		n = gap
+	case len(last.trail) > 0:
+		n = 2
+	case brk || len(last.tail) > 0:
+		n = 1
+	}
+	for range n {
+		p.buf.WriteByte('\n')
+	}
+}
+
+func (p *printer) groups(g [][]string, depth int) {
+	for i, lines := range g {
+		if i > 0 {
+			p.buf.WriteByte('\n')
+		}
+		p.comments(lines, depth)
+	}
+}
+
+func (p *printer) comments(lines []string, depth int) {
+	for _, l := range lines {
+		p.indentTo(depth)
+		p.buf.WriteString(l)
+		p.buf.WriteByte('\n')
+	}
+}
+
+// width is where a piece's = and comment go when its run is aligned
+type width struct{ key, code int }
+
+// widths lines up each run of keys on consecutive lines, each on one line and
+// short enough, no comment line among them
+func (p *printer) widths(all []*piece, depth int) []width {
+	w := make([]width, len(all))
+	if !p.align {
+		return w
+	}
+	codes := make([]string, len(all))
+	for i, pc := range all {
+		if pc.header || utf8.RuneCountInString(strings.Join(pc.key, ".")) > alignKey {
+			continue
+		}
+		if c := p.code(pc, depth, alignKey); utf8.RuneCountInString(c) <= alignCode {
+			codes[i] = p.code(pc, depth, 0)
+		}
+	}
+	for i := 0; i < len(all); {
+		j := i + 1
+		for j < len(all) && joins(all[j-1], all[j], codes[j-1], codes[j]) {
+			j++
+		}
+		if codes[i] == "" {
+			i = j
+			continue
+		}
+		var run width
+		for k := i; k < j; k++ {
+			run.key = max(run.key, utf8.RuneCountInString(strings.Join(all[k].key, ".")))
+		}
+		for k := i; k < j; k++ {
+			run.code = max(run.code, utf8.RuneCountInString(p.code(all[k], depth, run.key)))
+		}
+		for k := i; k < j; k++ {
+			w[k] = run
+		}
+		i = j
+	}
+	return w
+}
+
+// joins reports b on the line after a in one aligned run
+func joins(a, b *piece, ac, bc string) bool {
+	return ac != "" && bc != "" && !b.brk && len(b.heads)+len(b.lead) == 0 &&
+		len(a.tail)+len(a.trail) == 0 && !strings.Contains(ac+bc, "\n")
+}
+
+// code is a key's line before its comment, the key padded to keyWidth
+func (p *printer) code(pc *piece, depth, keyWidth int) string {
+	q := printer{indent: p.indent, depth: depth}
+	k := strings.Join(pc.key, ".")
+	q.buf.WriteString(k)
+	q.buf.WriteString(strings.Repeat(" ", max(0, keyWidth-utf8.RuneCountInString(k))))
+	q.buf.WriteString(" = ")
+	q.node(pc.value)
+	return q.buf.String()
+}
+
+func (p *printer) piece(pc *piece, depth int, w width) {
+	for _, g := range pc.heads {
+		p.comments(g, depth)
+		p.buf.WriteByte('\n')
+	}
+	p.comments(pc.lead, depth)
+	p.indentTo(depth)
+	code := "[" + strings.Join(pc.key, ".") + "]"
+	switch {
+	case pc.array:
+		code = "[" + code + "]"
+	case !pc.header:
+		code = p.code(pc, depth, w.key)
+	}
+	p.buf.WriteString(code)
+	if pc.note != "" {
+		p.buf.WriteString(strings.Repeat(" ", max(0, w.code-utf8.RuneCountInString(code))+1))
+		p.buf.WriteString(pc.note)
+	}
+	p.buf.WriteByte('\n')
+	p.comments(pc.tail, depth)
+	for _, g := range pc.trail {
+		p.buf.WriteByte('\n')
+		p.comments(g, depth)
+	}
+}
+
+// node writes a value; a container's elements on lines of their own go a level
+// deeper than the line its opening bracket is on, and it keeps a trailing comma
+// exactly when its closing bracket has a line of its own
+func (p *printer) node(n *node) {
+	if n.open == 0 {
+		p.buf.WriteString(n.text)
+		return
+	}
+	open, close, pad := "[", "]", ""
+	if n.open == TokenLBrace {
+		open, close, pad = "{", "}", " "
+	}
+	base := p.depth
+	p.buf.WriteString(open)
+	note := n.note
+	if note != "" {
+		p.buf.WriteString(" " + note)
+	}
+	for i, e := range n.items {
+		switch {
+		case e.broken || note != "":
+			for _, c := range e.before {
+				p.line(c, base+1)
+			}
+			p.line(line{blank: e.blank}, base+1)
+		case i == 0:
+			p.buf.WriteString(pad)
+		default:
+			p.buf.WriteByte(' ')
+		}
+		if e.key != nil {
+			p.buf.WriteString(strings.Join(e.key, ".") + " = ")
+		}
+		p.node(e.value)
+		if i < len(n.items)-1 || n.closeOwn {
+			p.buf.WriteByte(',')
+		}
+		if note = e.note; note != "" {
+			p.buf.WriteString(" " + note)
+		}
+	}
+	for _, c := range n.tail {
+		p.line(c, base+1)
+	}
+	if n.closeOwn || note != "" || len(n.tail) > 0 {
+		p.newline(base)
+	} else if len(n.items) > 0 {
+		p.buf.WriteString(pad)
+	}
+	p.buf.WriteString(close)
+}
+
+// line starts a line at depth inside a value, a blank one above it if asked,
+// and writes the comment it holds
+func (p *printer) line(l line, depth int) {
+	if l.blank {
+		p.buf.WriteByte('\n')
+	}
+	p.newline(depth)
+	p.buf.WriteString(l.text)
+}

format_test.go

@@ -0,0 +1,188 @@
+package toml
+
+import (
+	"bytes"
+	"fmt"
+	"reflect"
+	"slices"
+	"strings"
+	"testing"
+)
+
+func formats(t *testing.T, f Formatter, in, want string) {
+	t.Helper()
+	got, err := f.Format([]byte(in))
+	if err != nil {
+		t.Fatalf("format %q: %v", in, err)
+	}
+	if string(got) != want {
+		t.Fatalf("format of:\n%s\ngot:\n%s\nwant:\n%s", in, got, want)
+	}
+}
+
+// TestFormatSpacesAsMarshal: one space around = and after a comma, none inside
+// an array's brackets, one inside an inline table's; a value's lines a tab deeper
+// than the line it opens on, a trailing comma exactly when the closing bracket
+// has its own line; a blank line run kept as one; LF line ends, no trailing blank.
+func TestFormatSpacesAsMarshal(t *testing.T) {
+	formats(t, Formatter{}, "  a=[1,2,   3]\r\nb = {x=1,y = {z=[ ]}, w = { }}\r\n\r\n\r\n"+
+		"c = [\n    1, # one\n  2\n]\nd = [ [1,\n2], 3, ]   # note   \n[ t . \"u v\" ]\n[[ arr ]]\nk = {\na = 1 }\n\n",
+		"a = [1, 2, 3]\nb = { x = 1, y = { z = [] }, w = {} }\n\n"+
+			"c = [\n\t1, # one\n\t2,\n]\nd = [[1,\n\t2], 3] # note\n[t.\"u v\"]\n[[arr]]\nk = {\n\ta = 1 }\n")
+}
+
+// TestFormatKeepsCommentsWithTheirOwners: a comment stays with what fewer blank
+// lines part it from as a Document reads it, and the output's blank lines say so
+// again: two follow a group belonging to what is above it, after a blank line.
+// Comments inside a value keep their place among its elements.
+func TestFormatKeepsCommentsWithTheirOwners(t *testing.T) {
+	// The request's example: a dotted key under [section2] is section2's
+	formats(t, Formatter{}, "# comment s1\n[section1]\npart1=1 # comment s1p1\n\n#comment s2\n[section2]\n"+
+		"# comment s2p2\npart1=1\n\n# comment s1p2\nsection1.part2 = 5\n",
+		"# comment s1\n[section1]\npart1 = 1 # comment s1p1\n\n#comment s2\n[section2]\n"+
+			"# comment s2p2\npart1 = 1\n\n# comment s1p2\nsection1.part2 = 5\n")
+	formats(t, Formatter{}, "# file\n\n# also file\n\n\n[a]\nx = 1\n\n# about a\n\n\n# about b\n\n[b]\n"+
+		"y = [ # open\n\n  # first\n  1,\n\n  # second\n  2 # two\n  , 3,\n  # last\n\n]\n\n# end\n",
+		"# file\n\n# also file\n\n\n[a]\nx = 1\n\n# about a\n\n\n# about b\n\n[b]\n"+
+			"y = [ # open\n\t# first\n\t1,\n\n\t# second\n\t2, # two\n\t3,\n\t# last\n]\n\n# end\n")
+}
+
+// TestFormatWritesATableOnce: a header opening a table again gives its keys to
+// the first header, a blank line apart, its comments above them; a dotted key
+// reaching a table that has a header moves there, spelled from it, its own
+// comments along. Arrays of tables, and tables under them, stay as written.
+func TestFormatWritesATableOnce(t *testing.T) {
+	formats(t, Formatter{}, "# comment s1\n[section1]\npart1=1 # comment s1p1\n\n#comment s2\n[section2]\n"+
+		"# comment s2p2\npart1=1\n\n# comment s1p2\n[section1]\npart2 = 5\n",
+		"# comment s1\n[section1]\npart1 = 1 # comment s1p1\n\n# comment s1p2\npart2 = 5\n\n"+
+			"#comment s2\n[section2]\n# comment s2p2\npart1 = 1\n")
+	formats(t, Formatter{}, "# about part2\nsection1.part2 = 5\na.b.c = 1\n\n[a.b]\nd = 2\n[section1]\npart1 = 1\n[a]\nb.e = 3\n",
+		"[a.b]\nd = 2\nc = 1\ne = 3\n[section1]\npart1 = 1\n# about part2\npart2 = 5\n[a]\n")
+	formats(t, Formatter{}, "[[s]]\nx = 1\n[s.t]\ny = 1\n[[s]]\nx = 2\n[s.t]\ny = 2\n",
+		"[[s]]\nx = 1\n[s.t]\ny = 1\n[[s]]\nx = 2\n[s.t]\ny = 2\n")
+}
+
+// TestFormatterOptions: Indent sets a level's characters, IndentTables indents a
+// table a level per table it is under whose header comes first, Align lines up
+// the = and the comments of keys on consecutive lines, a comment or blank line
+// ending the run.
+func TestFormatterOptions(t *testing.T) {
+	formats(t, Formatter{Indent: "  ", IndentTables: true},
+		"[servers]\na = 1\n[servers.alpha]\nports = [\n80,\n]\n[servers.alpha.deep]\nx = 1\n[other.deep]\ny = 2\n",
+		"[servers]\na = 1\n  [servers.alpha]\n  ports = [\n    80,\n  ]\n    [servers.alpha.deep]\n    x = 1\n[other.deep]\ny = 2\n")
+	formats(t, Formatter{Align: true}, "[play]\nsfx = true # on\nmusic = true\nsfx_volume = 0.5 # half\n# lead\nmouse = true # m\n\nhud = false\nwait = true # w\n",
+		"[play]\nsfx        = true # on\nmusic      = true\nsfx_volume = 0.5  # half\n# lead\nmouse = true # m\n\nhud  = false\nwait = true  # w\n")
+}
+
+// TestFormatGrowsBoundedly: past 16 levels a line indents no deeper, and a key or
+// line too long to align stands alone, so however deep or long its source, the
+// result stays within a fixed multiple of it: a two-byte line such as "]" takes
+// 16 tabs and a comma at most
+func TestFormatGrowsBoundedly(t *testing.T) {
+	deep := "a = " + strings.Repeat("[\n", 999) + "1" + strings.Repeat("\n]", 999) + "\n"
+	var long strings.Builder
+	long.WriteString("\"" + strings.Repeat("k", 4096) + "\" = 1 # c\n")
+	for i := range 1000 {
+		fmt.Fprintf(&long, "k%d = %d # c\n", i, i)
+	}
+	for _, doc := range []string{deep, long.String()} {
+		out, err := Formatter{Align: true}.Format([]byte(doc))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(out) > 10*len(doc) {
+			t.Errorf("%d bytes grew to %d", len(doc), len(out))
+		}
+		if bytes.Contains(out, []byte(strings.Repeat("\t", 17))) {
+			t.Error("a line indents past 16 levels")
+		}
+	}
+}
+
+// TestFormatRefusesWhatItCannotKeep: a document Parse refuses, a comment holding
+// a rune that can hide the text after it, and an indent that is not up to 8
+// blanks are refused; such a rune in a string or quoted key is written escaped.
+func TestFormatRefusesWhatItCannotKeep(t *testing.T) {
+	for _, c := range []struct {
+		f   Formatter
+		doc string
+	}{
+		{Formatter{}, "a = \n"},
+		{Formatter{}, "a = 1 # \u202e reversed\n"},
+		{Formatter{}, "# \u2066isolate\na = 1\n"},
+		{Formatter{}, "# next\u2028line\na = 1\n"},
+		{Formatter{Indent: "# "}, "a = [\n1,\n]\n"},
+		{Formatter{Indent: strings.Repeat(" ", 9)}, "a = 1\n"},
+	} {
+		if out, err := c.f.Format([]byte(c.doc)); err == nil {
+			t.Errorf("format of %q taken:\n%s", c.doc, out)
+		}
+	}
+	formats(t, Formatter{}, "s = \"a\u202eb\"\n\"k\u2066\" = \"\\\u202e\"\n", "s = \"a\\u202Eb\"\n\"k\\u2066\" = \"\\\\\\u202E\"\n")
+}
+
+// FuzzFormat: whatever Parse takes Format writes, unless a comment holds a rune
+// that can hide text, with every option: the result reads as the document does,
+// formats to itself, keeps every comment once and has no CR or trailing blank.
+func FuzzFormat(f *testing.F) {
+	for _, s := range []string{
+		settingsDoc,
+		"# a\n\n# b\n\n\n[t]\nk = 1\n# t's\n\n\n# u\n[u]\nv = [ # o\n\n1, # one\n# c\n2\n,3]\n\n# end\n",
+		"x.y = 1\n[x]\nz = 2\n[w]\n[x] # again\nq = { a.b = 1, c = [ {d = 1} ] }\n",
+		"[[s]]\n[s.t]\na = 1\n[[s]]\n[s.t]\nb = 2\n",
+		"k = \"\\u00e9 \\\\ \\\" \\t\"\r\n\"q k\" = \"\u202e\"\r\n",
+		"a = [[[[1]]],\n[[2,\n3]]]\nb = {\nc = {\nd = 1,\n},\n}\n",
+		"\n\n# only\n\n# comments\n\n",
+		"a=1",
+		"[a]\n[a.b]\n[a]\nc = 1\n",
+		"[A]\n\n#\n\n[A]",
+		"[a]\n[a.A]\n[a]#",
+		"[a]\n[a.A]\n[a.A]#",
+		"[\"a b\".c]\nx = 1\n[\"a b\"]\nc.y = 2\n",
+		"[A]\n\n#\n\n\n[A0]\n[A]#",
+		"# p\n\n[a]\nx = 1\n# t\n\n# u\n\n\n[b]\n# lb\n[a] # n\n# la\ny = 2\n\n# v\n\n\n[a.c]\n[a]\n\n# w\n",
+	} {
+		f.Add([]byte(s))
+	}
+	f.Fuzz(func(t *testing.T, data []byte) {
+		tree, err := NewParser(data).Parse()
+		if err != nil || hiddenInComments(data) != nil {
+			return
+		}
+		for _, fm := range []Formatter{{}, {Indent: "  ", IndentTables: true, Align: true}} {
+			out, err := fm.Format(data)
+			if err != nil {
+				t.Fatalf("%+v refused %q: %v", fm, data, err)
+			}
+			if got, _ := NewParser(out).Parse(); !reflect.DeepEqual(got, tree) {
+				t.Fatalf("%+v changed the values:\n%s", fm, out)
+			}
+			if again, err := fm.Format(out); err != nil || !bytes.Equal(again, out) {
+				t.Fatalf("%+v does not format to itself (%v):\n%s\nthen:\n%s", fm, err, out, again)
+			}
+			if a, b := comments(data), comments(out); !slices.Equal(a, b) {
+				t.Fatalf("%+v comments %q became %q:\n%s", fm, a, b, out)
+			}
+			if bytes.Contains(out, []byte("\r")) || bytes.Contains(out, []byte(" \n")) || bytes.Contains(out, []byte("\t\n")) ||
+				bytes.Contains(out, []byte("\n\n\n\n")) || len(out) > 0 && out[len(out)-1] != '\n' {
+				t.Fatalf("%+v wrote a stray blank:\n%q", fm, out)
+			}
+			if len(out) > 48*len(data)+64 {
+				t.Fatalf("%+v grew %d bytes to %d", fm, len(data), len(out))
+			}
+		}
+	})
+}
+
+// comments lists a document's comments, trailing blanks dropped, sorted
+func comments(data []byte) []string {
+	var out []string
+	lx := NewLexer(data)
+	for t := lx.NextToken(); t.Type != TokenEOF; t = lx.NextToken() {
+		if t.Type == TokenComment {
+			out = append(out, strings.TrimRight(t.Literal, " \t"))
+		}
+	}
+	slices.Sort(out)
+	return out
+}

fuzz_test.go

@@ -1,6 +1,7 @@
 package toml
 
 import (
+	"bytes"
 	"reflect"
 	"strings"
 	"testing"
@@ -87,5 +88,9 @@ func FuzzRoundTrip(f *testing.F) {
 		if !reflect.DeepEqual(normalize(m1), normalize(m2)) {
 			t.Fatalf("fixpoint violation:\nm1: %#v\nm2: %#v\nemitted:\n%s", m1, m2, out)
 		}
+		// Marshal lays a document out as Format does
+		if f, err := Format(out); err != nil || !bytes.Equal(f, out) {
+			t.Fatalf("Format rewrites Marshal's output (%v):\n%s\nto:\n%s", err, out, f)
+		}
 	})
 }

lexer.go

@@ -2,6 +2,7 @@ package toml
 
 import (
 	"fmt"
+	"strings"
 	"unicode/utf8"
 )
 
@@ -10,15 +11,13 @@ type Lexer struct {
 	pos   int
 	line  int
 	col   int
-	width int
+	// lean is a parser's lexer: a comment or number carries no Literal, the parser
+	// reading what it needs from Pos and End, so neither allocates
+	lean bool
 }
 
 func NewLexer(input []byte) *Lexer {
-	return &Lexer{
+	return &Lexer{input: input, line: 1}
-		input: input,
-		line:  1,
-		col:   0,
-	}
 }
 
 func (l *Lexer) NextToken() Token {
@@ -35,17 +34,12 @@ func (l *Lexer) next() Token {
 	}
 
 	ch := l.peek()
-
+	switch ch {
-	if ch == '\n' {
+	case '\n':
 		l.advance()
 		return l.newToken(TokenNewline, "\n")
-	}
+	case '#':
-
-	if ch == '#' {
 		return l.readComment()
-	}
-
-	switch ch {
 	case '=':
 		l.advance()
 		return l.newToken(TokenEqual, "=")
@@ -72,27 +66,16 @@ func (l *Lexer) next() Token {
 	}
 
 	// Number: digit or sign+digit
-	if isDigit(ch) {
+	if isDigit(ch) || (ch == '+' || ch == '-') && isDigit(rune(l.byteAt(l.pos+1))) {
 		return l.readNumber()
 	}
 	if ch == '+' {
-		if isDigit(l.peekAt(1)) {
-			return l.readNumber()
-		}
 		// Lone + is invalid TOML
 		l.advance()
 		return l.newToken(TokenError, "unexpected character: +")
 	}
-	if ch == '-' {
+	// A bare key may start with a hyphen, not a digit
-		if isDigit(l.peekAt(1)) {
+	if isAlpha(ch) || ch == '_' || ch == '-' {
-			return l.readNumber()
-		}
-		// Bare key can start with hyphen
-		return l.readIdent()
-	}
-
-	// Identifier: alpha or underscore
-	if isAlpha(ch) || ch == '_' {
 		return l.readIdent()
 	}
 
@@ -101,165 +84,96 @@ func (l *Lexer) next() Token {
 }
 
 func (l *Lexer) readNumber() Token {
-	start := l.pos
+	start, startCol := l.pos, l.col
-	startCol := l.col
+	i := start
-
+	if c := l.input[i]; c == '+' || c == '-' {
-	// Optional sign
+		i++
-	if l.peek() == '+' || l.peek() == '-' {
-		l.advance()
 	}
 
-	// Check radix prefix
+	// Radix prefix
-	if l.peek() == '0' {
+	if l.byteAt(i) == '0' {
-		next := l.peekAt(1)
+		var digit func(rune) bool
-		switch next {
+		var none string
+		switch l.byteAt(i + 1) {
 		case 'x', 'X':
-			l.advance() // '0'
+			digit, none = isHexDigit, "invalid hex: no digits after prefix"
-			l.advance() // 'x'
-			return l.readHex(start, startCol)
 		case 'o', 'O':
-			l.advance()
+			digit, none = isOctalDigit, "invalid octal: no digits after prefix"
-			l.advance()
-			return l.readOctal(start, startCol)
 		case 'b', 'B':
-			l.advance()
+			digit, none = isBinaryDigit, "invalid binary: no digits after prefix"
-			l.advance()
+		}
-			return l.readBinary(start, startCol)
+		if digit != nil {
+			i += 2
+			j := i
+			for j < len(l.input) && digit(rune(l.input[j])) {
+				j++
+			}
+			l.skipTo(j)
+			if j == i {
+				return Token{Type: TokenError, Literal: none, Line: l.line, Col: startCol}
+			}
+			return l.number(TokenInteger, start, startCol)
 		}
 	}
 
-	// Integer part
+	typ := TokenInteger
-	for isDigit(l.peek()) {
+	i = l.digits(i)
-		l.advance()
-	}
-
-	isFloat := false
-
 	// Fractional part: '.' followed by digit
-	if l.peek() == '.' && isDigit(l.peekAt(1)) {
+	if l.byteAt(i) == '.' && isDigit(rune(l.byteAt(i+1))) {
-		isFloat = true
+		typ, i = TokenFloat, l.digits(i+1)
-		l.advance() // consume '.'
+		if l.byteAt(i) == '.' && isDigit(rune(l.byteAt(i+1))) {
-
+			l.skipTo(i)
-		for isDigit(l.peek()) {
-			l.advance()
-		}
-
-		// Multi-dot: another '.' followed by digit is fatal
-		if l.peek() == '.' && isDigit(l.peekAt(1)) {
 			return Token{Type: TokenError, Literal: "invalid number: multiple decimal points", Line: l.line, Col: startCol}
 		}
 	}
-
 	// Exponent: e/E followed by optional sign and digits
-	if l.peek() == 'e' || l.peek() == 'E' {
+	if c := l.byteAt(i); c == 'e' || c == 'E' {
-		next := l.peekAt(1)
+		j := i + 1
-		validExp := isDigit(next) || ((next == '+' || next == '-') && isDigit(l.peekAt(2)))
+		if c := l.byteAt(j); c == '+' || c == '-' {
-		if validExp {
+			j++
-			isFloat = true
+		}
-			l.advance() // 'e'
+		if isDigit(rune(l.byteAt(j))) {
-			if l.peek() == '+' || l.peek() == '-' {
+			typ, i = TokenFloat, l.digits(j)
-				l.advance()
-			}
-			for isDigit(l.peek()) {
-				l.advance()
-			}
 		}
 	}
-
+	l.skipTo(i)
-	lit := string(l.input[start:l.pos])
+	return l.number(typ, start, startCol)
-
-	// Validate float format: 1.e2 is invalid (no digit after dot before exponent)
-	if isFloat {
-		if err := validateFloat(lit); err != nil {
-			return Token{Type: TokenError, Literal: err.Error(), Line: l.line, Col: startCol}
-		}
-		return Token{Type: TokenFloat, Literal: lit, Line: l.line, Col: startCol}
-	}
-	return Token{Type: TokenInteger, Literal: lit, Line: l.line, Col: startCol}
 }
 
-func validateFloat(lit string) error {
+func (l *Lexer) number(typ TokenType, start, col int) Token {
-	// Find dot position (if any)
+	t := Token{Type: typ, Line: l.line, Col: col}
-	dotIdx := -1
+	if !l.lean {
-	expIdx := -1
+		t.Literal = string(l.input[start:l.pos])
-	for i := 0; i < len(lit); i++ {
-		if lit[i] == '.' {
-			dotIdx = i
-		}
-		if lit[i] == 'e' || lit[i] == 'E' {
-			expIdx = i
-			break
-		}
 	}
-
+	return t
-	if dotIdx >= 0 {
-		// Check digit before dot (ignoring sign)
-		start := 0
-		if lit[0] == '+' || lit[0] == '-' {
-			start = 1
-		}
-		if dotIdx == start {
-			return fmt.Errorf("invalid float %q: no digit before decimal point", lit)
-		}
-
-		// Check digit after dot
-		afterDot := dotIdx + 1
-		endFrac := len(lit)
-		if expIdx > 0 {
-			endFrac = expIdx
-		}
-		if afterDot >= endFrac {
-			return fmt.Errorf("invalid float %q: no digit after decimal point", lit)
-		}
-	}
-	return nil
 }
 
-func (l *Lexer) readHex(start, startCol int) Token {
+// digits is the offset past the decimal digits at i
-	if !isHexDigit(l.peek()) {
+func (l *Lexer) digits(i int) int {
-		return Token{Type: TokenError, Literal: "invalid hex: no digits after prefix", Line: l.line, Col: startCol}
+	for i < len(l.input) && isDigit(rune(l.input[i])) {
+		i++
 	}
-	for isHexDigit(l.peek()) {
+	return i
-		l.advance()
-	}
-	return Token{Type: TokenInteger, Literal: string(l.input[start:l.pos]), Line: l.line, Col: startCol}
-}
-
-func (l *Lexer) readOctal(start, startCol int) Token {
-	if !isOctalDigit(l.peek()) {
-		return Token{Type: TokenError, Literal: "invalid octal: no digits after prefix", Line: l.line, Col: startCol}
-	}
-	for isOctalDigit(l.peek()) {
-		l.advance()
-	}
-	return Token{Type: TokenInteger, Literal: string(l.input[start:l.pos]), Line: l.line, Col: startCol}
-}
-
-func (l *Lexer) readBinary(start, startCol int) Token {
-	if !isBinaryDigit(l.peek()) {
-		return Token{Type: TokenError, Literal: "invalid binary: no digits after prefix", Line: l.line, Col: startCol}
-	}
-	for isBinaryDigit(l.peek()) {
-		l.advance()
-	}
-	return Token{Type: TokenInteger, Literal: string(l.input[start:l.pos]), Line: l.line, Col: startCol}
 }
 
 func (l *Lexer) readIdent() Token {
-	start := l.pos
+	start, i := l.pos, l.pos
-	for l.pos < len(l.input) {
+	for i < len(l.input) {
-		ch := l.peek()
+		if c := rune(l.input[i]); isAlpha(c) || isDigit(c) || c == '_' || c == '-' {
-		if isAlpha(ch) || isDigit(ch) || ch == '_' || ch == '-' {
+			i++
-			l.advance()
 		} else {
 			break
 		}
 	}
-	lit := string(l.input[start:l.pos])
+	l.skipTo(i)
-	if lit == "true" || lit == "false" {
+	switch lit := l.input[start:i]; string(lit) {
-		return l.newToken(TokenBool, lit)
+	case "true":
+		return l.newToken(TokenBool, "true")
+	case "false":
+		return l.newToken(TokenBool, "false")
+	default:
+		return l.newToken(TokenIdent, string(lit))
 	}
-	return l.newToken(TokenIdent, lit)
 }
 
 func (l *Lexer) newToken(typ TokenType, literal string) Token {
@@ -272,11 +186,12 @@ func (l *Lexer) newToken(typ TokenType, literal string) Token {
 
 func (l *Lexer) advance() rune {
 	if l.pos >= len(l.input) {
-		l.width = 0
 		return 0
 	}
-	r, w := utf8.DecodeRune(l.input[l.pos:])
+	r, w := rune(l.input[l.pos]), 1
-	l.width = w
+	if r >= utf8.RuneSelf {
+		r, w = utf8.DecodeRune(l.input[l.pos:])
+	}
 	l.pos += w
 	if r == '\n' {
 		l.line++
@@ -287,38 +202,43 @@ func (l *Lexer) advance() rune {
 	return r
 }
 
+// skipTo advances to i over single-byte runes on one line
+func (l *Lexer) skipTo(i int) {
+	l.col += i - l.pos
+	l.pos = i
+}
+
 func (l *Lexer) peek() rune {
 	if l.pos >= len(l.input) {
 		return 0
 	}
+	if c := l.input[l.pos]; c < utf8.RuneSelf {
+		return rune(c)
+	}
 	r, _ := utf8.DecodeRune(l.input[l.pos:])
 	return r
 }
 
-func (l *Lexer) peekAt(n int) rune {
+// byteAt is the input byte at i, 0 past the end
-	pos := l.pos
+func (l *Lexer) byteAt(i int) byte {
-	for i := 0; i < n && pos < len(l.input); i++ {
+	if i < len(l.input) {
-		_, w := utf8.DecodeRune(l.input[pos:])
+		return l.input[i]
-		pos += w
 	}
-	if pos >= len(l.input) {
+	return 0
-		return 0
-	}
-	r, _ := utf8.DecodeRune(l.input[pos:])
-	return r
 }
 
 // skipWhitespace takes a CR only as part of CRLF: a lone one would let a
 // terminal hide what precedes it on the line, so it lexes as an error.
 func (l *Lexer) skipWhitespace() {
-	for l.pos < len(l.input) {
+	i := l.pos
-		ch := l.peek()
+	for i < len(l.input) {
-		if ch == ' ' || ch == '\t' || ch == '\r' && l.peekAt(1) == '\n' {
+		if c := l.input[i]; c == ' ' || c == '\t' || c == '\r' && l.byteAt(i+1) == '\n' {
-			l.advance()
+			i++
 		} else {
 			break
 		}
 	}
+	l.skipTo(i)
 }
 
 // readComment refuses control characters but tab, as TOML does: an ESC or a
@@ -327,21 +247,49 @@ func (l *Lexer) readComment() Token {
 	l.advance() // '#'
 	start := l.pos
 	for l.pos < len(l.input) {
-		ch := l.peek()
+		c := l.input[l.pos]
-		if ch == '\n' || ch == '\r' && l.peekAt(1) == '\n' {
+		if c == '\n' || c == '\r' && l.byteAt(l.pos+1) == '\n' {
 			break
 		}
-		if ch < 0x20 && ch != '\t' || ch == 0x7f {
+		if c < 0x20 && c != '\t' || c == 0x7f {
-			return l.newToken(TokenError, fmt.Sprintf("control character %q in comment", ch))
+			return l.newToken(TokenError, fmt.Sprintf("control character %q in comment", rune(c)))
 		}
-		l.advance()
+		w := 1
+		if c >= utf8.RuneSelf {
+			_, w = utf8.DecodeRune(l.input[l.pos:])
+		}
+		l.pos += w
+		l.col++
+	}
+	if l.lean {
+		return l.newToken(TokenComment, "")
 	}
 	return l.newToken(TokenComment, string(l.input[start:l.pos]))
 }
 
+// readString takes a string without escapes in one allocation, the string itself
 func (l *Lexer) readString() Token {
 	l.advance() // opening quote
-	var result []byte
+	start := l.pos
+	for i := start; i < len(l.input); i++ {
+		c := l.input[i]
+		if c == '"' {
+			l.col += utf8.RuneCount(l.input[start:i]) + 1
+			l.pos = i + 1
+			return l.newToken(TokenString, string(l.input[start:i]))
+		}
+		if c == '\\' || c < 0x20 && c != '\t' || c == 0x7f {
+			break
+		}
+	}
+	return l.readEscaped()
+}
+
+// readEscaped decodes a string's escapes into one buffer sized to its source,
+// which an escape only shortens
+func (l *Lexer) readEscaped() Token {
+	var result strings.Builder
+	result.Grow(l.stringBound())
 	escaped := false
 
 	for l.pos < len(l.input) {
@@ -351,7 +299,7 @@ func (l *Lexer) readString() Token {
 		}
 		if ch == '"' && !escaped {
 			l.advance()
-			return l.newToken(TokenString, string(result))
+			return l.newToken(TokenString, result.String())
 		}
 		if ch == '\\' && !escaped {
 			escaped = true
@@ -361,21 +309,21 @@ func (l *Lexer) readString() Token {
 		if escaped {
 			switch ch {
 			case '"':
-				result = append(result, '"')
+				result.WriteByte('"')
 			case '\\':
-				result = append(result, '\\')
+				result.WriteByte('\\')
 			case 'n':
-				result = append(result, '\n')
+				result.WriteByte('\n')
 			case 't':
-				result = append(result, '\t')
+				result.WriteByte('\t')
 			case 'r':
-				result = append(result, '\r')
+				result.WriteByte('\r')
 			case 'b':
-				result = append(result, '\b')
+				result.WriteByte('\b')
 			case 'f':
-				result = append(result, '\f')
+				result.WriteByte('\f')
 			case 'e':
-				result = append(result, 0x1b)
+				result.WriteByte(0x1b)
 			case 'x', 'u', 'U':
 				n := 2
 				switch ch {
@@ -399,29 +347,38 @@ func (l *Lexer) readString() Token {
 				if !utf8.ValidRune(code) {
 					return l.newToken(TokenError, fmt.Sprintf("invalid unicode escape: U+%X is not a scalar value", code))
 				}
-				var buf [utf8.UTFMax]byte
+				result.WriteRune(code)
-				w := utf8.EncodeRune(buf[:], code)
-				result = append(result, buf[:w]...)
 				escaped = false
 				continue // hex digits already consumed; skip trailing l.advance()
 			default:
 				// Unknown escape, invalid in every TOML version: preserve backslash and full rune
-				result = append(result, '\\')
+				result.WriteByte('\\')
-				var buf [utf8.UTFMax]byte
+				result.WriteRune(ch)
-				n := utf8.EncodeRune(buf[:], ch)
-				result = append(result, buf[:n]...)
 			}
 			escaped = false
 		} else {
-			// Get actual width of current character, don't use stale l.width
 			_, w := utf8.DecodeRune(l.input[l.pos:])
-			result = append(result, l.input[l.pos:l.pos+w]...)
+			result.Write(l.input[l.pos : l.pos+w])
 		}
 		l.advance()
 	}
 	return l.newToken(TokenError, "unterminated string")
 }
 
+// stringBound is the source length of the string at pos, through its closing
+// quote, a line's end or the input's
+func (l *Lexer) stringBound() int {
+	for i := l.pos; i < len(l.input); i++ {
+		switch c := l.input[i]; {
+		case c == '\\':
+			i++
+		case c == '"', c < 0x20 && c != '\t', c == 0x7f:
+			return i - l.pos
+		}
+	}
+	return len(l.input) - l.pos
+}
+
 func isDigit(r rune) bool {
 	return r >= '0' && r <= '9'
 }

parser.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"slices"
 	"strconv"
 	"unicode/utf8"
 )
@@ -12,10 +13,13 @@ import (
 // Parser parses TOML tokens into a map[string]any
 type Parser struct {
 	root    map[string]any
-	lexer   *Lexer
+	lexer   Lexer
-	current any // Pointer to the current map or slice of maps being populated (scope)
+	current map[string]any // the table keys are assigned to
-	// Identity set of inline-table maps (immutable per TOML spec)
+	// Inline tables a key's value made, immutable per TOML spec: a few in place,
-	frozen    map[uintptr]bool
+	// a set past them. One in an array no key can reach, so it is not held.
+	frozen    [4]uintptr
+	nfrozen   int
+	frozenSet map[uintptr]bool
 	curToken  Token
 	peekToken Token
 	// Value nesting depth (arrays / inline tables)
@@ -23,6 +27,9 @@ type Parser struct {
 	// Tables (maps) a document may create; 0 means DefaultMaxTables
 	MaxTables int
 	tables    int
+	// Elements of the arrays being parsed, an inner array's above its parent's,
+	// so each array is one allocation of its final length
+	values []any
 }
 
 // Recursion bound for parseValue -> parseArray/parseInlineTable
@@ -34,15 +41,10 @@ const maxValueDepth = 1000
 const DefaultMaxTables = 1 << 16
 
 func NewParser(input []byte) *Parser {
-	l := NewLexer(input)
+	p := &Parser{lexer: Lexer{input: input, line: 1, lean: true}, root: make(map[string]any)}
-	p := &Parser{
-		lexer:  l,
-		root:   make(map[string]any),
-		frozen: make(map[uintptr]bool),
-	}
-	p.nextToken()
-	p.nextToken()
 	p.current = p.root
+	p.nextToken()
+	p.nextToken()
 	return p
 }
 
@@ -96,7 +98,7 @@ func (p *Parser) parseStatement() error {
 	case TokenError:
 		return fmt.Errorf("lexing error line %d: %s", p.curToken.Line, p.curToken.Literal)
 	default:
-		return fmt.Errorf("unexpected token line %d: %s", p.curToken.Line, p.curToken.String())
+		return fmt.Errorf("unexpected token line %d: %s", p.curToken.Line, p.show(p.curToken))
 	}
 }
 
@@ -112,7 +114,8 @@ func (p *Parser) parseTableDeclaration() error {
 	p.nextToken() // consume [
 
 	// Parse Key (dotted)
-	keys, err := p.parseKeyParts()
+	var buf [4]string
+	keys, err := p.parseKeyParts(buf[:0])
 	if err != nil {
 		return err
 	}
@@ -172,7 +175,7 @@ func (p *Parser) setTableScope(keys []string, isArrayOfTables bool) error {
 				if val, exists := currentMap[key]; exists {
 					if m, ok := val.(map[string]any); ok {
 						// Inline tables cannot be reopened
-						if p.frozen[reflect.ValueOf(m).Pointer()] {
+						if p.isFrozen(m) {
 							return fmt.Errorf("cannot extend inline table %q at line %d", key, p.curToken.Line)
 						}
 						targetMap = m
@@ -194,7 +197,7 @@ func (p *Parser) setTableScope(keys []string, isArrayOfTables bool) error {
 			if val, exists := currentMap[key]; exists {
 				if m, ok := val.(map[string]any); ok {
 					// Inline tables cannot be extended via sub-tables
-					if p.frozen[reflect.ValueOf(m).Pointer()] {
+					if p.isFrozen(m) {
 						return fmt.Errorf("cannot extend inline table %q at line %d", key, p.curToken.Line)
 					}
 					ptr = m
@@ -219,15 +222,16 @@ func (p *Parser) setTableScope(keys []string, isArrayOfTables bool) error {
 	return nil
 }
 
-func (p *Parser) parseKeyValuePair(scope any) error {
+func (p *Parser) parseKeyValuePair(scope map[string]any) error {
 	// Parse Key (dotted allowed: a.b.c = 1)
-	keys, err := p.parseKeyParts()
+	var buf [4]string
+	keys, err := p.parseKeyParts(buf[:0])
 	if err != nil {
 		return err
 	}
 
 	if p.curToken.Type != TokenEqual {
-		return fmt.Errorf("expected '=' after key at line %d, got %s", p.curToken.Line, p.curToken.String())
+		return fmt.Errorf("expected '=' after key at line %d, got %s", p.curToken.Line, p.show(p.curToken))
 	}
 	p.nextToken() // consume =
 
@@ -235,20 +239,12 @@ func (p *Parser) parseKeyValuePair(scope any) error {
 	if err != nil {
 		return err
 	}
-
-	// Assign value to scope
 	return p.assignValue(scope, keys, val)
 }
 
-func (p *Parser) assignValue(scope any, keys []string, val any) error {
+// assignValue sets keys under scope to val, freezing val if it is an inline table
-	ptr := scope
+func (p *Parser) assignValue(scope map[string]any, keys []string, val any) error {
-
+	currentMap := scope
-	// If scope is map, easy. If scope is not map, error.
-	currentMap, ok := ptr.(map[string]any)
-	if !ok {
-		return fmt.Errorf("scope is not a map")
-	}
-
 	for i, key := range keys {
 		if i == len(keys)-1 {
 			// Final key, assign value
@@ -256,11 +252,14 @@ func (p *Parser) assignValue(scope any, keys []string, val any) error {
 				return fmt.Errorf("duplicate key %q at line %d", key, p.curToken.Line)
 			}
 			currentMap[key] = val
+			if m, ok := val.(map[string]any); ok {
+				p.freeze(m)
+			}
 		} else {
 			// Intermediate, ensure map
 			if existing, exists := currentMap[key]; exists {
 				if m, ok := existing.(map[string]any); ok {
-					if p.frozen[reflect.ValueOf(m).Pointer()] {
+					if p.isFrozen(m) {
 						return fmt.Errorf("cannot extend inline table %q at line %d", key, p.curToken.Line)
 					}
 					currentMap = m
@@ -280,12 +279,38 @@ func (p *Parser) assignValue(scope any, keys []string, val any) error {
 	return nil
 }
 
-func (p *Parser) parseKeyParts() ([]string, error) {
+// freeze marks an inline table no later key may extend
-	var keys []string
+func (p *Parser) freeze(m map[string]any) {
+	k := reflect.ValueOf(m).Pointer()
+	if p.frozenSet == nil && p.nfrozen < len(p.frozen) {
+		p.frozen[p.nfrozen] = k
+		p.nfrozen++
+		return
+	}
+	if p.frozenSet == nil {
+		p.frozenSet = make(map[uintptr]bool)
+		for _, f := range p.frozen {
+			p.frozenSet[f] = true
+		}
+	}
+	p.frozenSet[k] = true
+}
+
+func (p *Parser) isFrozen(m map[string]any) bool {
+	k := reflect.ValueOf(m).Pointer()
+	if p.frozenSet != nil {
+		return p.frozenSet[k]
+	}
+	return slices.Contains(p.frozen[:p.nfrozen], k)
+}
+
+// parseKeyParts appends a dotted key's parts to keys, which callers keep on
+// their stack
+func (p *Parser) parseKeyParts(keys []string) ([]string, error) {
 	for {
 		// Rule: Tokens identified as Numbers are forbidden as keys
 		if p.curToken.Type == TokenInteger || p.curToken.Type == TokenFloat {
-			return nil, fmt.Errorf("numeric keys are forbidden: %s", brief(p.curToken.Literal))
+			return nil, fmt.Errorf("numeric keys are forbidden: %s", brief(string(p.lexer.input[p.curToken.Pos:p.curToken.End])))
 		}
 
 		if p.curToken.Type == TokenString {
@@ -296,7 +321,7 @@ func (p *Parser) parseKeyParts() ([]string, error) {
 		}
 
 		if p.curToken.Type != TokenIdent && p.curToken.Type != TokenString {
-			return nil, fmt.Errorf("expected key, got %s", p.curToken.String())
+			return nil, fmt.Errorf("expected key, got %s", p.show(p.curToken))
 		}
 
 		if len(keys) >= maxValueDepth {
@@ -328,16 +353,18 @@ func (p *Parser) parseValue() (any, error) {
 		p.nextToken()
 		return val, nil
 	case TokenInteger:
-		val, err := p.parseInteger(p.curToken.Literal)
+		lit := p.lexer.input[p.curToken.Pos:p.curToken.End]
+		val, err := parseInteger(lit)
 		if err != nil {
-			return nil, fmt.Errorf("invalid integer %s at line %d: %w", brief(p.curToken.Literal), p.curToken.Line, numCause(err))
+			return nil, fmt.Errorf("invalid integer %s at line %d: %w", brief(string(lit)), p.curToken.Line, numCause(err))
 		}
 		p.nextToken()
 		return val, nil
 	case TokenFloat:
-		val, err := strconv.ParseFloat(p.curToken.Literal, 64)
+		lit := p.lexer.input[p.curToken.Pos:p.curToken.End]
+		val, err := strconv.ParseFloat(string(lit), 64)
 		if err != nil {
-			return nil, fmt.Errorf("invalid float %s at line %d: %w", brief(p.curToken.Literal), p.curToken.Line, numCause(err))
+			return nil, fmt.Errorf("invalid float %s at line %d: %w", brief(string(lit)), p.curToken.Line, numCause(err))
 		}
 		p.nextToken()
 		return val, nil
@@ -354,46 +381,28 @@ func (p *Parser) parseValue() (any, error) {
 		// An unquoted value may be a secret: say what it is, not what it says
 		return nil, fmt.Errorf("unquoted value at line %d: strings need quotes", p.curToken.Line)
 	}
-	return nil, fmt.Errorf("unexpected value token %s at line %d", p.curToken.String(), p.curToken.Line)
+	return nil, fmt.Errorf("unexpected value token %s at line %d", p.show(p.curToken), p.curToken.Line)
 }
 
-func (p *Parser) parseInteger(lit string) (int64, error) {
+// parseInteger reads a lexed integer: a prefixed one as strconv's base 0 does, the
-	// Handle optional leading sign
+// rest in base 10, so a leading zero is not octal. Neither keeps lit.
-	negative := false
+func parseInteger(lit []byte) (int64, error) {
-	numLit := lit
+	digits := lit
-	if len(numLit) > 0 && (numLit[0] == '+' || numLit[0] == '-') {
+	if len(digits) > 0 && (digits[0] == '+' || digits[0] == '-') {
-		negative = numLit[0] == '-'
+		digits = digits[1:]
-		numLit = numLit[1:]
 	}
-
+	if len(digits) > 2 && digits[0] == '0' && !isDigit(rune(digits[1])) {
-	var val int64
+		return strconv.ParseInt(string(lit), 0, 64)
-	var err error
-
-	if len(numLit) > 2 && numLit[0] == '0' {
-		switch numLit[1] {
-		case 'x', 'X':
-			val, err = strconv.ParseInt(signPrefix(negative)+numLit[2:], 16, 64)
-		case 'o', 'O':
-			val, err = strconv.ParseInt(signPrefix(negative)+numLit[2:], 8, 64)
-		case 'b', 'B':
-			val, err = strconv.ParseInt(signPrefix(negative)+numLit[2:], 2, 64)
-		default:
-			val, err = strconv.ParseInt(lit, 10, 64)
-			return val, err
-		}
-		if err != nil {
-			return 0, err
-		}
-		return val, nil
 	}
-
+	return strconv.ParseInt(string(lit), 10, 64)
-	val, err = strconv.ParseInt(lit, 10, 64)
-	return val, err
 }
 
 func (p *Parser) parseArray() ([]any, error) {
 	p.nextToken() // consume [
-	arr := make([]any, 0)
+	base := len(p.values)
+	if p.values == nil {
+		p.values = make([]any, 0, 8)
+	}
 
 	for p.curToken.Type != TokenRBracket {
 		if p.curToken.Type == TokenNewline {
@@ -405,7 +414,7 @@ func (p *Parser) parseArray() ([]any, error) {
 		if err != nil {
 			return nil, err
 		}
-		arr = append(arr, val)
+		p.values = append(p.values, val)
 
 		for p.curToken.Type == TokenNewline {
 			p.nextToken()
@@ -417,6 +426,10 @@ func (p *Parser) parseArray() ([]any, error) {
 		}
 	}
 	p.nextToken() // consume ]
+	arr := make([]any, len(p.values)-base)
+	copy(arr, p.values[base:])
+	clear(p.values[base:])
+	p.values = p.values[:base]
 	return arr, nil
 }
 
@@ -434,7 +447,8 @@ func (p *Parser) parseInlineTable() (map[string]any, error) {
 		}
 
 		// Parse key = value
-		keys, err := p.parseKeyParts()
+		var buf [4]string
+		keys, err := p.parseKeyParts(buf[:0])
 		if err != nil {
 			return nil, err
 		}
@@ -464,17 +478,22 @@ func (p *Parser) parseInlineTable() (map[string]any, error) {
 		}
 	}
 	p.nextToken() // consume }
-
+	// The key this is the value of freezes it (assignValue), after the dotted
-	// Mark inline table immutable. Value.Pointer for maps is documented
+	// keys inside it ({a.b = 1}) built it
-	// stable for identity comparison. Nested inline tables self-mark on return;
-	// same-table dotted assignments happen before the mark, so intra-table
-	// dotted keys ({a.b = 1}) remain unaffected.
-	p.frozen[reflect.ValueOf(m).Pointer()] = true
 	return m, nil
 }
 
-// adjacent reports b starting right after the one-byte token a
+// show is t as an error names it; a number's text, which the lexer leaves out,
-func adjacent(a, b Token) bool { return a.Line == b.Line && b.Col == a.Col+1 }
+// read back from the input
+func (p *Parser) show(t Token) string {
+	if t.Type == TokenInteger || t.Type == TokenFloat {
+		t.Literal = string(p.lexer.input[t.Pos:t.End])
+	}
+	return t.String()
+}
+
+// adjacent reports b starting where a ends
+func adjacent(a, b Token) bool { return b.Pos == a.End }
 
 // brief quotes s for an error, cut to 32 bytes: a literal can be megabytes
 func brief(s string) string {
@@ -492,13 +511,6 @@ func numCause(err error) error {
 	return err
 }
 
-func signPrefix(negative bool) string {
-	if negative {
-		return "-"
-	}
-	return ""
-}
-
 // Keep the package's decimal-key restriction independent of numeric range.
 func numericKey(s string) bool {
 	if len(s) > 0 && (s[0] == '+' || s[0] == '-') {