From 12c65a42591eb4843b2b83df1cea7ecb73427012 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <rui314@gmail.com>
Date: Sat, 1 Sep 2018 07:48:46 +0000
Subject: [PATCH] Fix register allocator.

---
 9cc.h      | 13 +++++++++
 gen_ir.c   | 15 ++++++++++-
 irdump.c   | 44 ++++++++++++++++++++++++++----
 liveness.c | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 main.c     |  4 ++-
 regalloc.c |  6 +++++
 util.c     | 14 ++++++++++
 7 files changed, 168 insertions(+), 7 deletions(-)
 create mode 100644 liveness.c

diff --git a/9cc.h b/9cc.h
index 822819d..dae15f1 100644
--- a/9cc.h
+++ b/9cc.h
@@ -25,6 +25,8 @@ void vec_push(Vector *v, void *elem);
 void vec_pushi(Vector *v, int val);
 void *vec_pop(Vector *v);
 void *vec_last(Vector *v);
+bool vec_contains(Vector *v, void *elem);
+bool vec_union1(Vector *v, void *elem);
 
 typedef struct {
   Vector *keys;
@@ -354,6 +356,13 @@ typedef struct BB {
   int label;
   Vector *ir;
   Reg *param;
+
+  // For liveness analysis
+  Vector *succ;
+  Vector *pred;
+  Vector *def_regs;
+  Vector *in_regs;
+  Vector *out_regs;
 } BB;
 
 typedef struct {
@@ -387,6 +396,10 @@ typedef struct {
 
 void gen_ir(Program *prog);
 
+/// liveness.c
+
+void liveness(Program *prog);
+
 /// regalloc.c
 
 void alloc_regs(Program *prog);
diff --git a/gen_ir.c b/gen_ir.c
index 964d169..9ba3ec8 100644
--- a/gen_ir.c
+++ b/gen_ir.c
@@ -19,6 +19,11 @@ static BB *new_bb() {
   BB *bb = calloc(1, sizeof(BB));
   bb->label = nlabel++;
   bb->ir = new_vec();
+  bb->succ = new_vec();
+  bb->pred = new_vec();
+  bb->def_regs = new_vec();
+  bb->in_regs = new_vec();
+  bb->out_regs = new_vec();
   vec_push(fn->bbs, bb);
   return bb;
 }
@@ -430,10 +435,15 @@ static void gen_param(Var *var, int i) {
 void gen_ir(Program *prog) {
   for (int i = 0; i < prog->funcs->len; i++) {
     fn = prog->funcs->data[i];
-    out = new_bb();
 
     assert(fn->node->op == ND_FUNC);
 
+    // Add an empty entry BB to make later analysis easy.
+    out = new_bb();
+    BB *bb = new_bb();
+    jmp(bb);
+    out = bb;
+
     // Emit IR.
     Vector *params = fn->node->params;
     for (int i = 0; i < params->len; i++)
@@ -441,6 +451,9 @@ void gen_ir(Program *prog) {
 
     gen_stmt(fn->node->body);
 
+    // Make it always ends with a return to make later analysis easy.
+    new_ir(IR_RETURN)->r2 = imm(0);
+
     // Later passes shouldn't need the AST, so make it explicit.
     fn->node = NULL;
   }
diff --git a/irdump.c b/irdump.c
index 9238e19..e96b233 100644
--- a/irdump.c
+++ b/irdump.c
@@ -90,6 +90,44 @@ static char *tostr(IR *ir) {
   }
 }
 
+static void print_rel(char *name, Vector *v) {
+  if (v->len == 0)
+    return;
+  fprintf(stderr, " %s=", name);
+  for (int i = 0; i < v->len; i++) {
+    BB *bb = v->data[i];
+    if (i > 0)
+      fprintf(stderr, ",");
+    fprintf(stderr, ".L%d", bb->label);
+  }
+}
+
+static void print_regs(char *name, Vector *v) {
+  if (v->len == 0)
+    return;
+  fprintf(stderr, " %s=", name);
+  for (int i = 0; i < v->len; i++) {
+    Reg *r = v->data[i];
+    if (i > 0)
+      fprintf(stderr, ",");
+    fprintf(stderr, "r%d", regno(r));
+  }
+}
+
+static void print_bb(BB *bb) {
+  if (bb->param)
+    fprintf(stderr, ".L%d(r%d)", bb->label, regno(bb->param));
+  else
+    fprintf(stderr, ".L%d", bb->label);
+
+  print_rel("pred", bb->pred);
+  print_rel("succ", bb->succ);
+  print_regs("defs", bb->def_regs);
+  print_regs("in", bb->in_regs);
+  print_regs("out", bb->out_regs);
+  fprintf(stderr, "\n");
+}
+
 void dump_ir(Vector *irv) {
   for (int i = 0; i < irv->len; i++) {
     Function *fn = irv->data[i];
@@ -97,11 +135,7 @@ void dump_ir(Vector *irv) {
 
     for (int i = 0; i < fn->bbs->len; i++) {
       BB *bb = fn->bbs->data[i];
-
-      if (bb->param)
-        fprintf(stderr, ".L%d(r%d):\n", bb->label, regno(bb->param));
-      else
-        fprintf(stderr, ".L%d:\n", bb->label);
+      print_bb(bb);
 
       for (int i = 0; i < bb->ir->len; i++) {
         IR *ir = bb->ir->data[i];
diff --git a/liveness.c b/liveness.c
new file mode 100644
index 0000000..70c80f1
--- /dev/null
+++ b/liveness.c
@@ -0,0 +1,79 @@
+// Liveness analysis.
+
+#include "9cc.h"
+
+// Fill bb->succ and bb->pred.
+static void add_edges(BB *bb) {
+  if (bb->succ->len > 0)
+    return;
+  assert(bb->ir->len);
+
+  IR *ir = bb->ir->data[bb->ir->len - 1];
+
+  if (ir->bb1) {
+    vec_push(bb->succ, ir->bb1);
+    vec_push(ir->bb1->pred, bb);
+    add_edges(ir->bb1);
+  }
+
+  if (ir->bb2) {
+    vec_push(bb->succ, ir->bb2);
+    vec_push(ir->bb2->pred, bb);
+    add_edges(ir->bb2);
+  }
+}
+
+// Initializes bb->def_regs.
+static void set_def_regs(BB *bb) {
+  if (bb->param)
+    vec_union1(bb->def_regs, bb->param);
+
+  for (int i = 0; i < bb->ir->len; i++) {
+    IR *ir = bb->ir->data[i];
+    if (ir->r0)
+      vec_union1(bb->def_regs, ir->r0);
+  }
+}
+
+// Back-propagate r in the call flow graph.
+static void propagate(BB *bb, Reg *r) {
+  if (!r || vec_contains(bb->def_regs, r))
+    return;
+
+  if (!vec_union1(bb->in_regs, r))
+    return;
+
+  for (int i = 0; i < bb->pred->len; i++) {
+    BB *pred = bb->pred->data[i];
+    if (vec_union1(pred->out_regs, r))
+      propagate(pred, r);
+  }
+}
+
+// Initializes bb->in_regs and bb->out_regs.
+static void visit(BB *bb, IR *ir) {
+  propagate(bb, ir->r1);
+  propagate(bb, ir->r2);
+  propagate(bb, ir->bbarg);
+
+  if (ir->op == IR_CALL)
+    for (int i = 0; i < ir->nargs; i++)
+      propagate(bb, ir->args[i]);
+}
+
+void liveness(Program *prog) {
+  for (int i = 0; i < prog->funcs->len; i++) {
+    Function *fn = prog->funcs->data[i];
+    add_edges(fn->bbs->data[0]);
+
+    for (int i = 0; i < fn->bbs->len; i++) {
+      BB *bb = fn->bbs->data[i];
+      set_def_regs(bb);
+
+      for (int i = 0; i < bb->ir->len; i++) {
+	IR *ir = bb->ir->data[i];
+	visit(bb, ir);
+      }
+    }
+  }
+}
diff --git a/main.c b/main.c
index fa146a9..46fd7be 100644
--- a/main.c
+++ b/main.c
@@ -38,11 +38,13 @@ int main(int argc, char **argv) {
   if (dump_ir1)
     dump_ir(prog->funcs);
 
-  alloc_regs(prog);
+  liveness(prog);
 
   if (dump_ir2)
     dump_ir(prog->funcs);
 
+  alloc_regs(prog);
+
   gen_x86(prog);
   return 0;
 }
diff --git a/regalloc.c b/regalloc.c
index 900a122..25465ac 100644
--- a/regalloc.c
+++ b/regalloc.c
@@ -77,7 +77,13 @@ static Vector *collect_regs(Function *fn) {
         for (int i = 0; i < ir->nargs; i++)
           set_last_use(ir->args[i], ic);
     }
+
+    for (int i = 0; i < bb->out_regs->len; i++) {
+      Reg *r = bb->out_regs->data[i];
+      set_last_use(r, ic);
+    }
   }
+
   return v;
 }
 
diff --git a/util.c b/util.c
index 75eb0d3..a4d30d0 100644
--- a/util.c
+++ b/util.c
@@ -47,6 +47,20 @@ void *vec_last(Vector *v) {
   return v->data[v->len - 1];
 }
 
+bool vec_contains(Vector *v, void *elem) {
+  for (int i = 0; i < v->len; i++)
+    if (v->data[i] == elem)
+      return true;
+  return false;
+}
+
+bool vec_union1(Vector *v, void *elem) {
+  if (vec_contains(v, elem))
+    return false;
+  vec_push(v, elem);
+  return true;
+}
+
 Map *new_map(void) {
   Map *map = malloc(sizeof(Map));
   map->keys = new_vec();