diff --git a/lang/compile.py b/lang/compile.py
index 4e68784..4673531 100755
--- a/lang/compile.py
+++ b/lang/compile.py
@@ -305,6 +305,57 @@
 		ir.append(sub_ir)
 	return ir
 
+## Register allocation
+
+def find_variables(ir):
+	vars = {'Return': 0}
+	var_count = 0
+	for node in ir:
+		if not isinstance(node, IRLoad):
+			continue
+		var = node.variable
+		if not vars.get(var):
+			var_count += 1
+			vars[var] = var_count
+	return vars
+
+def replace_variables(ir, variables):
+	new_ir = []
+	var_count = len(variables)
+	new_ir.append(IRAllocate(var_count))
+	for node in ir:
+		if isinstance(node, IRLoad):
+			reg = variables[node.variable]
+			new_ir.append(IRLoad(reg))
+		elif isinstance(node, IRStore):
+			reg = variables[node.variable]
+			new_ir.append(IRStore(reg))
+			continue
+		elif isinstance(node, IRReturn):
+			new_ir.append(IRDrop(var_count))
+			new_ir.append(IRReturn())
+		else:
+			new_ir.append(node)
+	return new_ir
+
+def registers_allocate_func(ir):
+	registers = find_variables(ir.statements)
+	new_ir = replace_variables(ir.statements, registers)
+	name = ir.name
+	return IRFunction(name, new_ir)
+
+def registers_allocate(ir):
+	new_ir = []
+	for node in ir:
+		sub_ir = None
+		if isinstance(node, IRFunction):
+			sub_ir = registers_allocate_func(node)
+		if not sub_ir:
+			print("Unknown registers node: %s" % (node))
+			return None
+		new_ir.append(sub_ir)
+	return new_ir
+
 ## Tests
 
 code = """
@@ -319,4 +370,5 @@
 lines = tokenize(code)
 ast = parse_toplevel(lines)
 ir = generate_ir(ast)
-print(ir)
+ir_reg = registers_allocate(ir)
+print(ir_reg)