feat: add the Document edit pipeline with comment-preserving write

Assisted-by: GLM 5.3 Flash

document_test.go

@@ -288,27 +288,131 @@ func TestParseMapIsTheValueTree(t *testing.T) {
 	}
 }
 
-func TestMarshalRejectsDocument(t *testing.T) {
-	// A Document is not a value to marshal: its order and comments would be
-	// dropped, and a struct walk would silently write nothing at all.
-	doc, err := Parse([]byte("a = 1\n"))
+func TestMarshalDocument(t *testing.T) {
+	// A Document writes back: the keys in written order, the comments above
+	// the lines and headers they belonged to, and inline tables inline again.
+	doc, err := Parse([]byte("# leading\na = 1 # trailing\n\n[t]\nb = \"x\"\n\ninline = { n = 1 }\n"))
 	if err != nil {
 		t.Fatalf("parse: %v", err)
 	}
-	if _, err := Marshal(doc); err == nil {
-		t.Fatal("expected an error for a Document")
-	} else if !strings.Contains(err.Error(), "Map()") {
-		t.Errorf("err = %v, want it to point at Map()", err)
-	}
-	if _, err := Marshal(*doc); err == nil {
-		t.Fatal("expected an error for a Document value")
-	}
-	// The tree marshals, which is the way through.
-	out, err := Marshal(doc.Map())
+	out, err := Marshal(doc)
 	if err != nil {
-		t.Fatalf("marshal of the tree: %v", err)
+		t.Fatalf("marshal of a Document: %v", err)
 	}
-	if want := "a = 1\n"; string(out) != want {
-		t.Errorf("output = %q, want %q", out, want)
+	want := "# leading\na = 1 # trailing\n\n[t]\nb = \"x\"\ninline = {n = 1}\n"
+	if string(out) != want {
+		t.Errorf("output:\n%q\nwant:\n%q", out, want)
+	}
+	// The written document parses back to the same values.
+	re, err := Parse(out)
+	if err != nil {
+		t.Fatalf("re-parse: %v", err)
+	}
+	if got := re.Map()["a"]; got != int64(1) {
+		t.Errorf("a = %#v", got)
+	}
+	if _, err := Marshal(*doc); err != nil {
+		t.Errorf("marshal of a Document value: %v", err)
 	}
 }
+
+func TestDocumentEditPipeline(t *testing.T) {
+	doc, err := Parse([]byte("host = \"db\"\nport = 5432\n\n# The cache section\ntimeout = 1.5\n"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Run("typed getters", func(t *testing.T) {
+		if s, ok := doc.GetString("host"); !ok || s != "db" {
+			t.Errorf("host = %q, %v", s, ok)
+		}
+		if i, ok := doc.GetInt("port"); !ok || i != 5432 {
+			t.Errorf("port = %d, %v", i, ok)
+		}
+		if f, ok := doc.GetFloat("timeout"); !ok || f != 1.5 {
+			t.Errorf("timeout = %g, %v", f, ok)
+		}
+		if _, ok := doc.GetBool("host"); ok {
+			t.Error("host claimed as bool")
+		}
+	})
+	t.Run("set keeps the position and the comments", func(t *testing.T) {
+		doc.Set("port", int64(9090))
+		if got := doc.Root().Keys(); !slices.Equal(got, []string{"host", "port", "timeout"}) {
+			t.Fatalf("keys = %v", got)
+		}
+		out, err := Marshal(doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !strings.Contains(string(out), "port = 9090") {
+			t.Errorf("output missing the new value:\n%s", out)
+		}
+	})
+	t.Run("a new key joins the end", func(t *testing.T) {
+		doc.Set("lang", "cs")
+		if got := doc.Root().Keys(); !slices.Equal(got, []string{"host", "port", "timeout", "lang"}) {
+			t.Fatalf("keys = %v", got)
+		}
+	})
+	t.Run("a set table keeps an order of its own", func(t *testing.T) {
+		sub := map[string]any{"z": int64(1), "a": int64(2)}
+		doc.Set("cache", sub)
+		out, err := Marshal(doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !strings.Contains(string(out), "[cache]\na = 2\nz = 1\n") {
+			t.Errorf("output missing the new table in order:\n%s", out)
+		}
+	})
+	t.Run("delete removes the key", func(t *testing.T) {
+		doc.Delete("lang")
+		if _, ok := doc.Get("lang"); ok {
+			t.Fatal("lang survived Delete")
+		}
+		out, err := Marshal(doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if strings.Contains(string(out), "lang") {
+			t.Errorf("output still names lang:\n%s", out)
+		}
+	})
+	t.Run("UnmarshalDocument decodes without reparsing", func(t *testing.T) {
+		type Cfg struct {
+			Host string `toml:"host"`
+			Port int    `toml:"port"`
+		}
+		var cfg Cfg
+		if err := UnmarshalDocument(doc, &cfg); err != nil {
+			t.Fatal(err)
+		}
+		if cfg.Host != "db" || cfg.Port != 9090 {
+			t.Errorf("decoded %+v", cfg)
+		}
+	})
+	t.Run("comments survive the round trip", func(t *testing.T) {
+		src := "# header comment\n[a]\n# key comment\nb = 2\n"
+		doc, err := Parse([]byte(src))
+		if err != nil {
+			t.Fatal(err)
+		}
+		out, err := Marshal(doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for _, want := range []string{"# header comment", "[a]", "# key comment", "b = 2"} {
+			if !strings.Contains(string(out), want) {
+				t.Errorf("output missing %q:\n%s", want, out)
+			}
+		}
+	})
+	t.Run("a nil document refuses to decode", func(t *testing.T) {
+		var cfg struct {
+			A int `toml:"a"`
+		}
+		if err := UnmarshalDocument(nil, &cfg); err == nil {
+			t.Error("UnmarshalDocument(nil) succeeded, want an error")
+		}
+	})
+}