This year I got an itch and signed up for an Implementation of PL course, where you build a compiler with your bare hands. I’m jotting down the process here. Honestly, this might be one of the most hardcore CS courses at my school.
The instructor let us choose between Racket and Python, but I just can’t get used to Racket’s parentheses that go on forever, so I went with Python. Luckily, the brand-new Python 3.10 introduced match case, which cut down the amount of code massively! It’s sooo good!!! (I suspect writing the compiler in Rust would be even more fun, though.)
The source code lives in this repo, with each milestone tagged. If you’re interested, go take a look at the code! I think the code quality is… passable?
Leaving this as a work in progress, I’ll fill it in bit by bit! It’s finals review season, so time to actually fill it in.
The course mainly focuses on the compilation process itself (the passes). The stuff that happens before that, like generating the AST, isn’t covered, since there are tools that can generate an AST automatically once you have a grammar definition.
Started: 2021-10-16
Last updated: 2021-12-12
Overview
Rome wasn’t built in a day. – Anonymous
Reinventing the wheel is so much fun. – Me
Features
The same goes for programming languages: we add new features one step at a time.
- $L_{Int}$: integers, arithmetic, and output (
print) - $L_{Var}$: adds variables and assignment, plus reading input (
input_int) - $L_{If}$: adds
booland theifstatement - $L_{While}$: adds
whileloops. There’s nofor, but it can be implemented on top ofwhile - $L_{Tup}$: adds
tuple;listcan be implemented on top of it - $L_{Fun}$: adds
function, with tail call optimization - $L_{Lambda}$: adds
Lambda, with closure support - $L_{Dyn}$: supports dynamic typing
Here is the final supported syntax:
Concrete Syntax
Abstract Syntax
It’s only a small subset of Python, but hey, it’s just a course project.
Passes
As we all learned in second grade, a compiler is made up of a bunch of passes. My implementation doesn’t do any optimization above the function level, so each function is compiled independently. That gives better control over compilation from a higher level.
class Compiler:
"""compile the whole program, and for each function, call methods in `CompileFunction`"""
# TODO: How to ref CompileFunction.arg_passing?
# num_arg_passing_regs = len(CompileFunction.arg_passing)
num_arg_passing_regs = 6
def __init__(self):
self.functions = []
# function -> CompileFunction instance
self.function_compilers = {}
# function -> {original_name: new_name}
self.function_limit_renames = {}
self.num_uniquified_counter = 0
class CompileFunction:
"""compile a single function"""
temp_count: int = 0
tup_temp_count: int = 0
# used for tracking static stack usage
normal_stack_count: int = 0
shadow_stack_count: int = 0
tuple_vars = []
# `calllq`: include first `arg_num` registers in its read-set R
arg_passing = [Reg(x) for x in arg_passing]
# `callq`: include all caller_saved registers in write-set W
caller_saved = [Reg(x) for x in caller_saved]
callee_saved = [Reg(x) for x in callee_saved]
builtin_functions = ['input_int', 'print', 'len']
builtin_functions = [Name(i) for i in builtin_functions]
def __init__(self, name: str):
self.name = name
self.basic_blocks = {}
# mappings from a single instruction to a set
self.read_set_dict = {}
self.write_set_dict = {}
self.live_before_set_dict = {}
self.live_after_set_dict = {}
# this list can be changed for testing spilling
self.allocatable = [Reg(x) for x in allocatable]
all_reg = [Reg('r11'), Reg('r15'), Reg('rsp'), Reg(
'rbp'), Reg('rax')] + self.allocatable
self.int_graph = UndirectedAdjList()
self.move_graph = UndirectedAdjList()
self.control_flow_graph = DirectedAdjList()
self.live_before_block = {}
self.prelude_label = self.name
# assign this when iterating CFG
self.conclusion_label = self.name + 'conclusion'
self.basic_blocks[self.conclusion_label] = []
# make the initial conclusion non-empty to avoid errors
# TODO: come up with a more elegant solution, maybe from `live_before_block`
# self.basic_blocks[self.conclusion_label] = [Expr(Call(Name('input_int'), []))]
# why need this?
self.sorted_control_flow_graph = []
self.used_callee = set()
self.stack_frame_size: int
self.shadow_stack_size: int
# Reserved registers
self.color_reg_map = {}
color_from = -5
for reg in all_reg:
self.color_reg_map[color_from] = reg
color_from += 1
shrink
Gather up all the statements scattered around the source file that don’t belong to any function and whip them into a main.
def shrink(self, p: Module) -> Module:
"""create main function, making the module body a series of function definitions"""
assert(isinstance(p, Module))
# main_args = arguments([], [], [], [], [])
main = FunctionDef('main', args=[], body=[],
decorator_list=[], returns=int)
new_module = []
for c in p.body:
match c:
case FunctionDef(name, _args, _body, _deco_list, _rv_type):
new_module.append(c)
self.functions.append(name)
# print("DEBUG, function: ", name, "args: ", args.args, body, _deco_list, _rv_type)
case stmt():
main.body.append(c)
main.body.append(Return(Constant(0)))
new_module.append(main)
print("DEBUG, new module: ", new_module)
return Module(new_module)
The current implementation doesn’t support global variables, but I don’t think it’d be too hard. You could put all the variable definitions scattered around the source file on the heap so other functions can access them.
uniquify
Functions can have variable shadowing. To make things unambiguous, we just rename each variable where needed.
def uniquify(self, p: Module) -> Module:
class Uniquify(NodeTransformer):
def __init__(self, outer: Compiler, mapping: dict):
self.outer_instance = outer
self.uniquify_mapping = mapping
super().__init__()
def visit_Lambda(self, node):
self.generic_visit(node)
match node:
case Lambda(args, body_expr):
new_mapping = self.uniquify_mapping.copy()
new_args = []
for v in args:
new_v = v + "_" + str(self.outer_instance.num_uniquified_counter)
# find the new name in the previous mapping
new_mapping[new_mapping[v]] = new_v
# delete the old mapping
del new_mapping[v]
new_args.append(new_v)
self.outer_instance.num_uniquified_counter += 1
new_uniquifier = Uniquify(self.outer_instance, new_mapping)
new_body_expr = new_uniquifier.visit(body_expr)
return Lambda(new_args, new_body_expr)
case _:
return node
def visit_Name(self, node):
self.generic_visit(node)
match node:
case Name(id):
if id in self.uniquify_mapping:
return Name(self.uniquify_mapping[id])
else:
return node
case _:
return node
def do_uniquify(stmts: list, uniquify_mapping: dict) -> list:
"""change the variable names of statements in place according to the uniquify_mapping"""
uniquifier = Uniquify(self, uniquify_mapping)
new_body = []
for s in stmts:
new_body.append(uniquifier.visit(s))
return new_body
assert(isinstance(p, Module))
for f in p.body:
assert(isinstance(f, FunctionDef))
uniquify_mapping = {}
new_args = []
for v in f.args:
print("DEBUG, v: ", v[0], "type: ", type(v[0]))
new_arg_name = v[0] + "_" + str(self.num_uniquified_counter)
uniquify_mapping[v[0]] = new_arg_name
new_args.append((new_arg_name, v[1], ))
self.num_uniquified_counter += 1
f.args = new_args
f.body = do_uniquify(f.body, uniquify_mapping)
return p
reveal functions
In Python, functions are first-class citizens and can be used like variables, so at the AST level you can’t tell them apart. But assigning a function to a variable means something different from assigning anything else (e.g. a number, or another variable’s value): the former actually takes the function’s address, not its body. So we need to tell the two apart by converting every reference to a function defined in the source from Name(XXX) to FunRef(XXX). (FunRef is an AST node the instructor hand-rolled.)
def reveal_functions(self, p: Module) -> Module:
"""change `Name(f)` to `FunRef(f)` for functions defined in the module"""
class RevealFunction(NodeTransformer):
def __init__(self, outer: Compiler):
self.outer_instance = outer
super().__init__()
def visit_Call(self, node):
self.generic_visit(node)
match node:
case Call(Name(f), args) if f in self.outer_instance.functions:
# what if f is a builtin function? guard needed
return Call(FunRef(f), args)
case _:
return node
assert(isinstance(p, Module))
# Why this does't work?
# new_body = RevealFunction(self).visit_Call(p)
# p.body = new_body
for f in p.body:
new_body = []
for s in f.body:
new_line = RevealFunction(self).visit_Call(s)
new_body.append(new_line)
# print("DEBUG, new node: ", ast.dump(n))
f.body = new_body
return p
Convert Assignments
TODO
Convert to Closure
TODO
Limit Function
When a function has more than six parameters, the extra ones can’t be passed in registers. The usual approach is to put them on the stack, but to enable tail call optimization, we put them on the shadow stack (actually the heap) instead. Concretely: create a tuple holding the sixth and all remaining arguments, and pass that tuple as the sixth argument.
def limit_functions(self, p: Module) -> Module:
"""limit functions to 6 arguments, anything more gets put into a 6th tuple-type argument"""
class LimitFunction(NodeTransformer):
# limit call sites & convert name of args
def __init__(self, outer: Compiler, mapping: dict = {}):
self.outer_instance = outer
self.mapping = mapping
super().__init__()
def visit_Name(self, node):
# substitute the >5th arguments with subscript of a tuple
self.generic_visit(node)
match node:
case Name(n) if n in self.mapping.keys():
return self.mapping[n]
case _:
return node
def visit_Call(self, node):
self.generic_visit(node)
match node:
case Call(FunRef(f), args) if len(args) > Compiler.num_arg_passing_regs:
# print("DEBUG, HIT in visit_FunRef: ", ast.dump(node))
new_args = args[:Compiler.num_arg_passing_regs - 1]
new_args.append(
Tuple(args[Compiler.num_arg_passing_regs - 1:], Load()))
return Call(FunRef(f), new_args)
case _:
return node
def args_need_limit(args):
if isinstance(args, list):
# print("DEBUG, args: ", args, type(args))
# print("DEBUG, argsAST: ", args[1][0])
return len(args) > Compiler.num_arg_passing_regs
# else:
# print("DEBUG, args: ", ast.dump(args))
return False
assert(isinstance(p, Module))
# limit defines
for f in p.body:
match f:
case FunctionDef(_name, args, _body, _deco_list, _rv_type) if args_need_limit(args):
# args = args.args
new_args = args[:5]
arg_tup = ('tup_arg', TupleType([a[1] for a in args[5:]]))
alias_mapping = {}
for i in range(5, len(args)):
alias_mapping[args[i][0]] = Subscript(
Name('tup_arg'), Constant(i - 5), Load())
print("DEBUG, alias_mapping: ", alias_mapping)
new_args.append(arg_tup)
print("DEBUG, new_args: ", new_args)
f.args = new_args
new_body = []
for s in f.body:
# new_line is new node
new_line = LimitFunction(
self, alias_mapping).visit_Name(s)
print("DEBUG, new_line: ", new_line)
new_body.append(new_line)
print("DEBUG, new_body: ", new_body)
f.body = new_body
case _:
assert(isinstance(f, FunctionDef))
new_body = []
for s in f.body:
new_line = LimitFunction(self).visit_Call(s)
new_body.append(new_line)
f.body = new_body
return p
Since tail call optimization wipes out the current function’s stack frame, arguments can no longer be passed to the callee via the stack. That’s why the extra arguments get dumped on the shadow stack.
That said, I feel like you could also wipe the current stack, then set up the callee’s arguments and jump straight into the callee (skipping its stack initialization). But the instructor pointed out that preparing the arguments might need stuff from the current stack frame, so things could get messed up. There always seems to be a way around it, though: e.g. first extend the current stack frame and build the callee’s arguments somewhere clean, then set up the callee’s frame by popping your own frame and copying the callee’s arguments over. No idea how production compilers do it; there may well be a more efficient trick.
Expose Allocation
We store tuples on the heap, and keep pointers to them on a shadow stack (root stack). This makes life easier for the garbage collector (more on that later). This pass only deals with tuples: it turns a tuple’s initialization into a series of statements, followed by an expression that yields the tuple at the end.
def expose_allocation_hide(self, t: Tuple) -> Begin:
# Autograder call `expose_allocation` in a wrong way, so this is hidden from the autograder
"""convert a tuple creation into a begin"""
assert(isinstance(t, Tuple))
content = t.elts
body = []
for i in range(len(content)):
if not CompileFunction.is_atm(content[i]):
print("DEBUG: ???")
temp_name = 'temp_tup' + \
str(self.tup_temp_count) + 'X' + str(i)
body.append(Assign([Name(temp_name)], content[i]))
tup_bytes = (len(content) + 1) * 8
if_cond = Compare(BinOp(GlobalValue('free_ptr'), Add(), Constant(tup_bytes)), [
Lt()], [GlobalValue('fromspace_end')])
# TODO: Expr(Constant(0)) OK here?
body.append(If(if_cond, [], [Collect(tup_bytes)]))
var = Name("pyc_temp_tup_" + str(self.tup_temp_count))
body.append(Assign([var], Allocate(len(content), t.has_type)))
for i in range(len(content)):
if not CompileFunction.is_atm(content[i]):
body.append(Assign([Subscript(var, Constant(i), Store())], Name(
'temp_tup' + str(self.tup_temp_count) + 'X' + str(i))))
else:
body.append(
Assign([Subscript(var, Constant(i), Store())], content[i]))
self.tup_temp_count += 1
return Begin(body, var)
One thing to watch out for: everything from allocating the memory to finishing the tuple’s construction has to be atomic. The allocated memory is uninitialized, so if some other operation (which also allocates) sneaks in between, you get memory corruption, most likely because some global pointers go flying off.
Remove Complex Op*
This pass breaks complex operations and operands down into several simpler operations. For example, the operands of an addition need to be atoms (e.g. constants/variables).
The code for this one is way too long, so check the repo for the details.
def rco_exp(self, e: expr, need_atomic: bool) -> typing.Tuple[expr, Temporaries]:
temps = []
# tail must be assigned in the match cases
if CompileFunction.is_atm(e):
"""nothing need to do if it's already an `atm`"""
return (e, temps)
match e:
case XXX:
tail = YYY
return (tail, temps)
def rco_stmt(self, s: stmt) -> List[stmt]:
result = []
temps = []
match s:
case XXX:
self.rco_exp(XXX)
# late binding
for binding in temps:
# print("DEBUG, binding: ", binding)
result.append(Assign([binding[0]], binding[1]))
This saves a return line in every case, I guess. The calls here need to be recursive.
All in all, there’s still plenty of room for improvement in the code, and some redundancy too, but I tried my best to strike a balance between readability and simplicity. At least I think the code I wrote is pretty damn awesome.
Explicate Control
This is where we have to lay out the program’s control flow sensibly. As we all learned in third grade, once a function has multiple control-flow paths, you need basic blocks to split up the branches and loops, and that’s exactly what explicate control does. This pass is pretty bloated too, so I won’t paste the code. At a high level, though, it breaks down into five main parts:
explicate_pred: handlesifandwhileloops. Nestedifs need to be handled here tooexplicate_assign: handles assignments.explicate_tail: handlesreturn. This is where tail call optimization can happen.explicate_effect: handles other expressions that only have side effects, e.g.print(), or aninput()that isn’t assigned to a variable (the latter is usually used to pause a program).explicate_stmt: dispatches a statement to one of the cases above.
On top of that, I implemented a lazy evaluation optimization, which cuts down on trivial basic blocks that contain nothing but an unconditional jmp.
Select Instruction
This is mostly brute force: find the right x86_64 instructions for each expression. Too long to paste.
I pulled a neat trick here. I don’t know if there’s an official name for it in compiler land, but I’ll borrow the name late binding: the rhs of every Assign statement is first bound to an Unnamed_Pyc_Var in select_expr(), and then bound according to the lhs via bound_unamed() in rco_stmt(). This massively reduces redundant code (the instructor’s code was way too complicated, it gave me a headache = =).
I wish the AST were designed to be a bit simpler; it’s still a little too complicated as is.
Register Allocation
Liveness Analysis
Interference Graph
Graph Coloring
Patch Instruction
Prelude and Conclusion
Something Interesting
Here are some other techniques I used. Not sure I’ll have time to go into detail, but I’m staking out the sections for now.
Dynamic Typing
Garbage Collection
Since tuples live on the heap, we need GC. For that, there’s a root stack (shadow stack) that holds all the tuple pointers.
Here is what memory looks like:
The reason all tuples get dumped on the shadow stack: the GC needs a starting point to collect garbage from, and the shadow stack is that starting point. To be precise, everything that needs garbage collection should go there.
The GC finds every element reachable from the shadow stack and makes sure they don’t get collected; anything it can’t reach is out of luck. Reachable elements include those pointed to by pointers in the root stack, plus, recursively, whatever the pointers inside those elements point to. That raises some questions: what if an element is pointed to twice? What if some pointers form a cycle?
So here’s what a tuple actually looks like in memory:
The tuple’s first double word stores some metadata, including:
- pointer mask: marks which fields in the tuple are pointers
- vector length: the tuple’s length
- forwarding: marks whether it’s already been processed while the GC is running
The GC algorithm is Two-Space Copying. It’s pretty self-explanatory from the name, so I won’t go on about it here; see the Wiki for details. I read through its C implementation, and it was a real pain. The code isn’t that complicated, but if I didn’t already know the algorithm, I’d have had a hard time telling it was a GC at all (probably just my lack of experience).
Bringing in GC means a few extra chores: initialize the shadow stack when the program starts; reserve a register (r15) as the shadow stack pointer; reserve another register (r11) as the pointer for tuple access (since x86 requires memory writes to go through a register, and rax is already used in instruction patching to avoid having both operands in memory, so it can’t be reused for tuples); and before allocating, check whether there’s enough space, and allocate more if not.