Fix whitespace handling for synthetic tokens (#731)

This PR makes a fix to the `ast/printer` to account for synthetic
tokens when checking for whitespace trivia, skipping over them
when consulting for blank space for non-synthetic tokens.

It also expands the `Edit` API for adding declarations before,
so that we can test the behaviour of adding synthetic tokens
before non-synthetic definitions/declarations.

Author: doriable

experimental/ast/edit/edit.go

@@ -90,9 +90,13 @@ type Edit struct {
 	//   - oneof, extend, group: message body
 	Insertions []ast.DeclAny
 
-	// Before is the destination anchor for KindMove: the moved decl
-	// is reinserted immediately before Before. Honored only by
-	// KindMove.
+	// Before is the destination anchor for positional inserts:
+	//   - KindAdd:  when non-zero, Insertions are placed immediately
+	//               before Before in Target's decl list (in the order
+	//               given). Before must be a current member of Target.
+	//               When zero, Insertions are appended (default).
+	//   - KindMove: the moved decl is reinserted immediately before
+	//               Before. Required.
 	Before ast.DeclAny
 }
 
@@ -128,21 +132,34 @@ func applyEdit(file *ast.File, edit Edit) error {
 	}
 }
 
-// applyAdd appends insertions to the target's decl list, validating
-// each insertion against the target container.
+// applyAdd appends insertions to the target's decl list (or inserts
+// them before edit.Before when set), validating each insertion against
+// the target container.
 func applyAdd(file *ast.File, edit Edit) error {
 	decls, container, err := targetDecls(file, edit.Target)
 	if err != nil {
 		return err
 	}
+	insertIdx := -1
+	if !edit.Before.IsZero() {
+		idx := indexOf(decls, edit.Before)
+		if idx < 0 {
+			return errors.New("before not found in target")
+		}
+		insertIdx = idx
+	}
 	for j, ins := range edit.Insertions {
 		if ins.IsZero() {
 			return fmt.Errorf("insertion[%d] is zero", j)
 		}
 		if err := validateInsertion(container, ins); err != nil {
 			return fmt.Errorf("insertion[%d]: %w", j, err)
 		}
-		seq.Append(decls, ins)
+		if insertIdx >= 0 {
+			decls.Insert(insertIdx+j, ins)
+		} else {
+			seq.Append(decls, ins)
+		}
 	}
 	return nil
 }

experimental/ast/edit/edit_test.go

@@ -15,6 +15,7 @@
 package edit_test
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 	"testing"
@@ -139,6 +140,12 @@ type editSpec struct {
 	Tag    string `yaml:"tag"`
 	Option string `yaml:"option"`
 	Value  string `yaml:"value"`
+	// Before names the existing decl in Target before which the
+	// insertion should land. Honored by `add_option`. The value is
+	// the simple name of an existing decl in the target body (or
+	// dotted path for `move_decl`). When empty, the insertion
+	// appends.
+	Before string `yaml:"before"`
 }
 
 // pendingDecls maps a dotted path to a [ast.DeclAny] that has been
@@ -184,10 +191,19 @@ func buildEdit(file *ast.File, pending pendingDecls, spec editSpec) (edit.Edit,
 			return edit.Edit{}, err
 		}
 		opt := createOptionDecl(stream, nodes, spec.Option, spec.Value)
+		var before ast.DeclAny
+		if spec.Before != "" {
+			anchor, err := resolveAnchor(file, target, spec.Before)
+			if err != nil {
+				return edit.Edit{}, fmt.Errorf("before %q: %w", spec.Before, err)
+			}
+			before = anchor
+		}
 		return edit.Edit{
 			Kind:       edit.KindAdd,
 			Target:     target,
 			Insertions: []ast.DeclAny{opt.AsAny()},
+			Before:     before,
 		}, nil
 
 	case "add_message":
@@ -271,6 +287,11 @@ func buildAdd(file *ast.File, pending pendingDecls, targetPath, name string, ins
 // method without an existing `{}` body, one is created and attached
 // so the resulting target has a body.
 func ensureOptionTarget(file *ast.File, pending pendingDecls, targetPath string) (ast.DeclAny, error) {
+	if targetPath == "" {
+		// File-level target: edit.Edit treats Target.IsZero() as the
+		// file's top-level decl list.
+		return ast.DeclAny{}, nil
+	}
 	if d, ok := pending.resolve(file, targetPath); ok {
 		return d, nil
 	}
@@ -357,6 +378,33 @@ func findDeclByPath(file *ast.File, targetPath string) (ast.DeclAny, bool) {
 	return ast.DeclAny{}, false
 }
 
+// resolveAnchor returns the decl named `name` in the target container.
+// For a file-level target (zero), searches file.Decls(); otherwise
+// searches the target def's body decls. Used to resolve `before:` in
+// add_option specs.
+func resolveAnchor(file *ast.File, target ast.DeclAny, name string) (ast.DeclAny, error) {
+	var decls seq.Indexer[ast.DeclAny]
+	if target.IsZero() {
+		decls = file.Decls()
+	} else {
+		def := target.AsDef()
+		if def.IsZero() || def.Body().IsZero() {
+			return ast.DeclAny{}, errors.New("target has no body")
+		}
+		decls = def.Body().Decls()
+	}
+	for d := range seq.Values(decls) {
+		dd := d.AsDef()
+		if dd.IsZero() {
+			continue
+		}
+		if defName(dd) == name {
+			return d, nil
+		}
+	}
+	return ast.DeclAny{}, fmt.Errorf("decl %q not found in target", name)
+}
+
 // findTopLevelDeclByName returns the file-level decl with the given
 // name.
 func findTopLevelDeclByName(file *ast.File, name string) (ast.DeclAny, bool) {

experimental/ast/edit/testdata/edits/add_options_before.yaml

@@ -0,0 +1,66 @@
+# Copyright 2020-2026 Buf Technologies, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Regression test for trivia/blankBefore handling when synthetic
+# decls are inserted ahead of existing source decls.
+#
+# Mirrors the downstream `bufformat` deprecation flow: a synthetic
+# `option deprecated = true;` is inserted at index 0 of OuterMessage's
+# body (before the existing fields/nested types) and at the file level
+# (before the first message). The printer must translate AST
+# iteration indices to source decl indices when consulting
+# trivia.blankBefore — otherwise the inserted decl shifts the index,
+# and the source-level blank-line layout drifts off by one around the
+# surrounding source decls.
+
+source: |
+  syntax = "proto3";
+
+  package test;
+
+  // Outer message.
+  message OuterMessage {
+    string name = 1;
+
+    // Nested message.
+    message NestedMessage {
+      string value = 1;
+    }
+
+    // Nested enum.
+    enum NestedEnum {
+      NESTED_ENUM_UNSPECIFIED = 0;
+      NESTED_ENUM_VALUE = 1;
+    }
+  }
+
+  // Top-level enum.
+  enum TopLevelEnum {
+    TOP_LEVEL_ENUM_UNSPECIFIED = 0;
+    TOP_LEVEL_ENUM_VALUE = 1;
+  }
+
+edits:
+  # File-scope: insert option before the first source decl after package.
+  - kind: add_option
+    target: ""
+    before: OuterMessage
+    option: deprecated
+    value: "true"
+  # Body-scope: insert option before the first field of OuterMessage.
+  - kind: add_option
+    target: OuterMessage
+    before: name
+    option: deprecated
+    value: "true"

experimental/ast/edit/testdata/edits/add_options_before.yaml.txt

@@ -0,0 +1,28 @@
+syntax = "proto3";
+
+package test;
+
+option deprecated = true;
+
+// Outer message.
+message OuterMessage {
+  option deprecated = true;
+  string name = 1;
+
+  // Nested message.
+  message NestedMessage {
+    string value = 1;
+  }
+
+  // Nested enum.
+  enum NestedEnum {
+    NESTED_ENUM_UNSPECIFIED = 0;
+    NESTED_ENUM_VALUE = 1;
+  }
+}
+
+// Top-level enum.
+enum TopLevelEnum {
+  TOP_LEVEL_ENUM_UNSPECIFIED = 0;
+  TOP_LEVEL_ENUM_VALUE = 1;
+}

experimental/ast/printer/printer.go

@@ -178,10 +178,17 @@ func (p *printer) printFile(file *ast.File) {
 		// before sorting. After sortFileDeclsForFormat, the iteration
 		// order changes but trivia.hasBlankBefore is keyed by source
 		// position, so we need to translate sorted index -> source
-		// index when consulting it.
+		// index when consulting it. Synthetic decls (no source span)
+		// map to -1 so they do not consume a slot in blankBefore.
 		srcIdxByDecl := make(map[ast.DeclAny]int, len(sourceOrder))
-		for i, d := range sourceOrder {
-			srcIdxByDecl[d] = i
+		natCounter := 0
+		for _, d := range sourceOrder {
+			if d.Span().IsZero() {
+				srcIdxByDecl[d] = -1
+				continue
+			}
+			srcIdxByDecl[d] = natCounter
+			natCounter++
 		}
 		sorted := append([]ast.DeclAny(nil), sourceOrder...)
 		sortFileDeclsForFormat(sorted)
@@ -190,7 +197,7 @@ func (p *printer) printFile(file *ast.File) {
 			if idx, ok := srcIdxByDecl[d]; ok {
 				p.declSourceIdx[i] = idx
 			} else {
-				p.declSourceIdx[i] = i
+				p.declSourceIdx[i] = -1
 			}
 		}
 		defer func() { p.declSourceIdx = nil }()
@@ -346,18 +353,55 @@ func (p *printer) appendPending(tokens []token.Token) {
 //
 // Attached trivia (leading on first token, trailing on last) is handled by
 // [printer.printDecl] via [printer.printTokenAs].
+//
+// AST decls may include synthetic decls (those without a source span,
+// e.g. inserted via the edit package). Synthetic decls do not consume
+// entries in trivia.slots or trivia.blankBefore, so iteration indices
+// are translated through scopeSourceIdx before consulting trivia.
 func (p *printer) printScopeDecls(
 	trivia detachedTrivia,
 	decls seq.Indexer[ast.DeclAny],
 	scope scopeKind,
 ) {
+	sourceIdx := p.scopeSourceIdx(decls, scope)
+	lastSrc := -1
 	for i := range decls.Len() {
-		p.emitTriviaSlot(trivia, i)
-		gap := p.declGap(decls, trivia, i, scope)
+		src := sourceIdx[i]
+		if src >= 0 {
+			for s := lastSrc + 1; s <= src; s++ {
+				p.emitTriviaSlot(trivia, s)
+			}
+			if src > lastSrc {
+				lastSrc = src
+			}
+		}
+		gap := p.declGap(decls, trivia, i, src, scope)
 		p.printDecl(decls.At(i), gap)
 	}
 
-	p.emitRemainingTrivia(trivia, decls.Len())
+	p.emitRemainingTrivia(trivia, lastSrc+1)
+}
+
+// scopeSourceIdx returns a slice mapping each AST decl iteration index
+// to its source decl index, or -1 if the decl is synthetic (has no
+// source span). For the file scope under CanonicalizeFileOrder, the
+// pre-computed mapping from printFile is reused.
+func (p *printer) scopeSourceIdx(decls seq.Indexer[ast.DeclAny], scope scopeKind) []int {
+	n := decls.Len()
+	if scope == scopeFile && p.declSourceIdx != nil && len(p.declSourceIdx) == n {
+		return p.declSourceIdx
+	}
+	out := make([]int, n)
+	counter := 0
+	for i := range n {
+		if decls.At(i).Span().IsZero() {
+			out[i] = -1
+			continue
+		}
+		out[i] = counter
+		counter++
+	}
+	return out
 }
 
 // declGap computes the gap before declaration i in a scope.
@@ -366,10 +410,15 @@ func (p *printer) printScopeDecls(
 // comments (copyright headers at file level, or comments between '{'
 // and the first member at body level). For subsequent declarations, it
 // determines whether a blank line or regular newline separates them.
+//
+// srcIdx is the source decl index for decls.At(i), or -1 if the decl
+// is synthetic (no source span). Synthetic decls never consult
+// trivia.hasBlankBefore because they have no source position to
+// compare against.
 func (p *printer) declGap(
 	decls seq.Indexer[ast.DeclAny],
 	trivia detachedTrivia,
-	i int,
+	i, srcIdx int,
 	scope scopeKind,
 ) gapStyle {
 	if i == 0 {
@@ -393,9 +442,7 @@ func (p *printer) declGap(
 	// Within sorted sections (imports, file-level options),
 	// canonicalization reorders entries, so hasBlankBefore (keyed
 	// by source position) doesn't line up with the sorted output;
-	// suppress source blanks there. For body decls, translate the
-	// sorted index back to the source index before consulting
-	// hasBlankBefore.
+	// suppress source blanks there.
 	if scope == scopeFile {
 		prev, curr := rankDecl(decls.At(i-1)), rankDecl(decls.At(i))
 		if prev != curr && (curr == rankImport || curr == rankOption) {
@@ -407,18 +454,14 @@ func (p *printer) declGap(
 		if curr == rankImport || curr == rankOption {
 			return gapNewline
 		}
-		srcIdx := i
-		if p.declSourceIdx != nil && i < len(p.declSourceIdx) {
-			srcIdx = p.declSourceIdx[i]
-		}
-		if trivia.hasBlankBefore(srcIdx) {
+		if srcIdx >= 0 && trivia.hasBlankBefore(srcIdx) {
 			return gapBlankline
 		}
 		return gapNewline
 	}
 
 	// Body level: preserve blank lines from the original source.
-	if trivia.hasBlankBefore(i) {
+	if srcIdx >= 0 && trivia.hasBlankBefore(srcIdx) {
 		return gapBlankline
 	}