Golang代码生成问题：使用go/ast添加函数文档注释位置错误

在errmsg目录执行go generate时，生成的代码中所有Get开头函数的注释都集中在文件顶部，无法对应到各自的函数上方。尝试用go/ast包的CommentMap调整注释位置，但未解决问题。

当前生成的代码 #

// Code generated by gen_code. DO NOT EDIT.
// 提示内容："成功"
// 状态码：10000
// 提示内容："失败"
// 状态码：10001
package errmsg

const (
    TypeCodeSuccess = 10000
    TypeCodeFailed  = 10001
)

func GetSuccess() Resp {
    return Resp{Code: TypeCodeSuccess, Msg: TypeMsgSuccess}
}

func GetFailed() Resp {
    return Resp{Code: TypeCodeFailed, Msg: TypeMsgFailed}
}

尝试的处理代码片段 #

commentMap := ast.NewCommentMap(fset, file, file.Comments)
ast.Inspect(file, func(node ast.Node) bool {
    switch x := node.(type) {
    case *ast.FuncDecl:
        if strings.HasPrefix(x.Name.Name, "Get") {
            commentText := x.Doc.List[0].Text
            fmt.Println()
            fmt.Printf("comment:%s", commentText)
            commentMap[x] = []*ast.CommentGroup{{List: []*ast.Comment{{Text: commentText, Slash: token.Pos(int(x.Pos() - 1))}}}}
        }
    }
    return true
})

ast.Print(fset, file)
//os.Exit(1)
// 保存已分配的错误码
err = saveLastErrorCode(projectName, lastCode)
if err != nil {
    return err

期望生成的代码 #

// Code generated by gen_code. DO NOT EDIT.
package errmsg

const (
    TypeCodeSuccess = 10000
    TypeCodeFailed  = 10001
)

// 提示内容："成功"
// 状态码：10000
func GetSuccess() Resp {
    return Resp{Code: TypeCodeSuccess, Msg: TypeMsgSuccess}
}


// 提示内容："失败"
// 状态码：10001
func GetFailed() Resp {
    return Resp{Code: TypeCodeFailed, Msg: TypeMsgFailed}
}

方案一：生成函数时直接绑定注释（推荐） #

不要等到生成完整AST后再调整注释，在创建*ast.FuncDecl节点的阶段，直接将对应注释赋值给函数的Doc字段，这样生成的代码天然会把注释放在函数上方。

示例代码（适配你的生成逻辑）：

import (
    "go/ast"
    "strconv"
    "strings"
)

// 假设你有错误信息的结构体
type ErrDef struct {
    Code        int
    Msg         string
    ConstName   string // 比如TypeCodeSuccess
    FuncName    string // 比如GetSuccess
}

// 遍历错误定义生成函数
for _, errDef := range errDefs {
    // 构建函数对应的注释组
    commentGroup := &ast.CommentGroup{
        List: []*ast.Comment{
            {Text: "// 提示内容：\"" + errDef.Msg + "\""},
            {Text: "// 状态码：" + strconv.Itoa(errDef.Code)},
        },
    }

    // 创建函数声明
    funcDecl := &ast.FuncDecl{
        Doc:  commentGroup, // 直接绑定注释到函数
        Name: ast.NewIdent(errDef.FuncName),
        Type: &ast.FuncType{
            Results: &ast.FieldList{
                List: []*ast.Field{
                    {Type: ast.NewIdent("Resp")},
                },
            },
        },
        // 生成函数体
        Body: &ast.BlockStmt{
            List: []ast.Stmt{
                &ast.ReturnStmt{
                    Results: []ast.Expr{
                        &ast.CompositeLit{
                            Type: ast.NewIdent("Resp"),
                            Elts: []ast.Expr{
                                &ast.KeyValueExpr{
                                    Key:   ast.NewIdent("Code"),
                                    Value: ast.NewIdent(errDef.ConstName),
                                },
                                &ast.KeyValueExpr{
                                    Key:   ast.NewIdent("Msg"),
                                    Value: ast.NewIdent("TypeMsg" + strings.TrimPrefix(errDef.ConstName, "TypeCode")),
                                },
                            },
                        },
                    },
                },
            },
        },
    }

    // 将函数添加到文件声明中
    file.Decls = append(file.Decls, funcDecl)
}

方案二：事后调整AST注释（针对已生成的AST） #

如果必须在生成AST后调整注释，需要完成以下步骤：

1. 先将文件顶部的注释按函数分组（比如每个函数对应两行注释）
2. 从file.Comments中移除这些待移动的注释
3. 为每个Get开头的函数设置Doc字段，或更新CommentMap并应用修改

修正后的处理代码：

import (
    "go/ast"
    "go/token"
    "strings"
)

// 先收集文件顶部的注释组（假设前两行是GetSuccess的，接下来两行是GetFailed的）
var funcComments []*ast.CommentGroup
var remainingComments []*ast.CommentGroup

// 跳过自动生成的注释，处理后面的错误注释
skipAutoGen := false
for _, cg := range file.Comments {
    if strings.Contains(cg.Text(), "Code generated by") {
        skipAutoGen = true
        remainingComments = append(remainingComments, cg)
        continue
    }
    if skipAutoGen {
        // 每两行组成一个函数的注释组
        if len(funcComments) == 0 || len(funcComments[len(funcComments)-1].List) == 2 {
            funcComments = append(funcComments, &ast.CommentGroup{})
        }
        lastCG := funcComments[len(funcComments)-1]
        lastCG.List = append(lastCG.List, cg.List...)
    } else {
        remainingComments = append(remainingComments, cg)
    }
}

// 更新文件的Comments，只保留自动生成的注释
file.Comments = remainingComments

// 遍历函数，绑定对应注释
commentIdx := 0
ast.Inspect(file, func(node ast.Node) bool {
    if fd, ok := node.(*ast.FuncDecl); ok && strings.HasPrefix(fd.Name.Name, "Get") && commentIdx < len(funcComments) {
        fd.Doc = funcComments[commentIdx]
        commentIdx++
    }
    return true
})

// 重新生成CommentMap并应用（可选，确保注释关联正确）
commentMap := ast.NewCommentMap(fset, file, file.Comments)
commentMap.Apply(file, true)

关键注意点 #

• 不要忽略commentMap.Apply(file, true)：如果使用CommentMap修改映射，必须调用此方法将修改同步到AST节点
• 避免重复注释：调整时要确保原注释从文件顶部移除，否则会同时出现两次
• 优先在生成节点时绑定注释：这是最直接且不易出错的方式，不需要事后调整AST

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