Go Generate and AST
I recently have been playing around with Go’s built in generator for the standard go compiler. I love the idea of a language having reflection, but what I love more is a way for the language to have a way to generate boiler plate code. The benefit of generated code over reflection is performance and compile time checks. In reflection you’re often fiddling with non compile checking strings, you may have fields or functions that get completely stripped from not being used, and other things you just have to infer.
Go Generate
The key to generating Go code is to use the built in //go:generate
comment. You can use this anywhere in the file, but I like to use it as a decorator for things like structs. For example, in the following:
//go:generate go run ../generators/serialized/main.go
type SerializableThing struct {
number int
name string
position [4]float32
}
In this case you’ll notice I’ve placed //go:generate go run ../generators/serialized/main.go
right above the struct. The reason for this is that the file, line number, and package are passed as environment variables (GOFILE
, GOLINE
, and GOPACKAGE
respectively). This allows us to snipe the exact struct based on the provided line number in the file.
Another great feature you may have noticed is that I didn’t need to compile a separate executable constantly, I can just use go run ../generators/serialized/main.go
. This allows me to have the generator code directly in my project and to call it up whenever it is needed. Once you’ve got all this setup, you can use:
fs := token.NewFileSet()
ast, err := parser.ParseFile(fs, filePath, orFileSrc, parser.ParseComments)
From here, the world is your oyster, you can go ahead and write out your code generator. Note that I like to use parser.ParseComments
as I put extra generation context in comments. Just remember to put a comment at the top of your generated file like // Code generated by "serialized"; DO NOT EDIT.
, this way the go tools can help with warning you or others about not editing the file.
Walking the AST
As a bonus, I’ll go over a little bit on how you can walk the AST from this generated code. Let’s say that we just wanted to get the specific struct within the file that is directly below the generate comment. Here is a simple way of doing that:
file := os.Getenv("GOFILE")
//pkg := os.Getenv("GOPACKAGE")
fs := token.NewFileSet()
pos := 1 // The linear position of this generate comment
lineNum, _ := strconv.Atoi(os.Getenv("GOLINE"))
fileSrc, _ := os.ReadFile(file)
reader := strings.NewReader(string(fileSrc))
for lineNum > 0 {
if r, _, _ := reader.ReadRune(); r == '\n' {
lineNum--
}
pos++
}
astRes, _ := parser.ParseFile(fs, "", fileSrc, parser.ParseComments)
var structSpec *ast.TypeSpec = nil
for _, d := range astRes.Decls {
if d.Pos() == token.Pos(pos) {
if g, ok := d.(*ast.GenDecl); ok {
structSpec = g.Specs[0].(*ast.TypeSpec)
}
}
}
log.Println(structSpec.Name.Name)
This wasn’t pretty, but it’ll get the job done. Baislly you need to find the position of the comment by looking through the source code. Once the position is found, you need to match it up to a position of a declaration in the ast. If you were to want to get all structs within the entire file, rather than just the one directly under the comment, you can just remove all the position stuff, which means you don’t need to read the file source and you can just use the file directly.
file := os.Getenv("GOFILE")
//pkg := os.Getenv("GOPACKAGE")
fs := token.NewFileSet()
astRes, _ := parser.ParseFile(fs, file, "", parser.ParseComments)
allStructs := []*ast.TypeSpec{}
for _, d := range astRes.Decls {
if g, ok := d.(*ast.GenDecl); ok {
allStructs = append(allStructs, g.Specs[0].(*ast.TypeSpec))
}
}
log.Println(len(allStructs))