Eli Bendersky's website - Go tooling

Basic Go tooling for generics

2022-09-17T07:53:00-07:00

Continuing the recent theme of blog posts about Go tooling, this post will briefly discuss how analysis tools interact with one of Go's newest features - generics. As usual, this post provides a motivating example and discusses a couple of approaches to writing a tool. The full code is on GitHub.

As a motivating example, let's find all calls to generic functions in a module, and the instantiated type arguments for that call. We'll use this code sample:

package main

import "golang.org/x/exp/slices"

func foo() {
  var fs []float64

  newfs := slices.Clone(fs)
  handle(newfs, len(newfs))
}

func bar() {
  var fs []string
  handle[string, int](fs, 10)
}

func main() {
  foo()
  bar()
}

func handle[T any, N ~int](k []T, n N) {
}

Note that there are several different calls to generic functions here:

A call to slices.Clone in foo that uses type inference.
A call to handle in foo that uses type inference.
A call to handle in bar that uses explicit type instantiation.

First attempt: use types.Info.Instances directly

When rummaging through the documentation of the go/types package, one encounters some new fields added to the Info structure; in particular, the Instances field:

// Instances maps identifiers denoting generic types or functions to their
// type arguments and instantiated type.
//
// For example, Instances will map the identifier for 'T' in the type
// instantiation T[int, string] to the type arguments [int, string] and
// resulting instantiated *Named type. Given a generic function
// func F[A any](A), Instances will map the identifier for 'F' in the call
// expression F(int(1)) to the inferred type arguments [int], and resulting
// instantiated *Signature.
//
// Invariant: Instantiating Uses[id].Type() with Instances[id].TypeArgs
// results in an equivalent of Instances[id].Type.
Instances map[*ast.Ident]Instance

This seems super useful for our task! However, on a closer look, some challenges become apparent. Note that the key of this map is an *ast.Ident, which points to the identifier representing the call in the AST. For the call to handle in foo, the AST looks as follows:

Note that the Ident node is a child of the CallExpr node. Given the identifier, how do we find the actual call (since we want the arguments)?

This turns out to be non-trivial! Take a look at my previous post on the subject of Parent links in Go ASTs for details; none of the techniques described there is perfect for our needs here:

Keeping manual track of parents or using Inspector.WithStack won't help because we're not actually traversing the AST. We have immediate access to the *ast.Ident for generic calls.
astutil.PathEnclosingInterval is quite promising but awkward to use in this case, mostly because it requires a root *ast.File which we don't necessarily have. It also takes positions as inputs and does the AST traversal itself so using it feels somewhat backwards here.

Running into all this trouble, a reasonable conclusion is that this field was not designed to be used in this particular way.

Second attempt: traverse the AST and use Instances for lookups

Let's put shortcuts aside and try to approach it in the usual tool-writing way: we'll traverse the AST and examine the type information where necessary. I won't paste the usual tool scaffolding (read about it here); we start with the place where we have a *packages.Package in hand:

func processPackage(pkg *packages.Package) {
  if len(pkg.Errors) > 0 {
    for _, e := range pkg.Errors {
      fmt.Printf("Error: %s\n", e)
    }
    os.Exit(1)
  }

  walkAstAndJoinInstances(pkg)
}

The overall approach is: recursively walk the AST with Inspect, and when we encounter a call expression examine the call target, checking it against the Instances field. There are some complications here to be aware of, though. The Fun field of ast.CallExpr is not necessarily an identifier. It can be a SelectorExpr for qualified calls like slices.Clone; it can also be an IndexListExpr for calls that have explicit type parameters like the second call to handle in our sample.

So our AST analysis has to aware of these options, and we end up with something like:

func walkAstAndJoinInstances(pkg *packages.Package) {
  for _, fileAst := range pkg.Syntax {
    ast.Inspect(fileAst, func(n ast.Node) bool {
      if cexpr, ok := n.(*ast.CallExpr); ok {
        var id *ast.Ident
        switch fn := cexpr.Fun.(type) {
        case *ast.Ident:
          id = fn
        case *ast.SelectorExpr:
          id = fn.Sel
        case *ast.IndexListExpr:
          if sel, ok := fn.X.(*ast.SelectorExpr); ok {
            id = sel.Sel
          } else {
            id = fn.X.(*ast.Ident)
          }
        }

        if id != nil {
          if v, ok := pkg.TypesInfo.Instances[id]; ok {
            fmt.Println("call", id)
            fmt.Printf("  instantiation type=%v args=%v\n", v.Type, getListOfTypes(v.TypeArgs))
          }
        }
      }
      return true
    })
  }
}

The getListOfTypes function is a utility to extract a slice of types from types.TypeList:

func getListOfTypes(tl *types.TypeList) []types.Type {
  var sl []types.Type
  for i := 0; i < tl.Len(); i++ {
    sl = append(sl, tl.At(i))
  }
  return sl
}

If we run our tool on the code sample, we get:

$ go run find-generic-calls.go -- samplemodules/use-generics/
call Clone
  instantiation type=func(s []float64) []float64 args=[[]float64 float64]
call handle
  instantiation type=func(k []float64, n int) args=[float64 int]
call handle
  instantiation type=func(k []string, n int) args=[string int]

Mission accomplished! We've been able to identify calls to generic functions, both with and without type inference, and to access their arguments.

A taste for more

While the tool developed in this post is very basic, hopefully it provides a gentle introduction to writing tools with generic code in mind. For a much more detailed discussion, I strongly recommend going over the guide at https://github.com/golang/exp/tree/master/typeparams/example; if you plan to write more sophisticated analyzers that work on generic code, it will prove invaluable.

Parent links in Go ASTs

2022-09-03T05:37:00-07:00

A common question that comes up when writing AST-based tooling in Go is "how do I find a node's parent?", or more generally, "what are all the ancestors of a given node?"

Go's AST package does not come with parent links out of the box; this is a common choice in AST libraries [1]. Parent links are rarely needed, but they carry significant overheads in terms of memory usage and CPU time. Typical AST nodes can be fairly compact and adding a pointer to each node can increase its size considerably. Keeping this pointer populated properly also takes CPU time; furthermore, in GC'd languages like Go, parent links are a significant source of pointers to scan and keep track of. All of these are fixed costs imposed on any analyzer, just for the sake of the few that may need this information.

But say your analysis really does need parent links? What do you do then? This post explores a few options.

Problem statement

For the sake of discussion, we need a concrete problem to focus on. To keep it simple, we'll use something artificial: suppose we want to find all the binary multiplication expressions that are nested inside other binary expressions. For example:

m = y + z * x

The z * x sub-expression is represented by an ast.BinaryExpr node, which is nested inside another ast.BinaryExpr (the +). Here's a graphical AST dump from my new favorite AST viewer:

We'll consider the following code as fulfilling our condition too:

m = x + foo(y*z)

The multiplication is nested inside the addition, even though the multiplication's immediate parent is a CallExpr.

Approach #1: ad-hoc

The first approach that comes to mind is "just use some flags during traversal"; I'll call it the ad-hoc approach, and it's a perfectly valid way to solve the problem. Let's start with the usual tool scaffolding (the full code is on GitHub):

var fset = token.NewFileSet()

func main() {
  const mode packages.LoadMode = packages.NeedName |
    packages.NeedTypes |
    packages.NeedSyntax |
    packages.NeedTypesInfo

  flag.Parse()
  if flag.NArg() != 1 {
    log.Fatal("Expecting a single argument: directory of module")
  }

  cfg := &packages.Config{Fset: fset, Mode: mode, Dir: flag.Args()[0]}
  pkgs, err := packages.Load(cfg, "./...")
  if err != nil {
    log.Fatal(err)
  }

  for _, pkg := range pkgs {
    processPackage(pkg)
  }
}

func processPackage(pkg *packages.Package) {
  if len(pkg.Errors) > 0 {
    for _, e := range pkg.Errors {
      fmt.Printf("Error: %s\n", e)
    }
    os.Exit(1)
  }

  discoverNodeParentsAdhoc(pkg)
}

As usual, our AST traversal employs ast.Inspect, which visits the entire AST in depth-first order. We can cleverly maintain a "I have a ast.BinaryExpr parent" state through the traversal, but it will be much easier to invert the logic, like this:

func discoverNodeParentsAdhoc(pkg *packages.Package) {
  for _, fileAst := range pkg.Syntax {
    ast.Inspect(fileAst, func(n ast.Node) bool {
      if bexpr, ok := n.(*ast.BinaryExpr); ok {
        findBinMulChild(bexpr.X)
        findBinMulChild(bexpr.Y)
        return false
      }
      return true
    })
  }
}

func findBinMulChild(ancestor ast.Node) {
  ast.Inspect(ancestor, func(n ast.Node) bool {
    if bexpr, ok := n.(*ast.BinaryExpr); ok && bexpr.Op == token.MUL {
      fmt.Printf("found BinaryExpr(*) as a child of another binary expr: %v\n",
        fset.Position(n.Pos()))
    }
    return true
  })
}

For each ast.BinaryExpr, we start another recursive visit of both its children. If this visit finds a binary * - bingo! Note that we don't do double work here, since after calling findBinMulChild on the left and right children we return false. That's the beauty of ad-hoc approaches - they are, well.. ad-hoc :-) We only need to do as much work as our analysis requires. The obvious disadvantage is that every task is different and we can't always reuse a lot between them.

Next, let's try something more generic.

Approach #2: keeping manual track of parents

Recall the original problem this post started with: AST nodes have no parent links. If we had a way to follow a chain of parents, solving the task would be easy: once we get to a * node, we simply go up the chain of parents and see if there are any ast.BinaryExprs in it.

Well, it turns out to be pretty easy to build a list of parents manually. The key insight is that ast.Inspect helps us by calling our traversal function with nil at the end of each node's children. We can use that for bookkeeping:

func discoverNodeParentsManualStack(pkg *packages.Package) {
  for _, fileAst := range pkg.Syntax {
    var ancestors []ast.Node
    ast.Inspect(fileAst, func(n ast.Node) bool {
      if bexpr, ok := n.(*ast.BinaryExpr); ok && bexpr.Op == token.MUL {
        // Walk the ancestor stack to find if one of them is also a BinaryExpr
        for i := len(ancestors) - 1; i >= 0; i-- {
          if _, ok := ancestors[i].(*ast.BinaryExpr); ok {
            fmt.Printf("found BinaryExpr(*) as a child of another binary expr: %v\n",
              fset.Position(n.Pos()))
            break
          }
        }
      }

      if n == nil {
        // Pop, since we're done with this node and its children.
        ancestors = ancestors[:len(ancestors)-1]
      } else {
        // Push this node on the stack, since its children will be visited
        // next.
        ancestors = append(ancestors, n)
      }
      return true
    })
  }
}

This code keeps a stack of the node's ancestors in a slice. Since the traversal is DFS, whenever out function is called we push the current node onto the stack (it will serve as an ancestor for subsequent calls), and when called with nil (meaning that we're done with this a node's children) we pop it off.

When an actual * node is found, it's straightforward to walk the stack upwards looking for an ast.BinaryExpr [2].

Approach #3: using Inspector.WithStack

The previous section showed that it's fairly easy to maintain a full stack of ancestors while traversing the AST, but even this minimal effort is unnecessary. The Go tooling packages come with some useful helpers, and this is one use case that's already covered.

The ast/inspector package provides helper functions for AST traversal. One capability that's very useful for our scenario is WithStack. You can read its documentation yourself, so without further ado here's how we solve our concrete problem using it:

func discoverNodeParentsWithStack(pkg *packages.Package) {
  insp := inspector.New(pkg.Syntax)
  insp.WithStack(nil, func(n ast.Node, push bool, stack []ast.Node) bool {
    if bexpr, ok := n.(*ast.BinaryExpr); push && ok && bexpr.Op == token.MUL {
      for i := len(stack) - 2; i >= 0; i-- {
        if _, ok := stack[i].(*ast.BinaryExpr); ok {
          fmt.Printf("found BinaryExpr(*) as a child of another binary expr: %v\n",
            fset.Position(n.Pos()))
          break
        }
      }
    }
    return true
  })
}

Approach #4: with PathEnclosingInterval

Finally, there's yet another tool we could use - PathEnclosingInterval from the ast/astutil package. This function is a bit unusual, in that it takes a token position, and returns all the AST nodes that contain that position in the source. It's a very specialized API that's mostly useful for text editors, but it can also solve our task for this post:

func discoverNodeParentsPathInterval(pkg *packages.Package) {
  for _, fileAst := range pkg.Syntax {
    ast.Inspect(fileAst, func(n ast.Node) bool {
      if bexpr, ok := n.(*ast.BinaryExpr); ok && bexpr.Op == token.MUL {
        path, _ := astutil.PathEnclosingInterval(fileAst, bexpr.Pos(), bexpr.End())
        for i := len(path) - 1; i >= 0; i-- {
          if _, ok := path[i].(*ast.BinaryExpr); ok && path[i] != bexpr {
            fmt.Printf("found BinaryExpr(*) as a child of another binary expr: %v\n",
              fset.Position(n.Pos()))
            break
          }
        }
      }
      return true
    })
  }
}

To be clear: I don't think this approach makes a whole lot of sense in this scenario, and it's much better to use WithStack. I'm just including it for completeness.

[1]	In my pycparser project, this is a FAQ.

[2]	For this particular analysis the order of stack traversal doesn't matter, so our loop could go from 0 to `len-1` instead.

Error handling with Go tooling

2022-08-06T06:24:00-07:00

Most tutorials on Go tooling (and probably most other tooling) tend to focus on the happy path - the input code is perfectly valid and contains no errors. This is often a reasonable assumption, because we tend to run tools on existing code bases, and these are likely to compile correctly in their steady state. However, sometimes knowing how to handle errors is important, and good tooling infrastructure will enable that.

In this quick post I'll discuss how Go tooling handles potentially erroneous code. The full code is available on GitHub.

Our sample tool starts with the basic setup described in my previous post (Step 0), and for now just a bare-bones processPackage that dumps the package's AST:

func processPackage(pkg *packages.Package) {
  for _, fileAst := range pkg.Syntax {
    ast.Print(fset, fileAst)
  }
}

Let's look at an input file with some errors in it:

package main

func util(x int) {
}

func main() {
  util(5 6)

  var s string = 5.5
}

Time for a brief self-evaluation :-) Most Go programmers will immediately identify two errors in this code; somewhat more experienced gophers will also find the third. Give yourself a few moments to play the compiler here - the answers are further down in the post.

Despite this code having some errors, if we run our tool on it - it will happily dump an AST with no complaints! If you're new to Go tooling, this can be somewhat surprising; wouldn't you expect the tool to error out somehow?

If you look closely at the dumped AST you'll notice that it's somewhat malformed, which is reasonable given the errors in the code. The call util(5 6) is syntactically incorrect - there's an extra parameter in a function call, without a separating comma. The AST dumped for this call is:

X: *ast.CallExpr {
.  Fun: *ast.Ident {
.  .  NamePos: sample-module/main.go:9:2
.  .  Name: "util"
.  .  Obj: *(obj @ 23)
.  }
.  Lparen: sample-module/main.go:9:6
.  Args: []ast.Expr (len = 1) {
.  .  0: *ast.BasicLit {
.  .  .  ValuePos: sample-module/main.go:9:7
.  .  .  Kind: INT
.  .  .  Value: "5"
.  .  }
.  }

The parser chose to include the first parameter and ignore the second. Why does it bother to produce an AST at all here, even for code that doesn't parse correctly?

The answer is that tooling is often used in things like IDEs. Imagine you have a long file open in an IDE and there's an error somewhere in the middle [1]; had the IDE stopped at the error, it would not syntax highlight the rest of the code, or offer any intellisense features for it like "jump to definition". This would not be great for ergonomics. Therefore tools attempt to recover from compilation errors and continue churning through the code, to provide maximal utility.

This isn't always possible and some syntax errors confuse the tool entirely; we're all familiar with what happens when we forget a closing ) or } somewhere.

Error reports from Go tooling

Back to the original goal of the post... given a tool we're writing, is there a way to know that there were errors in the code before we start analyzing the AST?

Yes. The packages.Package structure has a couple of fields set by XTGP for this purpose. The most important one is Errors. Let's add this code to our processPackage function:

if len(pkg.Errors) > 0 {
  fmt.Printf("package %v has %v errors\n", pkg.PkgPath, len(pkg.Errors))

  for _, e := range pkg.Errors {
    var errtype string
    switch e.Kind {
    case packages.ListError:
      errtype = "listing/driver"
    case packages.ParseError:
      errtype = "parser"
    case packages.TypeError:
      errtype = "type checker"
    default:
      errtype = "unknown"
    }
    fmt.Printf("Error [%v]: %s\n", errtype, e)
  }
}

When invoked on the erroneous module, we'll see:

package example.com has 3 errors
Error [parser]: sample-module/main.go:9:9: missing ',' in argument list
Error [type checker]: sample-module/main.go:11:17: cannot use 5.5 (untyped float constant) as string value in variable declaration
Error [type checker]: sample-module/main.go:11:6: s declared but not used
<AST dump, if still enabled>

Each error has the type packages.Error. An additional field that's set (but only if the packages.NeedTypes flag is set while configuring the tool) is IllTyped:

// IllTyped indicates whether the package or any dependency contains errors.
// It is set only when Types is set.
IllTyped bool

How should this be used?

OK, so now we know how to see what errors were encountered when parsing the input module. What's next?

What you do with this information is entirely dependent on the nature of your tool. As mentioned earlier, most tools assume that the code they're working on is already correct; in this case, erroring out if len(pkg.Errors) > 0 may be a good idea since we don't want to produce phantom results from incorrect code. Other tools will specifically not care about errors and will try to churn through as much AST as possible, even for partially incorrect code. YMMV.

[1]

Why would there be an error in the middle? Consider that IDEs often repeatedly parse the file, even while we're typing, to be able to provide intellisense features on the fly. While the code is being typed in, it's often un-parsable or has type errors (think of being in the middle of a parameter list fo a function call).

This also highlights an interesting topic in compiler frontends - error recovery. While classical compilers can be forgiven for printing out a few errors and bailing out, tools really do need to recover as quickly as possible to be able to process the rest of the code correctly, even if there's some issue in the middle.

Exploring function parameter types with Go tooling

2022-07-09T06:24:00-07:00

I enjoy hacking on Go tooling (Go programs that analyze and manipulate other Go programs), and a few weeks ago an interesting Stack Overflow question caught my attention. Given a function with a parameter of some structure type:

func foo(param SomeType) {}

How can we examine the fields of SomeType?

The reason I found this question interesting is because it's very representative of the problems Go tool writers typically face:

Find something in the code (e.g. a function)
Analyze this thing (e.g. check its parameters)
Find something else based on the analysis (e.g. the fields of the parameters' types)

Let's see how to write such a tool; the post presents the tool in pieces, with the full code on GitHub.

Step 0: setting up

As usual, we'll be using the x/tools/go/packages package [1]. In this day an age, we expect the input to be a Go module, and the tool will look for a path to the module as its sole command-line argument:

var fset = token.NewFileSet()

func main() {
  const mode packages.LoadMode = packages.NeedName |
    packages.NeedTypes |
    packages.NeedSyntax |
    packages.NeedTypesInfo

  flag.Parse()
  if flag.NArg() != 1 {
    log.Fatal("Expecting a single argument: directory of module")
  }

  cfg := &packages.Config{Fset: fset, Mode: mode, Dir: flag.Args()[0]}
  pkgs, err := packages.Load(cfg, "./...")
  if err != nil {
    log.Fatal(err)
  }

  for _, pkg := range pkgs {
    processPackage(pkg)
  }
}

After it loads all the code, main iterates over all the packages found in the input module and invokes processPackage for each.

Step 1: finding functions

This part is very simple. A package in Go consists of multiple source files; each file has an AST representation, and we we can use ast.Inspect to find all the function declarations as follows [2]:

func processPackage(pkg *packages.Package) {
  for _, fileAst := range pkg.Syntax {
    ast.Inspect(fileAst, func(n ast.Node) bool {
      if funcDecl, ok := n.(*ast.FuncDecl); ok {
        processFuncDecl(funcDecl, pkg.TypesInfo)
      }

      return true
    })
  }
}

Note that we're extracting TypesInfo from the packages.Package struct. This is the type information populated by the type checker because we've set the packages.NeedTypesInfo flag in LoadMode. This is an important prerequisite for step 3.

Step 2: analyzing function parameters

With a function declaration in our hands, we can walk over its parameter list and examine the type of each parameter:

func processFuncDecl(fd *ast.FuncDecl, tinfo *types.Info) {
  fmt.Println("=== Function", fd.Name)
  for _, field := range fd.Type.Params.List {
    var names []string
    for _, name := range field.Names {
      names = append(names, name.Name)
    }
    fmt.Println("param:", strings.Join(names, ", "))
    processTypeExpr(field.Type, tinfo)
  }
}

Note that in Go, multiple parameters can share a type - as in func foo(a, b int); this code collects all the names of a field into a single slice for reporting. processTypeExpr is then invoked to process the actual field type.

Step 3: finding the types of struct fields

Analyzing the actual parameter type is the most involved part of our tool. There are two ways to reason about a type in Go source code - the basic syntax level, and a semantic level.

On the syntax level, a type is what the parser sees and is represented by ast.Expr. For example, for type *T, this would be an ast.StarExpr with whatever T is in its X field. Our processFuncDecl function focuses on this syntax level, and it passes an ast.Expr to processTypeExpr to represent the type.

On the semantic level, the tooling infrastructure runs a full type checker (from the stdlib go/types package) to determine a mapping between names and actual types in the code.

Thinking about it differently, the syntax level is what we see by just looking at the code - MyType is an identifier, but for the Go compiler (and tooling) it can carry a deeper (semantic) meaning - a reference to some type.

func processTypeExpr(e ast.Expr, tinfo *types.Info) {
  switch tyExpr := e.(type) {
  case *ast.StarExpr:
    fmt.Println("  pointer to...")
    processTypeExpr(tyExpr.X, tinfo)
  case *ast.ArrayType:
    fmt.Println("  slice or array of...")
    processTypeExpr(tyExpr.Elt, tinfo)
  default:
    switch ty := tinfo.Types[e].Type.(type) {
    case *types.Basic:
      fmt.Println("  basic =", ty.Name())
    case *types.Named:
      fmt.Println("  name =", ty)
      uty := ty.Underlying()
      fmt.Println("  type =", ty.Underlying())
      if sty, ok := uty.(*types.Struct); ok {
        fmt.Println("  fields:")
        processStructParamType(sty)
      }
      fmt.Println("  pos =", fset.Position(ty.Obj().Pos()))
    default:
      fmt.Println("  unnamed type =", ty)
      if sty, ok := ty.(*types.Struct); ok {
        fmt.Println("  fields:")
        processStructParamType(sty)
      }
    }
  }
}

This function is recursively peeling off pointers and arrays/slices off the type, but eventually ends up in the default clause of the switch. This is where it queries the types.Info that was passed in to see what type it's dealing with. There are "basic" types like string, named types like MyType and even unnamed types like struct {x int} [3].

The type mapping has some quirks that may be surprising when working with it for the first time. For example, ast.SelectorExpr types that are used to identify dot-separated names like http.Handler; these have mappings in types.Info (which maps AST nodes to their types). However, pointer types like *MyType do not have direct mappings, and the ast.StarExpr has to be peeled off first. With some experience and judicious use of AST dumping these rules aren't hard to figure out, though.

Finally, here's the helper for reporting all fields of a struct:

func processStructParamType(sty *types.Struct) {
  for i := 0; i < sty.NumFields(); i++ {
    field := sty.Field(i)
    fmt.Println("    ", field.Name(), field.Type())
  }
}

This is where the tool's final destination is. Whatever you need from these struct field types, you can do here.

Sample run

Suppose we have an input module with a single package main, with this file (main.go):

package main

import (
  "net/http"
  "strings"
)

func foo(x string, u User, up *User, h http.Handler, s struct{ b bool }) {
}

func bar(sb *strings.Builder) {
}

And an additional file defining the User type (user.go):

package main

type User struct {
  id   int
  name string
}

Now we can run our tool on this module, and get the following output:

=== Function foo
param: x
  basic = string
param: u
  name = example.com.User
  type = struct{id int; name string}
  fields:
     id int
     name string
  pos = 2022/go-tool-func-param-types/sample-module/user.go:3:6
param: up
  pointer to...
  name = example.com.User
  type = struct{id int; name string}
  fields:
     id int
     name string
  pos = 2022/go-tool-func-param-types/sample-module/user.go:3:6
param: h
  name = net/http.Handler
  type = interface{ServeHTTP(net/http.ResponseWriter, *net/http.Request)}
  pos = $GOROOT/src/net/http/server.go:86:1
param: s
  unnamed type = struct{b bool}
  fields:
     b bool
=== Function bar
param: sb
  pointer to...
  name = strings.Builder
  type = struct{addr *strings.Builder; buf []byte}
  fields:
     addr *strings.Builder
     buf []byte
  pos = $GOROOT/src/strings/builder.go:15:1

Notes:

The tool correctly sees through pointer types to get to the underlying type
For both named and unnamed struct types, the tool gets to the struct's fields and lists their names and types
For named types, the tool reports where they are defined (file, line and column) - both for user-defined types and types in the standard library

[1]	For details on the basic usage of `x/tools/go/packages`, see my earlier post on Writing multi-package analysis tools.

[2]	Alternatively, we could iterate directly over the `Decl` fields of the `packages.Package`. I prefer using the inherently recursive `ast.Inspect` because it works for a wider variety of scenarios and will find the declarations of unnamed functions within other functions as well.

[3]	Note: this description does not intend to be complete in the sense of covering the entire Go spec. If you encounter additional cases not covered here, do an AST dump of the nodes you run into and consult the documentation to see how to deal with them.

Rewriting Go source code with AST tooling

2021-10-02T05:48:00-07:00

Go is well-known for having great tooling for analyzing code written in the language, right in the standard library with the go/* packages (go/parser, go/ast, go/types etc.); in addition, the golang.org/x/tools module contains several supplemental packages that are even more powerful. I've used one of these packages to describe how to write multi-package analysis tools in a post from last year.

Here I want to write about a slightly different task: rewriting Go source code using AST-based tooling. I will begin by providing a quick introduction to how existing capabilities of the stdlib go/ast package can be used to find points of interest in an AST. Then, I'll show how some simple rewrites can be done with the go/ast package without requiring additional tooling. Finally, I'll discuss the limitations of this approach and the golang.org/x/tools/astutil package which provides much more powerful AST editing capabilities.

This post assumes some basic level of familiarity with ASTs (Abstract Syntax Trees) in general, and ASTs for Go in particular.

Finding points of interest in a Go AST

Throughout this post, we're going to be using the following simple Go snippet as our lab rat:

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
}

Let's start by finding all calls to the pred function in this code. The go/ast package provides two approaches for finding points of interest in the code. First, we'll discuss ast.Walk. The full code sample for this part is on GitHub. We begin by parsing the source code (which we'll be piping into standard input):

fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "src.go", os.Stdin, 0)
if err != nil {
  log.Fatal(err)
}

Now we create a new value implementing the ast.Visitor interface and call ast.Walk:

visitor := &Visitor{fset: fset}
ast.Walk(visitor, file)

Finally, the interesting part of the code is the Visitor type:

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
}

Our visitor is only interested in AST nodes of type CallExpr. Once it sees such a node, it checks the name of the called function, and reports matches. Note the type assertion on x.Fun; we only want to report calls when the function is referred to by an ast.Ident. In Go, we could call functions in other ways, like invoking anonymous functions directly - e.g. func(){}().

We have a FileSet stored in the visitor; this is only used here to report positions in the parsed code properly. To save space, the AST stores all position information in a single int (aliased as the token.Pos type), and the FileSet is required to translate these numbers into actual positions of the expected <filename>:line:column form.

Visualizing the Go AST

At this point it's worth mentioning some useful tools that help writing analyzers for Go ASTs. First and foremost, the go/ast package has a Print function that will emit an AST in a textual format. Here's how the full if statement in our code snippet would look if printed this way:

.  .  1: *ast.IfStmt {
.  .  .  If: 9:2
.  .  .  Cond: *ast.BinaryExpr {
.  .  .  .  X: *ast.BinaryExpr {
.  .  .  .  .  X: *ast.Ident {
.  .  .  .  .  .  NamePos: 9:5
.  .  .  .  .  .  Name: "x"
.  .  .  .  .  .  Obj: *(obj @ 72)
.  .  .  .  .  }
.  .  .  .  .  OpPos: 9:7
.  .  .  .  .  Op: >
.  .  .  .  .  Y: *ast.BasicLit {
.  .  .  .  .  .  ValuePos: 9:9
.  .  .  .  .  .  Kind: INT
.  .  .  .  .  .  Value: "2"
.  .  .  .  .  }
.  .  .  .  }
.  .  .  .  OpPos: 9:11
.  .  .  .  Op: &&
.  .  .  .  Y: *ast.CallExpr {
.  .  .  .  .  Fun: *ast.Ident {
.  .  .  .  .  .  NamePos: 9:14
.  .  .  .  .  .  Name: "pred"
.  .  .  .  .  .  Obj: *(obj @ 11)
.  .  .  .  .  }
.  .  .  .  .  Lparen: 9:18
.  .  .  .  .  Ellipsis: -
.  .  .  .  .  Rparen: 9:19
.  .  .  .  }
.  .  .  }
.  .  .  Body: *ast.BlockStmt {
.  .  .  .  Lbrace: 9:21
.  .  .  .  List: []ast.Stmt (len = 1) {
.  .  .  .  .  0: *ast.ReturnStmt {
.  .  .  .  .  .  Return: 10:3
.  .  .  .  .  .  Results: []ast.Expr (len = 1) {
.  .  .  .  .  .  .  0: *ast.BasicLit {
.  .  .  .  .  .  .  .  ValuePos: 10:10
.  .  .  .  .  .  .  .  Kind: INT
.  .  .  .  .  .  .  .  Value: "5"
.  .  .  .  .  .  .  }
.  .  .  .  .  .  }
.  .  .  .  .  }
.  .  .  .  }
.  .  .  .  Rbrace: 11:2
.  .  .  }

A somewhat more interactive way to explore this AST dump is using the web page at http://goast.yuroyoro.net/, where you can paste your source and get an AST dump with expandable and collapsible sections. This helps focus only on parts we're interested in; here's an extract from our AST:

(Update: https://astexplorer.net/ is an even nicer AST explorer)

Using the ast.Inspect API

Using ast.Walk for finding interesting nodes is pretty straightforward, but it requires scaffolding that feels a bit heavy for simple needs - defining a custom type that implements the ast.Visitor interface, and so on. Luckily, the go/ast package provides a lighter-weight API - Inspect; it only needs to be provided a closure. Here's our program to find calls to pred() rewritten with ast.Inspect:

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 actual AST node matching logic is the same, but the surrounding code is somewhat simpler. Unless there's a strong need to use ast.Walk specifically, ast.Inspect is the approach I recommend, and it's the one we'll be using in the next section to actually rewrite the AST.

Simple AST rewrites

To begin, it's important to highlight that the AST returned by the parser is a mutable object. It's a collection of node values interconnected via pointers to each other. We can change this set of nodes in any way we wish - or even create a wholly new set of nodes - and then use the go/format package to emit Go formatted source code back from the AST. The following program will simply emit back the Go program it's provided (though it will drop the comments with the default configuration):

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, back to rewriting that AST. Let's make a couple of changes:

We'll rename the function pred to pred2, and rename all the call sites to call the new function name.
We'll inject a printout into the beginning of each function body - emulating some sort of instrumentation we could add this way.

Given the original code snippet, the output will look like this (with the changed/new lines highlighted):

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: we're not adding an import of fmt here - this is left as an exercise for the reader]

The full code of our rewriting program is available here. It's using ast.Inspect to find the nodes it wants to operate on. Here's the renaming of the call sites:

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"
      }
    }
    // ...

If the function is called by an identifier, the code just appends "2" to the name. Again, we're not operating on some copy of the AST - this is the real, living AST we're editing here.

Now let's move on to the next case, where we're handing function declarations:

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

The first three lines in this case do the same as we did for the call sites - just rename the pred function to pred2. The rest of the code is adding the printout to the start of a function body.

That task is fairly easy to accomplish since each FuncDecl has a Body which is an ast.StmtList, which itself holds a slice of ast.Stmt in its List attribute. Out program prepends a new expression to this slice, in effect adding a new statement to the very beginning of the function body. The statement is a hand-crafted AST node. You must be thinking - how did I know how to build this node?

It's really not a big deal once you get the hang of it. Parsing small snippets of code and dumping their ASTs helps, as well as the detailed documentation of the go/ast package. I also found the go2ast tool very useful; it takes a piece of code and emits exactly the Go code needed to build its AST.

Finally, at the end of the program we emit back the modified AST:

fmt.Println("Modified AST:")
format.Node(os.Stdout, fset, file)

And this gets us the modified snippet shown at the beginning of this section.

Limitations of AST editing with Walk and Inspect

So far we've managed to rewrite the AST in a couple of interesting ways using ast.Inspect for finding the nodes. Can we do any kind of rewrite this way?

It turns out the answer to this question is no, or at least not easily. As a motivating example, consider the following task: we'd like to rewrite each call to pred() so that it's logically negated, or turns into !pred(). How do we do that?

It's worth spending a few minutes thinking about this question before reading on.

The issue is that when ast.Inspect (or ast.Walk) hands us an ast.Node, we can change the node's contents and its children, but we cannot replace the node itself. To replace the node itself, we'd need access to its parent, but ast.Inspect does not give us any way to access its parent. A different, slightly more technical way to think about it is: we get handed a node pointer by value, meaning that we can tweak the node it points to, but can't set the pointer to point to a different node. To achieve the latter, ast.Inspect would have to hand us a pointer to a pointer to the node.

This limitation was discussed several years ago, and finally in 2017 a new package appeared in the "extended stdlib" golang.org/x/tools module - astutil.

More powerful rewriting with `astutil`

The APIs astutil provides let us not only find nodes of interest in the AST, but also a way to replace the node itself, not just its contents. In fact, the package provides several useful helpers to delete, replace and insert new nodes through the Cursor type. A full walkthrough of the capabilities of astutil is outside the scope of this post, but I will show how to use it in order to implement our task of turning each pred() into !pred(). Here we go:

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

Instead of calling ast.Inspect, we call astutil.Apply, which also walks the AST recursively and gives our closure access to the node. Apply lets us register a callback for the node both before and after it was visited; in this case we only provide the after case.

Our closure identifies the call to pred in a way that should be similar by now. It then uses the Cursor type to replace this node by a new one which is just the same node wrapped in a unary NOT expression. Hidden in its implementation, the Cursor type does have access to the parent of each node, making it possible to replace the actual node with something else.