Fix multi-struct rewrite loss by deferring DST rendering to run()

dkorunic · dkorunic · commit e9871011f0a9 · 2026-03-04T15:01:58.000+01:00
Store dirty *dst.File pointers instead of pre-rendered []byte buffers so
that all struct rewrites within the same file accumulate on the shared DST
before final rendering. The previous approach called decorator.Fprint inside
betteralign() and stored the result in fixOps[fn], silently overwriting any
prior rewrite for the same file; only the last modified struct per file was
preserved.

Also hoist os.Getwd() out of the inspect.Preorder closure so it is called
once per analysis pass instead of once per AST node, and propagate its error
properly rather than printing to stderr and continuing silently.

Move the dNode lookup and hasIgnoreComment check to the top of betteralign()
so ignored structs are skipped before any expensive computation.

Pre-compute architecture suffix strings in removeOtherArches to avoid
repeated string construction inside the inner loop.
diff --git a/betteralign.go b/betteralign.go
@@ -181,11 +181,20 @@ func run(pass *analysis.Pass) (any, error) {
 
 	var dFile *dst.File
 
-	applyFixesFset := make(map[string][]byte)
+	dirtyFiles := make(map[string]*dst.File)
 	testFset := make(map[string]bool)
 	generatedFset := make(map[string]bool)
 	structMetaMap := make(map[token.Pos]structMeta)
 
+	var wd string
+	if len(fExcludeDirs) > 0 || len(fExcludeFiles) > 0 {
+		var err error
+		wd, err = os.Getwd()
+		if err != nil {
+			return nil, fmt.Errorf("%v: %w", ErrPreFilterFiles, err)
+		}
+	}
+
 	inspect.Preorder(nodeFilter, func(node ast.Node) {
 		fn := pass.Fset.File(node.Pos()).Name()
 
@@ -198,11 +207,6 @@ func run(pass *analysis.Pass) (any, error) {
 		}
 
 		if len(fExcludeDirs) > 0 || len(fExcludeFiles) > 0 {
-			wd, err := os.Getwd()
-			if err != nil {
-				fmt.Fprintf(os.Stderr, "%v %s: %v", ErrPreFilterFiles, fn, err)
-				return
-			}
 			relfn, err := filepath.Rel(wd, fn)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "%v %s: %v", ErrPreFilterFiles, fn, err)
@@ -288,17 +292,20 @@ func run(pass *analysis.Pass) (any, error) {
 		}
 
 		if tv, ok := pass.TypesInfo.Types[s]; ok {
-			betteralign(pass, s, tv.Type.(*types.Struct), dec, dFile, applyFixesFset, fn)
+			betteralign(pass, s, tv.Type.(*types.Struct), dec, dFile, dirtyFiles, fn)
 		}
 	})
 
-	if !apply {
-		return nil, nil
-	}
-
-	for fn, buf := range applyFixesFset {
-		if err := applyToFile(fn, buf); err != nil {
-			fmt.Fprintf(os.Stderr, "error applying fixes to %v: %v\n", fn, err)
+	if apply {
+		for fn, df := range dirtyFiles {
+			var buf bytes.Buffer
+			if err := decorator.Fprint(&buf, df); err != nil {
+				fmt.Fprintf(os.Stderr, "failed to print %s: %v\n", fn, err)
+				continue
+			}
+			if err := applyToFile(fn, buf.Bytes()); err != nil {
+				fmt.Fprintf(os.Stderr, "error applying fixes to %v: %v\n", fn, err)
+			}
 		}
 	}
 
@@ -308,8 +315,15 @@ func run(pass *analysis.Pass) (any, error) {
 var unsafePointerTyp = types.Unsafe.Scope().Lookup("Pointer").(*types.TypeName).Type()
 
 func betteralign(pass *analysis.Pass, aNode *ast.StructType, typ *types.Struct, dec *decorator.Decorator,
-	dFile *dst.File, fixOps map[string][]byte, fn string,
+	dFile *dst.File, dirtyFiles map[string]*dst.File, fn string,
 ) {
+	dNode := dec.Dst.Nodes[aNode].(*dst.StructType)
+
+	// Skip if explicitly ignored with magic comment substring.
+	if hasIgnoreComment(dNode.Fields) {
+		return
+	}
+
 	wordSize := pass.TypesSizes.Sizeof(unsafePointerTyp)
 	maxAlign := pass.TypesSizes.Alignof(unsafePointerTyp)
 
@@ -327,13 +341,6 @@ func betteralign(pass *analysis.Pass, aNode *ast.StructType, typ *types.Struct,
 		return
 	}
 
-	dNode := dec.Dst.Nodes[aNode].(*dst.StructType)
-
-	// Skip if explicitly ignored with magic comment substring.
-	if hasIgnoreComment(dNode.Fields) {
-		return
-	}
-
 	// Flatten the ast node since it could have multiple field names per list item while
 	// *types.Struct only have one item per field.
 	// TODO: Preserve multi-named fields instead of flattening.
@@ -362,19 +369,14 @@ func betteralign(pass *analysis.Pass, aNode *ast.StructType, typ *types.Struct,
 
 	dNode.Fields.List = reordered
 
-	var buf bytes.Buffer
-	if err := decorator.Fprint(&buf, dFile); err != nil {
-		return
-	}
-
 	pass.Report(analysis.Diagnostic{
 		Pos:            aNode.Pos(),
 		End:            aNode.Pos() + token.Pos(len("struct")),
 		Message:        message,
 		SuggestedFixes: nil,
 	})
 
-	fixOps[fn] = buf.Bytes()
+	dirtyFiles[fn] = dFile
 }
 
 func optimalOrder(str *types.Struct, sizes *gcSizes) (*types.Struct, []int) {
diff --git a/betteralign_test.go b/betteralign_test.go
@@ -29,12 +29,17 @@ func removeOtherArches(paths []string) []string {
 
 	delete(arches, runtime.GOARCH)
 
+	suffixes := make([]string, 0, len(arches))
+	for arch := range arches {
+		suffixes = append(suffixes, "_"+arch+".go")
+	}
+
 	var blacklist bool
 	for _, path := range paths {
 		blacklist = false
 
-		for arch := range arches {
-			if strings.Contains(path, strings.Join([]string{"_", arch, ".go"}, "")) {
+		for _, suffix := range suffixes {
+			if strings.Contains(path, suffix) {
 				blacklist = true
 				break
 			}
