将项目添加到Go AST后，注释乱序

Question

以下测试尝试使用AST将字段添加到结构中。 字段已正确添加，但注释未按顺序添加。 我收集了可能需要手动指定的位置，但是到目前为止，我已经找到了答案。 测试失败：http://play.golang.org/p/RID4N30FZK 这是代码：

package generator

import (
    "bytes"
    "fmt"
    "go/ast"
    "go/parser"
    "go/printer"
    "go/token"
    "testing"
)

func TestAst(t *testing.T) {

    source := `package a

// B comment
type B struct {
    // C comment
    C string
}`

    fset := token.NewFileSet()
    file, err := parser.ParseFile(fset, "", []byte(source), parser.ParseComments)
    if err != nil {
        t.Error(err)
    }

    v := &visitor{
        file: file,
    }
    ast.Walk(v, file)

    var output []byte
    buf := bytes.NewBuffer(output)
    if err := printer.Fprint(buf, fset, file); err != nil {
        t.Error(err)
    }

    expected := `package a

// B comment
type B struct {
    // C comment
    C string
    // D comment
    D int
    // E comment
    E float64
}
`

    if buf.String() != expected {
        t.Error(fmt.Sprintf("Test failed. Expected:
%s
Got:
%s", expected, buf.String()))
    }

    /*
    actual output = `package a

// B comment
type B struct {
    // C comment
    // D comment
    // E comment
    C   string
    D   int
    E   float64
}
`
    */

}

type visitor struct {
    file *ast.File
}

func (v *visitor) Visit(node ast.Node) (w ast.Visitor) {

    if node == nil {
        return v
    }

    switch n := node.(type) {
    case *ast.GenDecl:
        if n.Tok != token.TYPE {
            break
        }
        ts := n.Specs[0].(*ast.TypeSpec)
        if ts.Name.Name == "B" {
            fields := ts.Type.(*ast.StructType).Fields
            addStructField(fields, v.file, "int", "D", "D comment")
            addStructField(fields, v.file, "float64", "E", "E comment")
        }
    }

    return v
}

func addStructField(fields *ast.FieldList, file *ast.File, typ string, name string, comment string) {
    c := &ast.Comment{Text: fmt.Sprint("// ", comment)}
    cg := &ast.CommentGroup{List: []*ast.Comment{c}}
    f := &ast.Field{
        Doc:   cg,
        Names: []*ast.Ident{ast.NewIdent(name)},
        Type:  ast.NewIdent(typ),
    }
    fields.List = append(fields.List, f)
    file.Comments = append(file.Comments, cg)
}

Answer 1

I believe I have gotten it to work. As stated in my comment above, the main points required are:

1. Specifically set the buffer locations including the Slash and NamePos
2. Use token.File.AddLine to add new lines at specific offsets (calculated using the positions from item 1)
3. Overallocate the source buffer so token.File.Position (used by printer.Printer and token.File.Addline don't fail range checks on the source buffer

Code:

package main

import (
    "bytes"
    "fmt"
    "go/ast"
    "go/parser"
    "go/printer"
    "go/token"
    "testing"
)

func main() {
    tests := []testing.InternalTest{{"TestAst", TestAst}}
    matchAll := func(t string, pat string) (bool, error) { return true, nil }
    testing.Main(matchAll, tests, nil, nil)
}

func TestAst(t *testing.T) {

    source := `package a

// B comment
type B struct {
    // C comment
    C string
}`

    buffer := make([]byte, 1024, 1024)
    for idx,_ := range buffer {
        buffer[idx] = 0x20
    }
    copy(buffer[:], source)
    fset := token.NewFileSet()
    file, err := parser.ParseFile(fset, "", buffer, parser.ParseComments)
    if err != nil {
        t.Error(err)
    }

    v := &visitor{
        file: file,
        fset: fset,
    }
    ast.Walk(v, file)

    var output []byte
    buf := bytes.NewBuffer(output)
    if err := printer.Fprint(buf, fset, file); err != nil {
        t.Error(err)
    }

    expected := `package a

// B comment
type B struct {
    // C comment
    C   string
    // D comment
    D   int
    // E comment
    E   float64
}
`
    if buf.String() != expected {
        t.Error(fmt.Sprintf("Test failed. Expected:
%s
Got:
%s", expected, buf.String()))
    }

}

type visitor struct {
    file *ast.File
    fset *token.FileSet
}

func (v *visitor) Visit(node ast.Node) (w ast.Visitor) {

    if node == nil {
        return v
    }

    switch n := node.(type) {
    case *ast.GenDecl:
        if n.Tok != token.TYPE {
            break
        }
        ts := n.Specs[0].(*ast.TypeSpec)
        if ts.Name.Name == "B" {
            fields := ts.Type.(*ast.StructType).Fields
            addStructField(v.fset, fields, v.file, "int", "D", "D comment")
            addStructField(v.fset, fields, v.file, "float64", "E", "E comment")
        }
    }

    return v
}

func addStructField(fset *token.FileSet, fields *ast.FieldList, file *ast.File, typ string, name string, comment string) {
    prevField := fields.List[fields.NumFields()-1] 

    c := &ast.Comment{Text: fmt.Sprint("// ", comment), Slash: prevField.End() + 1}
    cg := &ast.CommentGroup{List: []*ast.Comment{c}}
    o := ast.NewObj(ast.Var, name)
    f := &ast.Field{
        Doc:   cg,
        Names: []*ast.Ident{&ast.Ident{Name: name, Obj: o, NamePos: cg.End() + 1}},
    }
    o.Decl = f
    f.Type = &ast.Ident{Name: typ, NamePos: f.Names[0].End() + 1}

    fset.File(c.End()).AddLine(int(c.End()))
    fset.File(f.End()).AddLine(int(f.End()))

    fields.List = append(fields.List, f)
    file.Comments = append(file.Comments, cg)
}

Example: http://play.golang.org/p/_q1xh3giHm

For Item (3), it is also important to set all the overallocated bytes to spaces (0x20), so that the printer doesn't complain about null bytes when processing them.