From 5b3991e81ce5af888d7915f7a346736d317a34a9 Mon Sep 17 00:00:00 2001
From: Mid <>
Date: Thu, 2 Jan 2025 18:04:57 +0200
Subject: [PATCH] Handle certain precoloring collisions

---
 src/ast.c |  1 +
 src/cg.c  | 66 ++++++++++++++++++++++++++++++++++++++-----------------
 2 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/src/ast.c b/src/ast.c
index 64db50c..089b679 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -100,6 +100,7 @@ int ast_expression_equal(AST *a, AST *b) {
 	return 0;
 }
 
+// This function may return three values: YES (1), NO (0) or UNKNOWN (-1).
 // ... Ew
 int ast_stmt_is_after(const AST *chunk, const AST *s1, const AST *s2) {
 	const AST *s = chunk->chunk.statementFirst;
diff --git a/src/cg.c b/src/cg.c
index 0d2c9c5..d897cea 100644
--- a/src/cg.c
+++ b/src/cg.c
@@ -165,6 +165,13 @@ static const char *xop(AST *tlc, AST *e) {
 	return xop_sz(tlc, e, type_size(e->expression.type));
 }
 
+static int ud_empty(VarTableEntry *a) {
+	assert(a->kind == VARTABLEENTRY_VAR);
+	assert(!a->data.var.usedefFirst == !a->data.var.usedefLast);
+	
+	return !a->data.var.usedefFirst;
+}
+
 void cg_chunk(CGState *cg, AST *a) {
 	AST *s = a->chunk.statementFirst;
 	
@@ -383,15 +390,9 @@ void cg_chunk(CGState *cg, AST *a) {
 	}
 }
 
-struct Spill2StackState {
-	AST *targetTLC;
-	VarTableEntry *target;
-};
-
 struct Spill2VarState {
 	VarTableEntry *target;
 };
-
 static void spill2var_visitor(AST **aptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
 	static size_t vidx = 0;
 	
@@ -476,6 +477,10 @@ static void spill2var_visitor(AST **aptr, AST *stmt, AST *stmtPrev, AST *chunk,
 	}
 }
 
+struct Spill2StackState {
+	AST *targetTLC;
+	VarTableEntry *target;
+};
 static void spill2stack_visitor(AST **aptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
 	struct Spill2StackState *this = ud;
 	
@@ -521,14 +526,17 @@ static void spill2stack_visitor(AST **aptr, AST *stmt, AST *stmtPrev, AST *chunk
 	}
 }
 
-static int ud_empty(VarTableEntry *a) {
-	assert(a->kind == VARTABLEENTRY_VAR);
-	assert(!a->data.var.usedefFirst == !a->data.var.usedefLast);
-	
-	return !a->data.var.usedefFirst;
-}
-
+struct PrecolorState {
+	ASTExprVar *mustBeSpilled;
+};
 static void precolor_visitor(AST **nptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
+	struct PrecolorState *this = ud;
+	
+	if(this->mustBeSpilled) {
+		/* Since something must be spilled first, we can't do anything else. Quit ASAP. */
+		return;
+	}
+	
 	AST *n = *nptr;
 	
 	if(n == tlc) {
@@ -544,9 +552,15 @@ static void precolor_visitor(AST **nptr, AST *stmt, AST *stmtPrev, AST *chunk, A
 			
 			VarTableEntry *vte = n->stmtReturn.val->exprVar.thing;
 			
-			assert(!vte->data.var.precolored);
+			const int requiredColor = COLOR_EAX;
 			
-			vte->data.var.color = COLOR_EAX;
+			if(vte->data.var.precolored && vte->data.var.color != requiredColor) {
+				// Precoloring collision!
+				this->mustBeSpilled = vte;
+				return;
+			}
+			
+			vte->data.var.color = requiredColor;
 			vte->data.var.precolored = true;
 		}
 	} else if(n->nodeKind == AST_STMT_ASSIGN && n->stmtAssign.to->nodeKind == AST_EXPR_CALL) {
@@ -554,9 +568,15 @@ static void precolor_visitor(AST **nptr, AST *stmt, AST *stmtPrev, AST *chunk, A
 		
 		VarTableEntry *vte = n->stmtAssign.what->exprVar.thing;
 		
-		assert(!vte->data.var.precolored);
+		const int requiredColor = COLOR_EAX;
 		
-		vte->data.var.color = COLOR_EAX;
+		if(vte->data.var.precolored && vte->data.var.color != requiredColor) {
+			// Precoloring collision!
+			this->mustBeSpilled = vte;
+			return;
+		}
+		
+		vte->data.var.color = requiredColor;
 		vte->data.var.precolored = true;
 	}
 }
@@ -626,7 +646,6 @@ struct CalleeSavedState {
 	// To make sure we don't process the same return statement to infinity
 	AST *lastProcessedReturn;
 };
-
 static void callee_saved_visitor(AST **nptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
 	struct CalleeSavedState *this = ud;
 	
@@ -753,11 +772,18 @@ static int comparator(const void *A, const void *B) {
 }
 int cg_go(AST *a) {
 	do {
-		generic_visitor(&a, NULL, NULL, a, a, NULL, precolor_visitor);
+		struct PrecolorState state = {};
+		generic_visitor(&a, NULL, NULL, a, a, &state, precolor_visitor);
 		
 		ast_usedef_reset(a);
 		
-		VarTableEntry *vte = get_precolor_spill_candidate(a);
+		// Spill?
+		VarTableEntry *vte = state.mustBeSpilled; // set if precoloring demands DIFFERENT colors for the same variable
+		if(!vte) {
+			// Find variables with the same color
+			vte = get_precolor_spill_candidate(a);
+		}
+		
 		if(vte) {
 			struct Spill2VarState state;
 			memset(&state, 0, sizeof(state));