// rewriteFile applies the rewrite rule 'pattern -> replace' to an entire file.
func rewriteFile(pattern, replace ast.Expr, p *ast.File) *ast.File {
cmap := ast.NewCommentMap(fileSet, p, p.Comments)
m := make(map[string]reflect.Value)
pat := reflect.ValueOf(pattern)
repl := reflect.ValueOf(replace)
var rewriteVal func(val reflect.Value) reflect.Value
rewriteVal = func(val reflect.Value) reflect.Value {
// don't bother if val is invalid to start with
if !val.IsValid() {
return reflect.Value{}
}
for k := range m {
delete(m, k)
}
val = apply(rewriteVal, val)
if match(m, pat, val) {
val = subst(m, repl, reflect.ValueOf(val.Interface().(ast.Node).Pos()))
}
return val
}
r := apply(rewriteVal, reflect.ValueOf(p)).Interface().(*ast.File)
r.Comments = cmap.Filter(r).Comments() // recreate comments list
return r
}
// This example illustrates how to remove a variable declaration
// in a Go program while maintaining correct comment association
// using an ast.CommentMap.
func ExampleCommentMap() {
// src is the input for which we create the AST that we
// are going to manipulate.
src := `
// This is the package comment.
package main
// This comment is associated with the hello constant.
const hello = "Hello, World!" // line comment 1
// This comment is associated with the foo variable.
var foo = hello // line comment 2
// This comment is associated with the main function.
func main() {
fmt.Println(hello) // line comment 3
}
`
// Create the AST by parsing src.
fset := token.NewFileSet() // positions are relative to fset
f, err := parser.ParseFile(fset, "src.go", src, parser.ParseComments)
if err != nil {
panic(err)
}
// Create an ast.CommentMap from the ast.File's comments.
// This helps keeping the association between comments
// and AST nodes.
cmap := ast.NewCommentMap(fset, f, f.Comments)
// Remove the first variable declaration from the list of declarations.
f.Decls = removeFirstVarDecl(f.Decls)
// Use the comment map to filter comments that don't belong anymore
// (the comments associated with the variable declaration), and create
// the new comments list.
f.Comments = cmap.Filter(f).Comments()
// Print the modified AST.
var buf bytes.Buffer
if err := format.Node(&buf, fset, f); err != nil {
panic(err)
}
fmt.Printf("%s", buf.Bytes())
// output:
// // This is the package comment.
// package main
//
// // This comment is associated with the hello constant.
// const hello = "Hello, World!" // line comment 1
//
// // This comment is associated with the main function.
// func main() {
// fmt.Println(hello) // line comment 3
// }
}