Adding "four" loops to the Go Compiler - Part 1
I gave a talk at GopherCon 2024 based on this post. The slides for the talk can be found here, and the generated SSA file can be found here.
Recently, I watched this stream where George Hotz adds "four" loops to Clang, which made me want to try adding "four" loops to Go as well. This write-up will describe the details of how to add both "four" loops and "unless" statements to the Go compiler. There will be a second write-up that will go over how to use the compiler tooling to further investigate these changes and see the optimizations the compiler back-end makes.
Disclaimer: This is just for learning purposes, I don't think these would actually be good additions to Go.
So what is a four loop? A four loop is just like a regular loop, except its post statement is executed 4 times at the end of each iteration instead of just once. For example, the following two code snippets are equivalent:
unless statements exist in languages like Ruby and Perl. They are meant to help application logic to read more naturally in English. They act as negated if statements, where the body is executed if the condition is false and skipped if it is true. For example, the following two code snippets are equivalent:
Putting these together, here are the programs that we'll use to test our changes:
which should both output:
To implement this, we'll first take a look at the phases of the Go compiler:
- Lexical Analysis and Parsing
- Type Checking
- IR Construction
- Middle End
- Walk
- SSA Generation
- Machine Code Generation
The first 3 phases are the "front-end", the next 2 are the "middle-end", and the last 2 are the "back-end". The front-end is responsible for taking the source code and turning it into an intermediate representation (IR) that the middle-end can work with. The middle-end is responsible for optimizing the IR and preparing it for the back-end. Finally, the back-end is responsible for further optimizing the IR and turning it into machine code. We'll go through each of these phases and see what changes are needed at each step.
Before we can do that, we first need to get the Go compiler source code and build it with ./make.bash. It may take up to a few minutes to build.
This should install the compiler in go/bin. You can verify it's working by running bin/go run demo1.go. You should also test it on bin/go run demo2.go to see that we get syntax errors, as the compiler doesn't recognize the four and unless keyword yet.
From this point onwards we can do quicker rebuilds of the compiler with
../bin/go install cmd/compile from the go/src directory.
I recommend using this alias:
Now we have everything set up to start making changes to the compiler. All the files referenced are in the go/src/cmd/compile/internal/ directory unless otherwise specified.
Lexical Analysis and Parsing
The first phase of the compiler is lexical analysis and parsing. This is where the source code is scanned character by character and tokenized. The first change we need to make is in syntax/tokens.go
Those comments are important, becaues they are used by the go generate tool create other code we need. To install it, run:
And then run it on the file we just changed:
which should update syntax/token_string.go. If we run the recompile alias from earlier and try run bin/go run demo1.go, we get a panic:
The scanner uses a perfect hash function based on the first two characters and the length of the token to generate a keyword map. Our addition of two new keywords breaks this hash function, but it can be worked around by increasing the size of the keyword map from 1 << 6 to 1 << 7.
Now the four and unless keywords can be scanned and tokenized, but we don't know how to parse them yet. The structure for both should be familiar:
To parse these, we need to make changes to syntax/nodes.go, syntax/parser.go, and syntax/walk.go.
Type Checking
The second phase of the compiler is type checking. It's done in several steps and does things like:
- Name resolution: mapping identifiers to language objects
- Constant folding: computing compile-time constants
- Type inference: computing the type of every expression and checking for compliance with language specification
The type checking changes we need to make are in types2/stmt.go.
At this point, if we recompile our code and run it on demo2.go we get:
which shows we've gone from a syntax error to an internal compiler error after these changes. If you place a call to Fdump in the parser, you can can take a look at what the syntax tree looks like before the internal compiler error is encountered.
The final phase of the front-end is IR construction, or as it's referred to in the Go compiler, "noding". This is where we convert the type checked syntax tree into an abstract syntax tree. There are quite a few changes needed in this phase. First, update the following files:
At this point we need to do some code generation again. This time, run:
which should update ir/op_string.go and ir/node_gen.go. There's still a few more files we need to update:
Then just run go generate sync.go, and that's everything that's needed to convert the concrete syntax tree into an abstract syntax tree.
Middle End
The middle end performs several optimization passes on the AST:
- Dead code elimination
- Interface method devirtualization
- Function inlining
- Escape analysis
There are escape analysis changes needed for both the four and unless statements.
Walk
The walk is the final pass over the AST. It has 2 steps:
- "Order" step: convert complex statements into simpler ones, introducing temporary variables and respecting order of evaluation.
- "Desugar" step: convert high level constructs into more primitive ones; e.g. unless statements into negated if statements.
The order step isn't very interesting:
But the desugar step is where we'll convert unless statements into negated if statements. We do this by wrapping the condition in a NOT expression, and then returning a normal if statement with the negated condition.
At this point our unless statement should be working correctly, but the four statement is still not. This concludes the middle end, and now the AST is passed on to the back-end phases of the compiler where the four statement will be implemented.
SSA Generation
The first phase of the back-end is SSA generation. This is when the AST is converted to an IR in Static Single Assignment form. You should read about SSA here and here. The program is rewritten so that each variable is only assigned once. This allows us to break the program up into blocks, where each block contains a set of values/operations and has only 1 exit point at the end of the block. The blocks form control flow graph for the program. For example, demo2.go's CFG looks like this:
The code to actually build these blocks for four and unless statements is below:
The relevant parts are:
where we generate SSA code for the four loop's post statement 4 times, and:
where we generate the SSA code for the unless condition, except we've set it to go to the end block if the condition is true, and to go to the then statement if the condition is false.
Having the program in this form makes it easier to perform optimizations and translate the program into machine code. A series of machine independent optimizations are performed before we move on to the final phase of the compiler.
Machine Code Generation
The final phase of the compiler is machine code generation. It starts with a pass called "lower", which begins rewriting the SSA with machine specific instructions. Then a series of machine dependent optimization passes are run:
Things like stack frame layout and pointer liveness analysis are also handled here. The final output is passed to the assembler to generate the actual machine code, which will then be linked and then can be executed. We don't have any changes to make to this phase of the compiler which means we are all done, and should be able to test our changes out now. Recompiling with all the changes and running bin/go run demo2.go should produce the expected output:
So we've successfully added both the four and unless statements to the Go compiler. In the next write-up, we'll take a look at some of the compiler tooling that Go provides to help us analyze our changes and see what sort of optimizations the compiler makes. Thanks for reading.