Instrumenting Go Code via AST - Part 1

Welcome! 

This is the second part of our AST blog post series, expanding on the subject of utilizing Go AST libraries to automate and improve your workflow.

In this post I’ll discuss a rather common problem that comes up while working with Go code and the way we’ve solved it by sprinkling a little bit of AST magic dust. Let’s dive in.

Problem: A struct with no interface 

Let’s say you are working on a large code base that was not built with interfaces in mind, meaning, there are structs and methods attached to those structs, but there is no interface describing it. This is a perfectly valid approach when you don’t need to mock/stub the method implementations provided by that struct, or there’s only one implementation of the same ‘contract’.

However, when these things are required we need to provide an interface.

Here’s a small code snippet to demonstrate:

type Person struct {
	Name string
}

func (p Person) Hello() string {
	return "Hello: " + p.Name
}

type Animal struct {
	Legs int
}

func (a Animal) Hello() string {
	return fmt.Sprintf("I have %d legs!", a.Legs)
}

func main() {
	p := Person{Name: "Fred"}
	a := Animal{Legs: 4}
	// ...
}

In this example both Person and Animal have the same Hello() method. If we wanted to store a list of both the Person and Animal structs, we would have to define it as:

list := []interface{}{p,a}

But this way we lose the type information of the list elements. This is where interfaces come in. Since both Person and Animal implement a method with the same signature, we can extract that signature into an interface and use it for storing items in a list:

type interface Hello {
	Hello() string
}
// ...
list := []Hello{p,a}
fmt.Printf("Person: [%v] Animal: [%v]\n", list[0].Hello(), list[1].Hello())

Awesome. Now let’s say the struct you are extracting the interface from is a big one. A really big one. With lots and lots of methods spread out in different .go files. Creating such an interface manually would be very laborious.

This problem is itching to get an AST treatment. Let’s get to it!

AST to the rescue! 

Let’s break down the task at hand into smaller, digestable parts:

1. Scan the source code for all methods implemented on a specific struct
2. Collect all those methods (and their comments, for clarity)
3. Create a new file that will contain an interface with all the collected methods inside
4. …
5. Profit?

Scanning the source code for all struct methods 

Here’s a short piece of code that scans a folder of .go source code to first find a package by name and then search for all methods that are bound to the struct we’re interested in.

fset := token.NewFileSet()
// 1. scan source code folder
pkgs, err := parser.ParseDir(fset, folder, nil, parser.AllErrors|parser.ParseComments)
if err != nil {
	log.Fatalf("Unable to parse %s folder", folder)
}
// 2. find the required package by name
var appPkg *ast.Package
for _, pkg := range pkgs {
	if pkg.Name == pkgName {
		appPkg = pkg
		break
	}
}
if appPkg == nil {
	log.Fatalf("Unable to find package %s", pkgName)
}

// 3. find all methods that are bound to the specific struct
for _, file := range appPkg.Files {
	ast.Inspect(file, func(n ast.Node) bool {
		if fun, ok := n.(*ast.FuncDecl); ok {
			// 4. Validate that method is exported and has a receiver
			if fun.Name.IsExported() && fun.Recv != nil && len(fun.Recv.List) == 1 {
					// 5. Check that the receiver is actually the struct we want
					if r, rok := fun.Recv.List[0].Type.(*ast.StarExpr); rok && r.X.(*ast.Ident).Name == structName {
						// we found it!
					}
				}
			}

		}
		return true
	})
}

Steps 1 and 2 are pretty straightforward, the interesting bits start at step 3: For each file in the package, we execute ast.Inspect to get all the AST nodes. For every node that is actually a function (checked by n.(*ast.FuncDecl)), we check if that function is:

• Exported (we are not interested in private methods)
• Has a receiver (it’s bound to a struct)
• Receiver’s type matches the struct we are interested in

Collecting functions and their comments 

Now that we can get a *ast.FuncDecl (let’s call it fun) for each function, we can collect all the information about it and reconstruct it:

1. Name: fun.Name.Name
2. Comments: fun.Doc.List
3. Parameters: fun.Type.Params.List
4. Return values: fun.Type.Results.List

Generate the interface file 

To generate the output interface, we’ll use Go’s template package. Let’s define a simple template:

const outputTemplate = `
// DO NOT EDIT, auto generated

package {{.Package}}

type {{.Name}} interface {
	{{.Content}}
}
`

And populate it:

sort.Strings(funcs)
out := bytes.NewBufferString("")

t := template.Must(template.New("").Parse(outputTemplate))
err = t.Execute(out, map[string]interface{}{
	"Content": strings.Join(funcs, "\n"),
	"Name":    ifName,
	"Package": pkgName,
})

We’re almost done! Our interface file is ready, but it’s missing a crucial part - imports. Luckily, there is a package for that™ - https://pkg.go.dev/golang.org/x/tools/imports.

Let’s give it whirl:

formatted, err := imports.Process(outputFile, out.Bytes(), &imports.Options{Comments: true})
if err != nil {
	log.Panic(err)
}
err = ioutil.WriteFile(outputFile, formatted, 0644)

Voila! You have an interface file that contains all the methods implemented on the struct.

struct2interface 

The process I’ve described is a rather generic one, so to make it fully repeatable, I’ve extracted it into a separate CLI utility called struct2interface. It can be found at https://github.com/reflog/struct2interface.

Conclusion 

Once again, AST came to the rescue and saved us lots and lots of manual labor. Hurray! In the next post of this series, I’ll describe how we combined our layers approach with AST to create a clean opentracing instrumentation we’ve started in the first part of these series.