Rewriting Go source with AST tooling
Go's standard library ships with powerful AST tooling in the go/* packages, and the golang.org/x/tools module adds even more capability. One task that's particularly well-served is rewriting source code. With a mutable AST, you can inspect, edit, and re-emit Go code programmatically — useful for codemods, instrumentation, and refactoring tools.
This walkthrough uses a small Go snippet as the subject, walking through finding and modifying nodes with the standard library, then showing why a different package is needed for certain kinds of edits.
package p
func pred() bool {
return true
}
func pp(x int) int {
if x > 2 && pred() {
return 5
}
var b = pred()
if b {
return 6
}
return 0
}
Finding nodes in an AST
The first step is locating points of interest. The go/ast package offers two mechanisms: ast.Walk, which requires implementing the ast.Visitor interface, and ast.Inspect, which takes a closure.
Starting with ast.Walk, you parse source into an *ast.File using parser.ParseFile, then walk it:
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "src.go", os.Stdin, 0)
if err != nil {
log.Fatal(err)
}
A visitor implementation handles each node:
type Visitor struct {
fset *token.FileSet
}
func (v *Visitor) Visit(n ast.Node) ast.Visitor {
if n == nil {
return nil
}
switch x := n.(type) {
case *ast.CallExpr:
id, ok := x.Fun.(*ast.Ident)
if ok {
if id.Name == "pred" {
fmt.Printf("Visit found call to pred() at %s\n", v.fset.Position(n.Pos()))
}
}
}
return v
}
type Visitor struct {
fset *token.FileSet
}
func (v *Visitor) Visit(node ast.Node) ast.Visitor {
if call, ok := node.(*ast.CallExpr); ok {
if ident, ok := call.Fun.(*ast.Ident); ok && ident.Name == "pred" {
fmt.Printf("%s: call to pred\n", v.fset.Position(call.Pos()))
}
}
return v
}
Note the type assertion on call.Fun: this only catches calls where the function is referenced by an identifier. The visitor stores a *token.FileSet to translate token.Pos values (which are just integers) into human-readable file:line:column positions.
For simple filtering, ast.Inspect is lighter weight — no interface implementation needed, just a function:
func main() {
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "src.go", os.Stdin, 0)
if err != nil {
log.Fatal(err)
}
ast.Inspect(file, func(n ast.Node) bool {
switch x := n.(type) {
case *ast.CallExpr:
id, ok := x.Fun.(*ast.Ident)
if ok {
if id.Name == "pred" {
fmt.Printf("Inspect found call to pred() at %s\n", fset.Position(n.Pos()))
}
}
}
return true
})
}
The matching logic is identical, but the scaffolding is simpler. Unless you specifically need ast.Walk, ast.Inspect is the recommended approach — and it's what we'll use for the initial rewrite examples.
Editing the AST and re-emitting code
The AST from the parser is fully mutable: nodes are connected via pointers, so you can change them in place, or even build entirely new node trees. The go/format package converts the modified AST back into formatted Go source. On its own, format.Node will round-trip code (though comments are dropped by default):
func main() {
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "src.go", os.Stdin, 0)
if err != nil {
log.Fatal(err)
}
format.Node(os.Stdout, fset, file)
}
Now consider two transformations on the example program:
- Rename
predtopred2, updating every call site. - Inject a
fmt.Printlnstatement at the top of each function body.
After the rewrite, the output looks like this (added lines marked):
package p
func pred2() bool {
fmt.Println("instrumentation")
return true
}
func pp(x int) int {
fmt.Println("instrumentation")
if x > 2 && pred2() {
return 5
}
var b = pred2()
if b {
return 6
}
return 0
}
(Note: this example deliberately avoids adding the fmt import — that's left as a separate exercise.)
Using ast.Inspect, renaming call sites is a matter of matching *ast.CallExpr nodes and appending to the identifier name:
ast.Inspect(file, func(n ast.Node) bool {
switch x := n.(type) {
case *ast.CallExpr:
id, ok := x.Fun.(*ast.Ident)
if ok {
if id.Name == "pred" {
id.Name += "2"
}
}
// ...
Since we're editing the actual AST, these changes are live; no copy is involved.
Handling function declarations follows the same pattern: rename, then prepend instrumentation:
case *ast.FuncDecl:
if x.Name.Name == "pred" {
x.Name.Name += "2"
}
newCallStmt := &ast.ExprStmt{
X: &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{
Name: "fmt",
},
Sel: &ast.Ident{
Name: "Println",
},
},
Args: []ast.Expr{
&ast.BasicLit{
Kind: token.STRING,
Value: `"instrumentation"`,
},
},
},
}
x.Body.List = append([]ast.Stmt{newCallStmt}, x.Body.List...)
Each *ast.FuncDecl has a Body, which contains a slice of statements in its List field. Prepending a hand-crafted ast.Stmt to that slice inserts the new statement at the top of the body. Building custom AST nodes takes some familiarity — tools like go2ast, which emit the Go code to construct a given snippet's AST, make this much easier. Finally, the modified tree is emitted back to source:
fmt.Println("Modified AST:")
format.Node(os.Stdout, fset, file)
Where Inspect falls short
Not every rewrite fits this pattern. Consider a transformation that wraps each pred() call in a logical NOT: !pred(). How would you implement that?
The problem is structural. When ast.Inspect delivers a node, it gives you a pointer to that node's contents, which you may mutate, but it does not give you the ability to replace the node itself. Replacement requires access to the parent — or, more precisely, a pointer to the pointer that refers to this node in the parent's child list. With ast.Inspect, you simply hold the inner pointer; you can't change what the parent points to.
This design limitation was noted years ago in the Go issue tracker. In 2017, the golang.org/x/tools/go/ast/astutil package arrived to fill the gap.
Replacing nodes with astutil
astutil extends the walker pattern so that callbacks receive a *astutil.Cursor, which internally tracks the parent chain. This makes it possible to replace, delete, or insert nodes at arbitrary points in the tree. The Apply function runs callbacks both before and after each node is visited.
Here's the full implementation of the NOT-wrapping transformation:
func main() {
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "src.go", os.Stdin, 0)
if err != nil {
log.Fatal(err)
}
astutil.Apply(file, nil, func(c *astutil.Cursor) bool {
n := c.Node()
switch x := n.(type) {
case *ast.CallExpr:
id, ok := x.Fun.(*ast.Ident)
if ok {
if id.Name == "pred" {
c.Replace(&ast.UnaryExpr{
Op: token.NOT,
X: x,
})
}
}
}
return true
})
fmt.Println("Modified AST:")
format.Node(os.Stdout, fset, file)
}
// applying a unary ! to each pred() call
func negatePreds(f *ast.File) {
astutil.Apply(f, nil, func(c *astutil.Cursor) bool {
call, ok := c.Node().(*ast.CallExpr)
if !ok {
return true
}
ident, ok := call.Fun.(*ast.Ident)
if !ok || ident.Name != "pred" {
return true
}
c.Replace(&ast.UnaryExpr{
Op: token.NOT,
X: call,
})
return true
})
}
The callback identifies the node the same way as before, but this time the replacement is performed via the cursor: the original CallExpr node is swapped out for a UnaryExpr that wraps it. Hidden inside Cursor is the parent reference, which makes the substitution possible.
That is the clean line between ast.Inspect and astutil: the former is fine for editing node contents, but when you need to replace whole nodes, you have to reach for astutil — or build your own parent-aware traversal.
