Comments out of order after adding item to Go AST

Question

The following test attempts to use AST to add fields to a struct. The fields are added correctly, but the comments are added out of order. I gather the position may need to be specified manually, but I've so far drawn a blank finding an answer.

Here's a failing test: http://play.golang.org/p/RID4N30FZK

Here's the code:

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:\n%s\nGot:\n%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:\n%s\nGot:\n%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.

Answer 2

I know that this answer might be a little late. But for the benefit of others, I found a reference to this library in the following GitHub issue

https://github.com/golang/go/issues/20744

The library is called dst and it can convert a go ast to dst and vice versa.

https://github.com/dave/dst

In ast, Comments are stored by their byte offset instead of attached to nodes. Dst solves this by attaching the comments to its respective nodes so that re-arranging nodes doesn't break the output/tree.

The library works as advertized and I haven't found any issues so far.

Note: There is also a subpackage called dst/dstutil which is compatible with golang.org/x/tools/go/ast/astutil