1 files changed, 70 insertions, 0 deletions

diff --git a/src/compile/compile.c b/src/compile/compile.c
index f552ae3..8e0e250 100644
--- a/src/compile/compile.c
+++ b/src/compile/compile.c
@@ -2856,6 +2856,18 @@ static int spill_cost_sort(struct alive_slot *a, struct alive_slot *b)
 	return a->cost < b->cost;
 }
 
+/* sort barriers according to starting address, smallest to largest */
+static int barrier_sort(struct barrier_tuple *a, struct barrier_tuple *b)
+{
+	if (a->start < b->start)
+		return -1;
+
+	if (a->start > b->start)
+		return 1;
+
+	return a->end - b->end;
+}
+
 /* slightly more parameters than I would like but I guess it's fine */
 static void calculate_alive(struct alive *alive, size_t idx,
 		size_t prio, size_t start, size_t end, size_t *rno,
@@ -2924,6 +2936,40 @@ static void linear_gpr_alloc(struct ejit_func *f)
 	}
 }
 
+static void extend_gpr_lifetime(struct gpr_stat *gpr, struct barriers *barriers, size_t bi)
+{
+	if (bi >= barriers_len(barriers))
+		return;
+
+	struct barrier_tuple *barrier = barriers_at(barriers, bi);
+	if (gpr->end < barrier->start)
+		return;
+
+	/* we cross a barrier */
+	gpr->end = gpr->end > barrier->end ? gpr->end : barrier->end;
+
+	/* check if we cross the next barrier as well,
+	 * hopefully gets optimized into a tail call */
+	return extend_gpr_lifetime(gpr, barriers, bi + 1);
+}
+
+static void extend_fpr_lifetime(struct fpr_stat *fpr, struct barriers *barriers, size_t bi)
+{
+	if (bi >= barriers_len(barriers))
+		return;
+
+	struct barrier_tuple *barrier = barriers_at(barriers, bi);
+	if (fpr->end < barrier->start)
+		return;
+
+	/* we cross a barrier */
+	fpr->end = fpr->end > barrier->end ? fpr->end : barrier->end;
+
+	/* check if we cross the next barrier as well,
+	 * hopefully gets optimized into a tail call */
+	return extend_fpr_lifetime(fpr, barriers, bi + 1);
+}
+
 /* there's a fair bit of repetition between this and the gpr case, hmm */
 static void assign_gprs(struct ejit_func *f)
 {
@@ -2937,8 +2983,20 @@ static void assign_gprs(struct ejit_func *f)
 	struct alive_slot a = {.r = -1, .cost = 0, .idx = 0};
 	alive_append(&alive, a);
 
+	/* barrier index, keeps track of earliest possible barrier we can be
+	 * dealing with. Since register start addresses grow upward, we can
+	 * fairly easily keep track of which barrier a register cannot cross */
+	size_t bi = 0;
 	for (size_t gi = 0; gi < gpr_stats_len(&f->gpr); ++gi) {
 		struct gpr_stat *gpr = gpr_stats_at(&f->gpr, gi);
+
+		extend_gpr_lifetime(gpr, &f->barriers, bi);
+		if (bi < barriers_len(&f->barriers)) {
+			struct barrier_tuple *barrier = barriers_at(&f->barriers, bi);
+			if (gpr->start >= barrier->start)
+				bi++;
+		}
+
 		calculate_alive(&alive, gi,
 				gpr->prio, gpr->start, gpr->end, &gpr->rno,
 				&f->gpr, gpr_dead);
@@ -2989,8 +3047,17 @@ static void assign_fprs(struct ejit_func *f)
 	struct alive_slot a = {.r = -1, .cost = 0, .idx = 0};
 	alive_append(&alive, a);
 
+	size_t bi = 0;
 	for (size_t fi = 0; fi < fpr_stats_len(&f->fpr); ++fi) {
 		struct fpr_stat *fpr = fpr_stats_at(&f->fpr, fi);
+
+		extend_fpr_lifetime(fpr, &f->barriers, bi);
+		if (bi < barriers_len(&f->barriers)) {
+			struct barrier_tuple *barrier = barriers_at(&f->barriers, bi);
+			if (fpr->start >= barrier->start)
+				bi++;
+		}
+
 		calculate_alive(&alive, fi,
 				fpr->prio, fpr->start, fpr->end, &fpr->fno,
 				&f->fpr, fpr_dead);
@@ -3035,6 +3102,9 @@ bool ejit_compile(struct ejit_func *f, bool use_64, bool im_scawed)
 	if (!init_jit())
 		return false;
 
+	/* sort barriers so they can be used to extend register life times in
+	 * loops */
+	barriers_sort(&f->barriers, barrier_sort);
 	assign_gprs(f);
 	assign_fprs(f);