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

Assisted-by: GLM 5.3

document_test.go

@@ -4,6 +4,7 @@
 package interpres
 
 import (
+	"reflect"
 	"slices"
 	"strings"
 	"testing"
@@ -416,3 +417,143 @@ func TestDocumentEditPipeline(t *testing.T) {
 		}
 	})
 }
+
+// TestMarshalDocumentRoundTrips pins that a parsed document written back
+// re-parses to the same tree: arrays of tables keep exactly one header per
+// element, dotted keys hold their line position without swallowing the keys
+// after them, inline tables inside value arrays keep their written order,
+// and comments travel with their statements.
+func TestMarshalDocumentRoundTrips(t *testing.T) {
+	tests := []struct {
+		name string
+		src  string
+	}{
+		{"array of tables", "[[items]]\nname = \"a\"\n\n[[items]]\nname = \"b\"\n"},
+		{"array of tables with comments", "# about items\n[[items]] # first\nname = \"a\"\n"},
+		{"dotted key before a later key", "a.b = 1\nc = 2\n"},
+		{"dotted keys grouped", "a.b = 1\na.c = 2\nd = 3\n"},
+		{"dotted key with a nested leaf", "a.b.c = 1\nz = 2\n"},
+		{"header section after a dotted key", "a.b = 1\n\n[a.x]\ny = 2\n"},
+		{"inline tables in a value array keep order", "arr = [{y = 1, x = 2}, {second = true, first = false}]\n"},
+		{"nested array of tables", "[[items]]\nn = 1\n\n[items.sub]\nk = \"v\"\n"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			doc, err := Parse([]byte(tt.src))
+			if err != nil {
+				t.Fatalf("Parse: %v", err)
+			}
+			out, err := Marshal(doc)
+			if err != nil {
+				t.Fatalf("Marshal: %v", err)
+			}
+			reparsed, err := Parse(out)
+			if err != nil {
+				t.Fatalf("re-parse of %q: %v", out, err)
+			}
+			if !reflect.DeepEqual(doc.Map(), reparsed.Map()) {
+				t.Errorf("round trip changed the tree:\nin:  %#v\nout: %#v", doc.Map(), reparsed.Map())
+			}
+			if got, want := reparsed.Root().Keys(), doc.Root().Keys(); !slices.Equal(got, want) {
+				t.Errorf("root keys = %v, want %v", got, want)
+			}
+		})
+	}
+}
+
+// TestMarshalDocumentArrayComments pins where the comments of an array of
+// tables land: above and beside the [[header]] itself.
+func TestMarshalDocumentArrayComments(t *testing.T) {
+	doc, err := Parse([]byte("# element one\n[[items]] # trailing\nname = \"a\"\n"))
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+	out, err := Marshal(doc)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+	want := "# element one\n[[items]] # trailing\nname = \"a\"\n"
+	if string(out) != want {
+		t.Errorf("output = %q, want %q", out, want)
+	}
+}
+
+// TestTableSetReplacesTableNode pins that Set over a key holding a table
+// rebuilds the node, so the new map's keys are the ones written.
+func TestTableSetReplacesTableNode(t *testing.T) {
+	doc, err := Parse([]byte("[cache]\nz = 1\n"))
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+	doc.Set("cache", map[string]any{"a": int64(2)})
+	out, err := Marshal(doc)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+	want := "[cache]\na = 2\n"
+	if string(out) != want {
+		t.Errorf("output = %q, want %q", out, want)
+	}
+}
+
+// TestTableSetNestedArraysOfTables pins that a value set through the edit API
+// carries its arrays of tables into the header form.
+func TestTableSetNestedArraysOfTables(t *testing.T) {
+	doc, err := Parse([]byte("x = 1\n"))
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+	doc.Set("t", map[string]any{"items": []map[string]any{{"n": int64(1)}, {"n": int64(2)}}})
+	out, err := Marshal(doc)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+	if !strings.Contains(string(out), "[[t.items]]") {
+		t.Errorf("output = %q, want the array of tables under a header", out)
+	}
+}
+
+// TestTableSetCyclicMapErrors pins that a cyclic map set through the edit API
+// reaches the depth limit instead of the stack.
+func TestTableSetCyclicMapErrors(t *testing.T) {
+	doc, err := Parse([]byte("x = 1\n"))
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+	m := map[string]any{}
+	m["self"] = m
+	doc.Set("cyclic", m)
+	if _, err := Marshal(doc); err == nil || !strings.Contains(err.Error(), "nests deeper") {
+		t.Errorf("err = %v, want the depth-limit complaint", err)
+	}
+}
+
+// TestDocumentNilSafety pins that the nil document answers its readers
+// instead of panicking, the contract Root already carries.
+func TestDocumentNilSafety(t *testing.T) {
+	var doc *Document
+	if doc.Map() != nil {
+		t.Errorf("Map = %v", doc.Map())
+	}
+	if doc.Footer() != nil {
+		t.Errorf("Footer = %v", doc.Footer())
+	}
+	doc.SetFooter([]string{"x"})
+	if e, ok := doc.Get("k"); e != nil || ok {
+		t.Errorf("Get = %v, %v", e, ok)
+	}
+	if _, ok := doc.GetString("k"); ok {
+		t.Error("GetString on a nil document reports a value")
+	}
+	if _, ok := doc.GetTable("k"); ok {
+		t.Error("GetTable on a nil document reports a value")
+	}
+	doc.Set("k", 1)
+	doc.Delete("k")
+	if keys := doc.Root().Keys(); keys != nil {
+		t.Errorf("Keys = %v", keys)
+	}
+	if doc.Root().Entries() != nil {
+		t.Errorf("Entries = %v", doc.Root().Entries())
+	}
+}