add callsite register saving

5 files changed, 311 insertions, 37 deletions

diff --git a/src/abi.c b/src/abi.c
index 80f34e4..fede365 100644
--- a/src/abi.c
+++ b/src/abi.c
@@ -30,7 +30,9 @@ static struct insn rewrite_retval(struct insn n)
 	assert(n.in[1].class == IMM);
 	int64_t nth_retval = n.in[1].v;
 	assert(nth_retval >= 0 && nth_retval < 25);
-	return insn_create(MOVE, I27, n.out, nth_ar(nth_retval), noclass());
+	n = insn_create(MOVE, I27, n.out, nth_ar(nth_retval), noclass());
+	set_insn_flags(&n, CALL_TEARDOWN);
+	return n;
 }
 
 static struct insn rewrite_arg(struct insn n)
@@ -40,12 +42,18 @@ static struct insn rewrite_arg(struct insn n)
 	int64_t nth_arg = n.in[1].v;
 	assert(nth_arg >= 0 && nth_arg < 25);
 
-	if (n.in[0].class == TMP || n.in[0].class == REG)
-		return insn_create(MOVE, I27, nth_ar(nth_arg), n.in[0],
+	if (n.in[0].class == TMP || n.in[0].class == REG) {
+		n = insn_create(MOVE, I27, nth_ar(nth_arg), n.in[0],
 		                   noclass());
-	else if (n.in[0].class == IMM || n.in[0].class == REF)
-		return insn_create(COPY, I27, nth_ar(nth_arg), n.in[0],
+		set_insn_flags(&n, CALL_SETUP);
+		return n;
+	}
+	else if (n.in[0].class == IMM || n.in[0].class == REF) {
+		n = insn_create(COPY, I27, nth_ar(nth_arg), n.in[0],
 		                   noclass());
+		set_insn_flags(&n, CALL_SETUP);
+		return n;
+	}
 
 	assert("illegal arg type");
 	abort();
diff --git a/src/asm.c b/src/asm.c
index fd4d68a..7eac389 100644
--- a/src/asm.c
+++ b/src/asm.c
@@ -95,26 +95,50 @@ static const char *rname(struct val v)
 	abort();
 }
 
+static size_t frame_size(struct fn *f)
+{
+	return 3 * f->max_callee_save + 3 * f->max_call_save + 3 * 2;
+}
+
+static size_t callee_save_fploc(size_t i)
+{
+	return 3 * (i + 1) + 3 * 2;
+}
+
+static size_t call_save_fploc(struct fn *f, size_t i)
+{
+	return 3 * (f->max_callee_save + 1) + 3 * i + 3 * 2;
+}
+
+static size_t ra_fploc()
+{
+	return 6;
+}
+
+static size_t fp_fploc()
+{
+	return 3;
+}
+
 static void save_state(struct fn *f, FILE *o)
 {
 	/* stack frame:
 	 * last frame
 	 * ra
-	 * s0
-	 * ...
-	 * local variables
+	 * callee save
+	 * local variables/spills
 	 * <- sp
 	 */
-	fprintf(o, "st w fp, -3(sp)\n");
+	fprintf(o, "st w fp, -%zi(sp)\n", fp_fploc());
+	fprintf(o, "mv fp, sp\n");
 	if (f->has_calls)
-		fprintf(o, "st w ra, -6(fp)\n");
+		fprintf(o, "st w ra, -%zi(fp)\n", ra_fploc());
 
-	fprintf(o, "mv fp, sp\n");
-	fprintf(o, "addi sp, sp, -%zi\n", f->max_callee_save * 3 + 6);
+	fprintf(o, "addi sp, sp, -%zi\n", frame_size(f));
 
 	for (size_t i = 0; i < f->max_callee_save; ++i) {
 		fprintf(o, "st w s%zi, -%zi(fp)\n",
-		        i, 3 * i + 9);
+		        i, callee_save_fploc(i));
 	}
 }
 
@@ -122,13 +146,13 @@ static void restore_state(struct fn *f, FILE *o)
 {
 	for (size_t i = 0; i < f->max_callee_save; ++i) {
 		fprintf(o, "ld w s%zi, -%zi(fp)\n",
-		        i, 3 * i + 9);
+		        i, callee_save_fploc(i));
 	}
 
 	if (f->has_calls)
-		fprintf(o, "ld w ra, -6(fp)\n");
+		fprintf(o, "ld w ra, -%zi(fp)\n", ra_fploc());
 
-	fprintf(o, "ld w fp, -3(fp)\n");
+	fprintf(o, "ld w fp, -%zi(fp)\n", fp_fploc());
 	fprintf(o, "mv sp, fp\n");
 }
 
@@ -202,7 +226,25 @@ static void output_call(struct insn n, FILE *o)
 	abort();
 }
 
-static void output_insn(struct insn n, FILE *o)
+static void output_save(struct insn n, FILE *o, struct fn *f)
+{
+	struct val r = n.in[0];
+	assert(r.class == REG);
+	int64_t i = n.in[1].v;
+	fprintf(o, "st w %s, -%zi(fp)\n",
+			rname(r), call_save_fploc(f, i));
+}
+
+static void output_restore(struct insn n, FILE *o, struct fn *f)
+{
+	struct val r = n.in[0];
+	assert(r.class == REG);
+	int64_t i = n.in[1].v;
+	fprintf(o, "ld w %s, -%zi(fp)\n",
+			rname(r), call_save_fploc(f, i));
+}
+
+static void output_insn(struct insn n, FILE *o, struct fn *f)
 {
 	/* one insn directly matches one or more assembly instructions,
 	 * we may be missing out on certain optimizations by not using some kind
@@ -213,6 +255,8 @@ static void output_insn(struct insn n, FILE *o)
 	case SUB: output_sub(n, o); break;
 	case COPY: output_copy(n, o); break;
 	case CALL: output_call(n, o); break;
+	case SAVE: output_save(n, o, f); break;
+	case RESTORE: output_restore(n, o, f); break;
 	default: fprintf(stderr, "unimplemented insn: %s\n", op_str(n.type));
 		abort();
 	}
@@ -285,7 +329,7 @@ void output(struct fn *f, FILE *o)
 
 		foreach_insn(i, b->insns) {
 			struct insn n = insn_at(b->insns, i);
-			output_insn(n, o);
+			output_insn(n, o, f);
 		}
 
 		if (b->s2)
diff --git a/src/nodes.c b/src/nodes.c
index 05f96fe..d5a22da 100644
--- a/src/nodes.c
+++ b/src/nodes.c
@@ -144,6 +144,8 @@ void destroy_block(struct blk *b)
 {
 	vec_destroy(&b->insns);
 	vec_destroy(&b->params);
+	vec_destroy(&b->args1);
+	vec_destroy(&b->args2);
 	free(b);
 }
 
@@ -188,6 +190,9 @@ void dump_insn(struct insn i)
 		dump_val(i.in[1]);
 	}
 
+	if (i.flags)
+		printf("*");
+
 	printf("\n");
 }
 
diff --git a/src/regalloc.c b/src/regalloc.c
index 68e68a2..ffefeb5 100644
--- a/src/regalloc.c
+++ b/src/regalloc.c
@@ -9,7 +9,11 @@
  * Idea is that callee-save cost is paid once at the start of the function call,
  * after that they're free. Argument registers are generally important to keep
  * free to make sure function calls in loops etc. don't have to shuffle
- * registers around as much */
+ * registers around as much. The preferred order can probably be bikeshedded to death,
+ * eventually would be kind of cool to add some heuristic for which category
+ * might be best suited for a specific location. 
+ * One possible one would be that callee-save should be preferred if the
+ * lifetime overlaps a function call? */
 static const int64_t tr_map[] = {
 	RT0, RT1, RT2, RT3, RT4, RT5, RT6, RT7, RT8, RT9,
 	RT10, RT11, RT12, RT13, RT14, RT15, RT16, RT17, RT18, RT19,
@@ -69,7 +73,7 @@ static struct val get_hint(struct vec *hints, struct val from)
 	if ((int64_t)vec_len(hints) <= from.r)
 		return noclass();
 
-	return rewrite_tmp(hints, from);
+	return reg_at(*hints, from.r);
 }
 
 struct lifetime {
@@ -94,7 +98,7 @@ static void add_def(struct vec *lifetimes, struct val v, size_t i)
 	/* I guess a def isn't techincally a 'use' but good enough, we assume
 	 * that unused variables have already been removed by some earlier stage
 	 * (not currently true but that's the intention) */
-	lifetime_at(*lifetimes, v.r) = (struct lifetime){v, i, 0, 1};
+	lifetime_at(*lifetimes, v.r) = (struct lifetime){v, i, i, 1};
 }
 
 static void add_use(struct vec *lifetimes, struct val v, size_t i)
@@ -108,7 +112,7 @@ static void add_use(struct vec *lifetimes, struct val v, size_t i)
 }
 
 static void collect_lifetimes(struct blk *b, struct vec *hints,
-                              struct vec *lifetimes)
+                              struct vec *lifetimes, struct vec *rmap, struct vec *calls)
 {
 	foreach_blk_param(pi, b->params) {
 		struct val v = blk_param_at(b->params, pi);
@@ -130,13 +134,21 @@ static void collect_lifetimes(struct blk *b, struct vec *hints,
 		/* collect some early hints */
 		if (i.type == MOVE) {
 			/* input arguments, retvals */
-			if (i.out.class == TMP && i.in[0].class == REG)
+			if (i.out.class == TMP && i.in[0].class == REG) {
 				add_hint(hints, i.out, i.in[0]);
+				/* retvals must be treated as reserving that
+				 * virtual register */
+				add_rewrite_rule(rmap, i.out, i.in[0]);
+			}
 
 			if (i.out.class == REG && i.in[0].class == TMP)
 				add_hint(hints, i.in[0], i.out);
 		}
 
+		/* collect calls, to be used later */
+		if (i.type == CALL)
+			vec_append(calls, &pos);
+
 		pos++;
 	}
 
@@ -157,24 +169,29 @@ static void collect_lifetimes(struct blk *b, struct vec *hints,
 	}
 }
 
-static void build_active(struct vec *active, struct vec *lifetimes, size_t i)
+static  void build_active_between(struct vec *active, struct vec *lifetimes, size_t start, size_t end)
 {
-	struct lifetime ref = lifetime_at(*lifetimes, i);
 	foreach_lifetime(li, *lifetimes) {
 		struct lifetime l = lifetime_at(*lifetimes, li);
 		if (l.used == 0)
 			continue;
 
-		if (l.end < ref.start)
+		if (l.end < start)
 			continue;
 
-		if (l.start > ref.end)
+		if (l.start > end)
 			continue;
 
 		vec_append(active, &l);
 	}
 }
 
+static void build_active(struct vec *active, struct vec *lifetimes, size_t i)
+{
+	struct lifetime ref = lifetime_at(*lifetimes, i);
+	build_active_between(active, lifetimes, ref.start, ref.end);
+}
+
 static void build_reserved(struct vec *reserved, struct vec *active,
                            struct vec *rmap)
 {
@@ -212,14 +229,51 @@ static struct val find_free_reg(struct vec *reserved)
 	}
 
 	/* handle spill case later */
+	/* right now I'm thinking that spills should behave like SAVE/RESTORE,
+	 * i.e. they get an index and the spilled number of registers is counted
+	 * somewhere and added to the frame size like call registers. */
 	assert(0 &&
 	       "ran out of registers, time to implement proper spill handling");
 	abort();
 }
 
-static void build_rmap(struct vec *hints, struct vec *lifetimes,
+static size_t highest_sreg(struct val f)
+{
+	assert(f.class == REG);
+	switch (f.r) {
+		case RS0: return 1;
+		case RS1: return 2;
+		case RS2: return 3;
+		case RS3: return 4;
+		case RS4: return 5;
+		case RS5: return 6;
+		case RS6: return 7;
+		case RS7: return 8;
+		case RS8: return 9;
+		case RS9: return 10;
+		case RS10: return 11;
+		case RS11: return 12;
+		case RS12: return 13;
+		case RS13: return 14;
+		case RS14: return 15;
+		case RS15: return 16;
+		case RS16: return 17;
+		case RS17: return 18;
+		case RS18: return 19;
+		case RS19: return 20;
+		case RS20: return 21;
+		case RS21: return 22;
+		case RS22: return 23;
+		case RS23: return 24;
+	}
+
+	return 0;
+}
+
+static size_t build_rmap(struct vec *hints, struct vec *lifetimes,
                        struct vec *rmap)
 {
+	size_t max_callee_save = 0;
 	struct vec active = vec_create(sizeof(struct lifetime));
 	struct vec reserved = vec_create(sizeof(struct val));
 
@@ -228,7 +282,11 @@ static void build_rmap(struct vec *hints, struct vec *lifetimes,
 		if (l.used == 0)
 			continue;
 
+		if (has_rewrite_rule(rmap, l.v))
+			continue;
+
 		vec_reset(&active);
+		vec_reset(&reserved);
 		build_active(&active, lifetimes, li);
 		build_reserved(&reserved, &active, rmap);
 
@@ -244,12 +302,20 @@ static void build_rmap(struct vec *hints, struct vec *lifetimes,
 		 * least used register, but that's a bit more complicated than
 		 * just this linear scan */
 		struct val f = find_free_reg(&reserved);
+		if (highest_sreg(f) > max_callee_save)
+			max_callee_save = highest_sreg(f);
+
 		add_rewrite_rule(rmap, l.v, f);
 	}
+
+	vec_destroy(&active);
+	vec_destroy(&reserved);
+
+	return max_callee_save;
 }
 
 /* has_calls is kind of in an iffy place but I guess it's fine */
-static void do_rewrites(struct blk *b, struct vec *rmap, struct fn *f)
+static void do_rewrites(struct blk *b, struct vec *rmap)
 {
 	foreach_blk_param(pi, b->params) {
 		struct val p = blk_param_at(b->params, pi);
@@ -267,10 +333,6 @@ static void do_rewrites(struct blk *b, struct vec *rmap, struct fn *f)
 		if (i.out.class == TMP)
 			i.out = rewrite_tmp(rmap, i.out);
 
-		/** @todo save caller-save at callsites here */
-		if (i.type == CALL)
-			f->has_calls = true;
-
 		/* write back changes, christ I forget this a lot */
 		insn_at(b->insns, ii) = i;
 	}
@@ -292,25 +354,160 @@ static void do_rewrites(struct blk *b, struct vec *rmap, struct fn *f)
 	}
 }
 
+static bool callee_save(struct val r)
+{
+	assert(r.class == REG);
+	switch (r.r) {
+		case RS0: return true;
+		case RS1: return true;
+		case RS2: return true;
+		case RS3: return true;
+		case RS4: return true;
+		case RS5: return true;
+		case RS6: return true;
+		case RS7: return true;
+		case RS8: return true;
+		case RS9: return true;
+		case RS10: return true;
+		case RS11: return true;
+		case RS12: return true;
+		case RS13: return true;
+		case RS14: return true;
+		case RS15: return true;
+		case RS16: return true;
+		case RS17: return true;
+		case RS18: return true;
+		case RS19: return true;
+		case RS20: return true;
+		case RS21: return true;
+		case RS22: return true;
+		case RS23: return true;
+	}
+
+	return false;
+}
+
+static void insn_insert_before_call(struct blk *b, struct insn save, size_t pos)
+{
+	assert((insn_at(b->insns, pos)).type == CALL);
+	struct insn setup;
+	do {
+		/* the first instruction in a block can be part of a call setup */
+		if (pos == 0)
+			break;
+
+		pos--;
+		setup = insn_at(b->insns, pos);
+	} while (has_insn_flag(setup, CALL_SETUP));
+
+	/* the do-while loop technically went one too far, fix */
+	insn_insert(b, save, pos == 0 ? 0 : pos + 1);
+}
+
+static void insn_insert_after_call(struct blk *b, struct insn restore, size_t pos)
+{
+	assert((insn_at(b->insns, pos)).type == CALL);
+	size_t max = vec_len(&b->insns) - 1;
+
+	struct insn setup;
+	do {
+		if (pos == max)
+			break;
+
+		pos++;
+		setup = insn_at(b->insns, pos);
+	} while (has_insn_flag(setup, CALL_TEARDOWN));
+
+	insn_insert(b, restore, pos);
+}
+
+static size_t do_call_saves(struct blk *b, struct vec *lifetimes, struct vec *rmap, struct vec *calls)
+{
+	size_t offset = 0;
+	size_t max_counter = 0;
+	struct vec active = vec_create(sizeof(struct lifetime));
+	struct vec reserved = vec_create(sizeof(struct val));
+
+	foreach_vec(ci, *calls) {
+		size_t call_pos = vect_at(size_t, *calls, ci);
+
+		vec_reset(&active);
+		vec_reset(&reserved);
+		build_active_between(&active, lifetimes, call_pos, call_pos);
+		build_reserved(&reserved, &active, rmap);
+
+		/* what follows is a slight bit of index counting, a bit
+		 * difficult to follow but not too bad */
+
+		/* add a save right before call, offset is how many
+		 * save/restores we've already added, very important */
+		size_t counter = 0;
+		foreach_val(ri, reserved) {
+			struct val r = val_at(reserved, ri);
+			assert(r.class == REG);
+			if (callee_save(r))
+				continue;
+
+			struct insn save = insn_create(SAVE, NOTYPE,
+							noclass(), r, imm_val(counter, I27));
+			insn_insert_before_call(b, save, call_pos + offset - 1);
+			offset++;
+			counter++;
+		}
+
+		call_pos += offset;
+
+		counter = 0;
+		foreach_val(ri, reserved) {
+			struct val r = val_at(reserved, ri);
+			if (callee_save(r))
+				continue;
+
+			/* add a restore right after call ~area~ */
+			/* hmm, restore should maybe put r as its output to be
+			 * more consistent... */
+			struct insn restore = insn_create(RESTORE, NOTYPE,
+							noclass(), r, imm_val(counter, I27));
+			insn_insert_after_call(b, restore, call_pos - 1);
+			counter++;
+			offset++;
+		}
+
+		if (counter > max_counter)
+			max_counter = counter;
+	}
+
+	vec_destroy(&active);
+	vec_destroy(&reserved);
+
+	return max_counter;
+}
+
 /* assumes ssa form */
 void regalloc(struct fn *f)
 {
 	struct vec hints = vec_create(sizeof(struct val));
 	struct vec lifetimes = vec_create(sizeof(struct lifetime));
 	struct vec rmap = vec_create(sizeof(struct val));
+	struct vec calls = vec_create(sizeof(size_t));
 
 	/* here it would probably make sense to iterate through lifetimes in
 	 * order of priority, like loop nesting/count etc */
 	foreach_blk(bi, f->blks) {
 		struct blk *b = blk_at(f->blks, bi);
 		vec_reset(&lifetimes);
+		vec_reset(&calls);
 		/** @todo collect hints from args/params */
-		collect_lifetimes(b, &hints, &lifetimes);
-		build_rmap(&hints, &lifetimes, &rmap);
-		do_rewrites(b, &rmap, f);
+		collect_lifetimes(b, &hints, &lifetimes, &rmap, &calls);
+		f->max_callee_save = build_rmap(&hints, &lifetimes, &rmap);
+		do_rewrites(b, &rmap);
+		f->max_call_save = do_call_saves(b, &lifetimes, &rmap, &calls);
 		/** @todo forward_hints(b, &hints, &rmap) */
+		if (vec_len(&calls) != 0)
+			f->has_calls = true;
 	}
 
+	vec_destroy(&calls);
 	vec_destroy(&hints);
 	vec_destroy(&lifetimes);
 	vec_destroy(&rmap);
diff --git a/src/vec.c b/src/vec.c
index e487162..7fb314e 100644
--- a/src/vec.c
+++ b/src/vec.c
@@ -49,6 +49,26 @@ void vec_append(struct vec *v, void *n)
 	memcpy(p, n, v->ns);
 }
 
-void vec_destroy(struct vec *v) {
+void vec_destroy(struct vec *v)
+{
 	free(v->buf);
 }
+
+void vec_insert(struct vec *v, void *n, size_t i)
+{
+	assert(i <= vec_len(v));
+	if (i == vec_len(v))
+		vec_append(v, n);
+
+	v->n++;
+	if (v->n >= v->s) {
+		v->s *= 2;
+		v->buf = realloc(v->buf, v->s * v->ns);
+	}
+
+	void *p = vec_at(v, i);
+	size_t elems = v->n - i - 1;
+	size_t bytes = elems * v->ns;
+	memmove(p + v->ns, p, bytes);
+	memcpy(p, n, v->ns);
+}