如何在Go AST正确位置插入代码并维持注释位置？

DevOps 小满

2026-6-17

问题：Go自修改程序中注释位置异常

我写了一个可自修改的Go程序，运行时读取源码，用go/parser解析AST，定位目标位置后插入新的go/ast.Expr，再用go/format格式化新AST。但插入后注释位置不符合预期：

原代码片段：

var passwd = []*PwEntry{}
// line 1 after
// line 2 after

我找到对应的ast.CompositeLit并将新Expr插入Elts，预期新代码在大括号内，注释留在大括号后，但格式化后注释跑到了代码中间：

var passwd = []*PwEntry{{Login:
// line 1 after
// line 2 after
"abc", Salt: []byte{214, 194, 249, 8, 11, 40, 37, 195, 65, 130, 142, 86, 68, 78, 185, 33}, Algo: 1, Hash: []byte{70, 8, 114, 178, 255, 193, 204, 112, 83, 209, 249, 153, 213, 253, 151, 47, 49, 99, 133, 139, 203, 184, 243, 15, 203, 16, 6, 235, 29, 236, 249, 57}}}

完整示例代码：

package main

import (
    "crypto/rand"
    "fmt"
    "github.com/peterh/liner"
    "go/ast"
    "go/format"
    "go/parser"
    "go/token"
    "golang.org/x/crypto/argon2"
    "os"
    "runtime"
)

//======================================================================
// Terminal input
//----------------------------------------------------------------------

func ReadLogin(term *liner.State) string {
    login, err := term.Prompt("Login: ")
    if err != nil {
        panic(err)
    }
    return login
}

func ReadPassword(term *liner.State) string {
    password, err := term.PasswordPrompt("Password: ")
    if err != nil {
        panic(err)
    }
    return password
}

//======================================================================
// Mutate source file
//----------------------------------------------------------------------

type Source struct {
    FileName string
    FileSet  *token.FileSet
    AstFile  *ast.File
}

func NewSource() (this *Source) {
    this = new(Source)
    _, fn, _, ok := runtime.Caller(0)
    if !ok {
        panic("Can not read source.")
    }
    this.FileName = fn
    this.FileSet = token.NewFileSet()
    return
}

func (this *Source) Read() {
    var err error
    this.AstFile, err = parser.ParseFile(
        this.FileSet, this.FileName, nil, parser.ParseComments)
    if err != nil {
        panic(err)
    }
}

func (this *Source) Dump() {
    ast.Print(this.FileSet, this.AstFile)
}

// Search for top level decleration "var passwd = []*PwEntry{...}".
func (this *Source) FindPasswdDecl() *ast.CompositeLit {
    for _, node := range this.AstFile.Decls {
        switch decl := node.(type) {
        case *ast.GenDecl:
            if decl.Tok == token.VAR {
                // fmt.Println("found: var")
                for _, spec := range decl.Specs {
                    switch vspec := spec.(type) {
                    case *ast.ValueSpec:
                        if len(vspec.Names) == 1 && vspec.Names[0].Name == "passwd" {
                            // fmt.Println("found: passwd")
                            if len(vspec.Values) == 1 {
                                switch cl := vspec.Values[0].(type) {
                                case *ast.CompositeLit:
                                    switch clt := cl.Type.(type) {
                                    case *ast.ArrayType:
                                        // fmt.Println("found: []")
                                        switch se := clt.Elt.(type) {
                                        case *ast.StarExpr:
                                            // fmt.Println("found: *")
                                            switch sx := se.X.(type) {
                                            case *ast.Ident:
                                                if sx.Name == "PwEntry" {
                                                    // fmt.Println("found: PwEntry")
                                                    return cl
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return nil
}

func AstString(s string) ast.Expr {
    return &ast.BasicLit{
        Kind: token.STRING,
        Value: fmt.Sprintf("%q", s)}
}

func AstInt(i int) ast.Expr {
    return &ast.BasicLit{
      Kind: token.INT,
        Value: fmt.Sprintf("%d", i)}
}

func AstByteArray(ba []byte) ast.Expr {
    cl := &ast.CompositeLit{
        Type: &ast.ArrayType{
            Elt: ast.NewIdent("byte")}}
    cl.Elts = make([]ast.Expr, len(ba))
    for i, b := range ba {
        cl.Elts[i] = AstInt(int(b))
    }
    return cl
}

func AstPwEntryElts(entry *PwEntry) []ast.Expr {
    return []ast.Expr{
        &ast.CompositeLit{
            Elts: []ast.Expr{
                &ast.KeyValueExpr{
                    Key: ast.NewIdent("Login"),
                    Value: AstString(entry.Login)},
                &ast.KeyValueExpr{
                    Key: ast.NewIdent("Salt"),
                    Value: AstByteArray(entry.Salt)},
                &ast.KeyValueExpr{
                    Key: ast.NewIdent("Algo"),
                    Value: AstInt(entry.Algo)},
                &ast.KeyValueExpr{
                    Key: ast.NewIdent("Hash"),
                    Value: AstByteArray(entry.Hash)}}}}
}

func (this *Source) SetPwEntry(entry *PwEntry) {
    cl := this.FindPasswdDecl()
    if cl.Elts == nil {
        cl.Elts = AstPwEntryElts(entry)
    } else {
        for i, elt := range cl.Elts {
            switch nelt := elt.(type) {
            case *ast.CompositeLit:
                if len(nelt.Elts) == 4 {
                    fmt.Println("found: 4 in %d", i)
                    return
                }
            }
        }
    }
}

func (this *Source) Write() {
    output, err := os.Create("tmp.go")
    if err != nil {
        panic(err)
    }
    err = format.Node(output, this.FileSet, this.AstFile)
    if err != nil {
        panic(err)
    }
}

//======================================================================
// Password entry
//----------------------------------------------------------------------

const SaltLength = 16
const (
    AlgoNil = iota
    AlgoArgon2id
)

type PwEntry struct {
    Login string
    Salt  []byte
    Algo  int
    Hash  []byte
}

func NewPwEntry() *PwEntry {
    this := new(PwEntry)
    this.Salt = make([]byte, SaltLength)
    _, err := rand.Read(this.Salt)
    if err != nil {
        panic(err)
    }
    return this
}

func (this *PwEntry) SetPassword(password []byte) {
    this.Algo = AlgoArgon2id
    this.Hash = argon2.IDKey(password, this.Salt, 1, 64*1024, 4, 32)
}

func (this *PwEntry) AlgoName() string {
    switch this.Algo {
    case AlgoNil:
        return "nil"
    case AlgoArgon2id:
        return "Argon2id"
    default:
        return "undefined"
    }
}

func (this *PwEntry) String() string {
    return fmt.Sprintf(`{"%s":"%s","%s":"%x","%s":"%s","%s":"%x"}`,
        "login", this.Login,
        "salt", this.Salt,
        "algo", this.AlgoName(),
        "hash", this.Hash,
    )
}

//======================================================================
// Password entries
//----------------------------------------------------------------------

var passwd = []*PwEntry{}
// line 1 after
// line 2 after

func SearchPwEntry(login string) *PwEntry {
    for _, entry := range passwd {
        if entry.Login == login {
            return entry
        }
    }
    return nil
}

func GetPwEntry(login string) *PwEntry {
    entry := SearchPwEntry(login)
    if entry == nil {
        entry = NewPwEntry()
        entry.Login = login
    }
    return entry
}

//======================================================================
// Main
//----------------------------------------------------------------------

func Passwd() {
    term := liner.NewLiner()
    defer term.Close()
    src := NewSource()
    defer src.Write()
    src.Read()
    //src.Dump()
    login := ReadLogin(term)
    entry := GetPwEntry(login)
    entry.SetPassword([]byte(ReadPassword(term)))
    fmt.Println(entry)
    src.SetPwEntry(entry)
}

func main() {
    Passwd()
}

解决方案

问题根源是Go的AST注释关联逻辑：原代码中注释绑定在var passwd声明语句（ast.GenDecl）的末尾，而非独立节点。修改CompositeLit的Elts后，go/format重新排版时，会错误地将注释关联到新插入的表达式内部。

要解决这个问题，需确保注释始终绑定到正确的AST节点，步骤如下：

定位注释所属的根节点：
解析源码时，parser.ParseComments会收集所有注释到AstFile.Comments，通过ast.CommentMap可建立注释与AST节点的关联。需要找到绑定在var passwd对应的ast.GenDecl节点后的注释。
修改AST后重新关联注释：
插入新表达式后，将原本绑定在CompositeLit或声明语句后的注释，重新绑定到整个ast.GenDecl节点，确保格式化时注释不会被插入表达式中间。

修改后的关键代码示例：

func (this *Source) SetPwEntry(entry *PwEntry) {
    cl := this.FindPasswdDecl()
    // 找到包含passwd变量声明的GenDecl节点
    var genDecl *ast.GenDecl
    for _, node := range this.AstFile.Decls {
        if d, ok := node.(*ast.GenDecl); ok {
            for _, spec := range d.Specs {
                if vs, ok := spec.(*ast.ValueSpec); ok {
                    if len(vs.Names) == 1 && vs.Names[0].Name == "passwd" {
                        genDecl = d
                        break
                    }
                }
            }
            if genDecl != nil {
                break
            }
        }
    }

    if cl.Elts == nil {
        cl.Elts = AstPwEntryElts(entry)
    } else {
        for i, elt := range cl.Elts {
            switch nelt := elt.(type) {
            case *ast.CompositeLit:
                if len(nelt.Elts) == 4 {
                    fmt.Printf("found: 4 in %d\n", i)
                    return
                }
            }
        }
    }

    // 重新关联注释到GenDecl节点
    if genDecl != nil {
        commentMap := ast.NewCommentMap(this.FileSet, this.AstFile, this.AstFile.Comments)
        // 获取原本绑定的注释
        comments := commentMap.Filter(genDecl).Comments()
        // 移除与CompositeLit的关联
        delete(commentMap, cl)
        // 将注释重新绑定到GenDecl
        commentMap[genDecl] = append(commentMap[genDecl], comments...)
        // 更新AstFile的注释集合
        this.AstFile.Comments = commentMap.Comments()
    }
}

此外，也可以通过调整CompositeLit的End位置标记，让格式化工具识别表达式的结束边界，但通过注释映射管理归属的方式更可靠。最后调用format.Node时，正确的注释映射会确保注释出现在声明语句末尾，符合预期。

内容的提问来源于stack exchange，提问作者ceving