Comentarios fuera de orden después de agregar el artículo a Go AST

abstract-syntax-tree

Question

La siguiente prueba intenta usar AST para agregar campos a una estructura. Los campos se agregan correctamente, pero los comentarios se agregan fuera de orden. Puede que sea necesario especificar la posición manualmente, pero hasta ahora he dejado un espacio en blanco para encontrar una respuesta.

Aquí hay una prueba fallida: http://play.golang.org/p/RID4N30FZK

Aquí está el código:

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

Creo que lo he conseguido para que funcione. Como se indica en mi comentario anterior, los principales puntos requeridos son:

1. Especifique específicamente las ubicaciones de los búferes, incluyendo Slash y NamePos
2. Use token.File.AddLine para agregar nuevas líneas en compensaciones específicas (calculadas usando las posiciones del elemento 1)
3. Ubique en general el búfer de origen, de modo que token.File.Position (utilizado por printer.Printer y token.File.Addline no fallan las comprobaciones de rango en el búfer de origen

Código:

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

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

Para el elemento (3) , también es importante configurar todos los bytes asignados a los espacios ( 0x20 ), para que la impresora no se queje de los bytes nulos cuando los procesa.