1 files changed, 94 insertions, 3 deletions

diff --git a/src/correct.c b/src/correct.c
index 8c0b249..6e99e3f 100644
--- a/src/correct.c
+++ b/src/correct.c
@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <qbt/opt.h>
 
 static size_t spill_ref(struct blk *b, size_t ii, struct insn i, size_t idx,
@@ -76,8 +77,89 @@ static size_t correct_store(struct blk *b, size_t ii, struct insn i, size_t ri)
 	return ri;
 }
 
-static size_t correct_insn(struct blk *b, size_t ii, struct insn i, size_t ri)
+static void add_rewrite_rule(struct vec *rmap, struct val from, struct val to)
 {
+	assert(from.class == TMP);
+	assert(to.class == TMP);
+
+	while ((int64_t)vec_len(rmap) <= from.r) {
+		struct val no = noclass();
+		vec_append(rmap, &no);
+	}
+
+	val_at(*rmap, from.r) = to;
+}
+
+static bool has_rewrite_rule(struct vec *rmap, struct val t)
+{
+	if (t.r >= (int64_t)vec_len(rmap))
+		return false;
+
+	struct val r = val_at(*rmap, t.r);
+	return r.class != NOCLASS;
+}
+
+static struct val rewrite_tmp(struct vec *rmap, struct val t)
+{
+	assert(t.class == TMP);
+	assert(t.r < (int64_t)vec_len(rmap));
+	struct val r = val_at(*rmap, t.r);
+	assert(r.class != NOCLASS);
+	return r;
+}
+
+static size_t correct_addr(struct blk *b, struct vec *rewrite_addrs, size_t ii, struct insn i, size_t ri)
+{
+	if (i.in[0].class == TMP) {
+		/** @todo this messes with the rest of the corrections, as the
+		 * store with the tmp is rewritten to load t1 first, overwriting
+		 * t0 */
+		/* addr could potentially be defined to move the register into
+		 * the location it's specifying? Not a particularly clean
+		 * solution but I guess it could work? */
+		add_rewrite_rule(rewrite_addrs, i.in[0], i.out);
+		return ri;
+	}
+
+	return ri;
+}
+
+static size_t load_rewrite(struct blk *b, struct vec *rewrite_addrs, size_t ii, struct insn i, size_t ri, size_t idx)
+{
+	struct val addr = rewrite_tmp(rewrite_addrs, i.in[idx]);
+	struct val tmp = tmp_val(ri++);
+	struct insn new = insn_create(LOAD, I27, tmp, addr, noclass(), 0);
+	i.in[idx] = tmp;
+	insn_at(b->insns, ii) = i;
+	insn_insert(b, new, ii);
+	return ri;
+}
+
+static size_t store_rewrite(struct blk *b, struct vec *rewrite_addrs, size_t ii, struct insn i, size_t ri)
+{
+	struct val addr = rewrite_tmp(rewrite_addrs, i.out);
+	struct val tmp = tmp_val(ri++);
+	struct insn new = insn_create(STORE, I27, tmp, addr, noclass(), 0);
+	i.out = tmp;
+	insn_at(b->insns, ii) = i;
+	insn_insert(b, new, ii + 1);
+	return ri;
+}
+
+static size_t correct_insn(struct blk *b, struct vec *rewrite_addrs, size_t ii, struct insn i, size_t ri)
+{
+	/* replace registers referencing rewritten addr */
+	if (i.in[0].class == TMP && has_rewrite_rule(rewrite_addrs, i.in[0]))
+		return load_rewrite(b, rewrite_addrs, ii, i, ri, 0);
+
+	if (i.in[1].class == TMP  && has_rewrite_rule(rewrite_addrs, i.in[1]))
+		return load_rewrite(b, rewrite_addrs, ii, i, ri, 1);
+
+	if (i.out.class == TMP && has_rewrite_rule(rewrite_addrs, i.out)) {
+		/* note no return */
+		store_rewrite(b, rewrite_addrs, ii, i, ri);
+	}
+
 	/* replace references with instructions */
 	if (i.type != CALL && i.in[0].class == REF)
 		return spill_ref(b, ii, i, 0, ri);
@@ -120,22 +202,28 @@ static size_t correct_insn(struct blk *b, size_t ii, struct insn i, size_t ri)
 	case STORE:
 		return correct_store(b, ii, i, ri);
 
+	case ADDR:
+		return correct_addr(b, rewrite_addrs, ii, i, ri);
+
 	default:
 	}
 
 	return ri;
 }
 
-void correct(struct fn *f, size_t ri)
+size_t correct(struct fn *f, size_t ri)
 {
 	/* some simpler corrections to make sure all instructions follow a
 	 * specific pattern. The textual version doesn't have these
 	 * restrictions, but they make our lives easier in the future. */
+
+	struct vec rewrite_addrs = vec_create(sizeof(struct val));
+
 	foreach_blk(bi, f->blks) {
 		struct blk *b = blk_at(f->blks, bi);
 		foreach_insn(ii, b->insns) {
 			struct insn i = insn_at(b->insns, ii);
-			ri = correct_insn(b, ii, i, ri);
+			ri = correct_insn(b, &rewrite_addrs, ii, i, ri);
 		}
 
 		if (b->cmp[0].class == IMM) {
@@ -150,4 +238,7 @@ void correct(struct fn *f, size_t ri)
 			b->cmp[1] = t;
 		}
 	}
+
+	vec_destroy(&rewrite_addrs);
+	return ri;
 }