Go AST 改写后怎样从 CommentMap 恢复文件注释列表

Go AST 改写时，真正容易丢的往往不是语法节点，而是节点旁边的注释。比较稳妥的做法是先用 ast.NewCommentMap 建立“节点—注释组”关系，替换节点时调用 Update，删除节点后对新树调用 Filter，最后用 Comments 恢复按源码位置排序的注释列表。

官方地址：https://pkg.go.dev/go/ast#CommentMap

CommentMap 不是另一份源码，也不会自动理解你所有的改写意图。它负责维护注释和 AST 节点的关联；节点替换、节点删除和最终输出，仍然要在改写代码中明确处理。

• 建立关系：解析时使用 parser.ParseComments，再创建 CommentMap。
• 替换关系：旧节点仍然承载注释时，用 Update(old, new) 迁移。
• 恢复列表：用 Filter(newFile) 丢弃已不在新树中的映射，再调用 Comments()。

先建立 AST 与注释的关联

parser.ParseFile 只有在解析模式包含 parser.ParseComments 时，才会把注释组放进 ast.File.Comments。拿到文件节点、文件集和注释列表后，ast.NewCommentMap 会根据源码位置把注释组关联到合适的节点。它倾向于选择较大的语法节点，例如赋值语句末尾的行尾注释通常会关联到整个赋值语句。

package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
)

func buildCommentMap(src string) (*token.FileSet, *ast.File, ast.CommentMap, error) {
	// ParseComments 很关键；没有它，File.Comments 不会保存注释组。
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, "input.go", src, parser.ParseComments)
	if err != nil {
		return nil, nil, nil, err
	}

	// CommentMap 保存的是节点到注释组的关联，不是注释文本的副本。
	cmap := ast.NewCommentMap(fset, file, file.Comments)
	fmt.Printf("mapped groups: %d\\n", len(cmap.Comments()))
	return fset, file, cmap, nil
}

这里的重点是保存三个对象的生命周期：FileSet 用来解释位置，File 是当前 AST，CommentMap 是改写期间的关联表。只把 file.Comments 复制到一个切片里并不能跟踪节点替换。

替换 AST 节点时用 Update 迁移注释

例如工具要把一个旧的声明节点换成新的声明节点。新节点即使出现在新树中，也不会自动继承旧节点在 CommentMap 中的映射。此时应在替换动作附近调用 cmap.Update(oldNode, newNode)。

Update 会删除旧键，把旧节点关联的注释组追加到新键；如果新节点本来已经有注释，也会保留已有列表。这样做的好处是迁移规则和 AST 改写保持在同一个位置，后续不需要根据行号猜注释应该归谁。

func replaceDecl(file *ast.File, cmap ast.CommentMap) {
	// 只演示迁移关系；实际项目通常会按自己的规则定位旧声明。
	var oldDecl *ast.GenDecl
	for _, decl := range file.Decls {
		if candidate, ok := decl.(*ast.GenDecl); ok && candidate.Tok.String() == "var" {
			oldDecl = candidate
			break
		}
	}
	if oldDecl == nil {
		return
	}

	// 新节点应由改写逻辑构造，并保留合法的 Pos/End 范围。
	newDecl := &ast.GenDecl{
		Tok:    oldDecl.Tok,
		TokPos: oldDecl.TokPos,
		Specs:  oldDecl.Specs,
	}

	// 先在映射表中迁移注释，再把新节点放入文件声明列表。
	cmap.Update(oldDecl, newDecl)
	for i, decl := range file.Decls {
		if decl == oldDecl {
			file.Decls[i] = newDecl
			break
		}
	}
}

这段示例只强调关联迁移，未把“如何构造业务上的新声明”伪装成通用算法。实际改写时还要同步更新父节点字段，并确保新节点的位置信息适合后续打印。

用 Filter 和 Comments 恢复文件注释列表

删除节点后，旧映射表可能还保留指向已删除节点的键。直接遍历原始 CommentMap，会得到已经不属于新树的关联。Filter 会检查新 AST 中实际存在的节点，只复制仍然有效的映射，返回一个新的 CommentMap。

得到过滤后的映射后，调用 Comments() 就能恢复注释组列表。官方实现会收集映射中的全部注释组并按 Pos() 排序，因此它适合作为文件级注释列表，交给后续的格式化或导出逻辑。若只想获取某个节点的注释，则直接读取 filtered[node]，不要把所有注释再拼回一个节点。

func commentsAfterRewrite(file *ast.File, cmap ast.CommentMap) []*ast.CommentGroup {
	// Filter 以新 AST 为准，清理删除节点留下的旧映射。
	filtered := cmap.Filter(file)

	// Comments 返回所有有效注释组，并按源码位置排序。
	groups := filtered.Comments()
	for _, group := range groups {
		// CommentGroup.Text() 适合得到展示文本；它不是原始行布局的替代品。
		fmt.Println(group.Text())
	}
	return groups
}

如果目标是重新生成完整 Go 源码，注释列表只是其中一部分，还需要把新 AST 交给 go/format 或 go/printer。如果目标是做索引、文档提取或注释统计，Comments() 返回的顺序化列表通常已经足够。

边界与常见误区

为什么只改 file.Decls 还不够？

因为 CommentMap 的键是具体的 AST 节点指针。把旧节点从切片里移除，再放入新节点，并不会自动复制映射；替换动作必须显式调用 Update。

Filter 会重新判断注释应该属于哪个节点吗？

不会。Filter 的职责是保留新树中仍存在的原有映射；如果你大幅改变源码位置或节点层级，需要重新解析并调用 NewCommentMap，不要把过滤当成重新归属算法。

Comments() 返回的是原始注释文本吗？

它返回 []*ast.CommentGroup，顺序按源码位置排列。需要展示文本时调用每个分组的 Text()；需要保留精确布局时，应结合位置和打印器处理。

最小处理顺序是什么？

可以记成：解析时保留注释，创建映射；替换节点时 Update；删除或重组完成后 Filter；最终读取文件注释时调用 Comments()。这四步分别对应建立、迁移、清理和输出。

小结

Go AST 改写后的注释恢复，核心不是重新扫描注释文本，而是维护节点身份的变化。NewCommentMap 负责初始关联，Update 负责替换迁移，Filter 负责清理失效节点，Comments 负责按源码顺序导出结果。把这四个动作放在 AST 生命周期的对应位置，注释列表就不会因为一次节点重建而悄悄丢失。