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)
}

Tehuantepec

I suspect you need to update the Comment Map for this to work correctly. – Understate

Here you can see some details of the actual and expected trees : play.golang.org/p/qv63Hu1xmP thanks to golang.org/pkg/go/ast/#Fprint . The main differences I see are Slash, NamePos, and Obj not set. I tried to fiddle with the positions, but could not get it right... – Dysuria

This one has me stumped... It looks like there is some kind of other bookeeping that needs to be done, as I was able to get Slash and NamePos to match (offset by 100 notwithstanding) in this: play.golang.org/p/pQodZncMjA -- and even adding in AddLine and CommentMap doesn't seem to help: play.golang.org/p/GGj2eDwDF- – Laurinda

I have made some progress on this: play.golang.org/p/rOqcPnovsx The tricks needed include: (1) Specifically setting the buffer locations including Slash and NamePos, (2) token.File.AddLine to add new lines at specific offsets, (3) Overallocating the source buffer so token.File.Position (used "under the covers") and token.File.AddLine don't fail range checks against the source buffer. But this introduces a bunch of serialization warnings/errors because of the mismatched size of the new tokens and the underlying buffer size. – Circumvolution

Answer 1

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

Suter

Yes! I asked this question and eventually got so annoyed with there not being an easy solution I created the dst package. I'll accept this answer! – Tehuantepec

Haha, I should have realized that Dave was your GitHub handle. – Suter

Answer 2

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.

Circumvolution

I think you've cracked it! Well done. I'm sure I'll have follow up questions as I add more functionality, but this is a great answer. Thanks! – Tehuantepec

This only works for this particular situation. If there were more code with newlines after the B struct than the length of the newly added fields and comments, then the calls to AddLine would be ignored (the impl. has not changed for the last nine years). See: play.golang.org/p/NsBL74Z5bb7 Additionally, adding comments at the bottom of the "file" make things go haywire... See: play.golang.org/p/g2vKIPm1OUx – Carnet