From 77a459ffd3871d04d8aa53fbcc07bdd32e8b3f41 Mon Sep 17 00:00:00 2001
From: Mid <245600-midn@users.noreply.gitlab.com>
Date: Thu, 15 Feb 2024 22:33:06 +0200
Subject: [PATCH] Loop guards

---
 src/ast.c              | 12 +++++----
 src/cg.c               | 48 +++++++++++++++++++-----------------
 src/dumberdowner.c     |  4 +--
 src/parse.c            | 56 ++++++++++++++++++++++++++++++++++++++++++
 src/vartable.c         |  1 +
 src/vartable.h         |  3 +++
 tests/loopregalloc.nct |  7 ++++++
 7 files changed, 101 insertions(+), 30 deletions(-)
 create mode 100644 tests/loopregalloc.nct

diff --git a/src/ast.c b/src/ast.c
index aaa69b7..3c11884 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -29,6 +29,13 @@ int ast_stmt_is_after(const AST *chunk, const AST *s1, const AST *s2) {
 	const AST *s = chunk->chunk.statementFirst;
 	
 	while(s) {
+		if(s->nodeKind == AST_STMT_LOOP) {
+			int i = ast_stmt_is_after(s->stmtLoop.body, s1, s2);
+			if(i != -1) {
+				return i;
+			}
+		}
+		
 		if(s == s1) {
 			return 0;
 		}
@@ -41,11 +48,6 @@ int ast_stmt_is_after(const AST *chunk, const AST *s1, const AST *s2) {
 			if(i != -1) {
 				return i;
 			}
-		} else if(s->nodeKind == AST_STMT_LOOP) {
-			int i = ast_stmt_is_after(s->stmtLoop.body, s1, s2);
-			if(i != -1) {
-				return i;
-			}
 		}
 		
 		s = s->statement.next;
diff --git a/src/cg.c b/src/cg.c
index ddd7ce7..6096220 100644
--- a/src/cg.c
+++ b/src/cg.c
@@ -178,32 +178,34 @@ void cg_chunk(AST *a) {
 				}
 			}
 		} else if(s->nodeKind == AST_STMT_ASSIGN) {
-		
-			if(s->stmtAssign.to->nodeKind == AST_EXPR_BINARY_OP && ast_expression_equal(s->stmtAssign.what, s->stmtAssign.to->exprBinOp.operands[0]) && (s->stmtAssign.to->exprBinOp.operator == BINOP_ADD || s->stmtAssign.to->exprBinOp.operator == BINOP_SUB) && s->stmtAssign.to->exprBinOp.operands[1]->nodeKind == AST_EXPR_PRIMITIVE && s->stmtAssign.to->exprBinOp.operands[1]->exprPrim.val == 1) {
 			
-				// inc or dec
+			if(s->stmtAssign.to) {
+				if(s->stmtAssign.to->nodeKind == AST_EXPR_BINARY_OP && ast_expression_equal(s->stmtAssign.what, s->stmtAssign.to->exprBinOp.operands[0]) && (s->stmtAssign.to->exprBinOp.operator == BINOP_ADD || s->stmtAssign.to->exprBinOp.operator == BINOP_SUB) && s->stmtAssign.to->exprBinOp.operands[1]->nodeKind == AST_EXPR_PRIMITIVE && s->stmtAssign.to->exprBinOp.operands[1]->exprPrim.val == 1) {
 				
-				static const char *instrs[] = {"inc", "dec"};
-				printf("%s %s %s\n", instrs[s->stmtAssign.to->exprBinOp.operator == BINOP_SUB], specexpr(s->stmtAssign.what), xop(s->stmtAssign.what));
-				
-			} else if(s->stmtAssign.what->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->nodeKind == AST_EXPR_BINARY_OP && s->stmtAssign.to->exprBinOp.operator == BINOP_ADD && s->stmtAssign.to->exprBinOp.operands[0]->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[1]->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[0]->exprVar.thing->kind == VARTABLEENTRY_VAR && s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing->kind == VARTABLEENTRY_VAR) {
-				
-				printf("lea %s, [%s + %s]\n",
-					xv(s->stmtAssign.what->exprVar.thing),
-					xv(s->stmtAssign.to->exprBinOp.operands[0]->exprVar.thing),
-					xv(s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing));
-				
-			} else if(s->stmtAssign.what->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->nodeKind == AST_EXPR_BINARY_OP && s->stmtAssign.to->exprBinOp.operator == BINOP_ADD && s->stmtAssign.to->exprBinOp.operands[0]->nodeKind == AST_EXPR_UNARY_OP && s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operator == UNOP_REF && s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operand->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[1]->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operand->exprVar.thing->kind == VARTABLEENTRY_SYMBOL && s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing->kind == VARTABLEENTRY_VAR) {
-				
-				printf("lea %s, [%s + %s]\n",
-					xv(s->stmtAssign.what->exprVar.thing),
-					s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operand->exprVar.thing->data.symbol.name,
-					xv(s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing));
-				
-			} else {
-			
-				printf("mov %s, %s\n", xop(s->stmtAssign.what), xop(s->stmtAssign.to));
+					// inc or dec
+					
+					static const char *instrs[] = {"inc", "dec"};
+					printf("%s %s %s\n", instrs[s->stmtAssign.to->exprBinOp.operator == BINOP_SUB], specexpr(s->stmtAssign.what), xop(s->stmtAssign.what));
+					
+				} else if(s->stmtAssign.what->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->nodeKind == AST_EXPR_BINARY_OP && s->stmtAssign.to->exprBinOp.operator == BINOP_ADD && s->stmtAssign.to->exprBinOp.operands[0]->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[1]->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[0]->exprVar.thing->kind == VARTABLEENTRY_VAR && s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing->kind == VARTABLEENTRY_VAR) {
+					
+					printf("lea %s, [%s + %s]\n",
+						xv(s->stmtAssign.what->exprVar.thing),
+						xv(s->stmtAssign.to->exprBinOp.operands[0]->exprVar.thing),
+						xv(s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing));
+					
+				} else if(s->stmtAssign.what->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->nodeKind == AST_EXPR_BINARY_OP && s->stmtAssign.to->exprBinOp.operator == BINOP_ADD && s->stmtAssign.to->exprBinOp.operands[0]->nodeKind == AST_EXPR_UNARY_OP && s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operator == UNOP_REF && s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operand->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[1]->nodeKind == AST_EXPR_VAR && s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operand->exprVar.thing->kind == VARTABLEENTRY_SYMBOL && s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing->kind == VARTABLEENTRY_VAR) {
+					
+					printf("lea %s, [%s + %s]\n",
+						xv(s->stmtAssign.what->exprVar.thing),
+						s->stmtAssign.to->exprBinOp.operands[0]->exprUnOp.operand->exprVar.thing->data.symbol.name,
+						xv(s->stmtAssign.to->exprBinOp.operands[1]->exprVar.thing));
+					
+				} else {
 				
+					printf("mov %s, %s\n", xop(s->stmtAssign.what), xop(s->stmtAssign.to));
+					
+				}
 			}
 			
 		} else if(s->nodeKind == AST_STMT_LOOP) {
diff --git a/src/dumberdowner.c b/src/dumberdowner.c
index 95cd95b..deee75a 100644
--- a/src/dumberdowner.c
+++ b/src/dumberdowner.c
@@ -158,14 +158,14 @@ static AST *varify(AST *tlc, AST *chunk, AST *stmtPrev, AST *stmt, AST *e) {
 	// Assemble ud-chain
 	
 	UseDef *ud[2];
-	
 	ud[0] = malloc(sizeof(UseDef));
+	ud[1] = malloc(sizeof(UseDef));
+	
 	ud[0]->stmt = (AST*) assign;
 	ud[0]->use = assign->what;
 	ud[0]->def = (AST*) assign;
 	ud[0]->next = ud[1];
 	
-	ud[1] = malloc(sizeof(UseDef));
 	ud[1]->stmt = stmt;
 	ud[1]->use = (AST*) ev[1];
 	ud[1]->def = (AST*) assign;
diff --git a/src/parse.c b/src/parse.c
index 2b884be..fa1f175 100644
--- a/src/parse.c
+++ b/src/parse.c
@@ -44,6 +44,11 @@ typedef struct {
 	// Used by pushstmt to assemble use-def chain
 	size_t udsToAddCount;
 	UseDefToAdd *udsToAdd;
+	
+	// Used to place guard variable uses after loops to stop reg allocation from fucking up
+	VarTable *loopScope;
+	size_t guardedVarCount;
+	ASTExprVar **guardedVars;
 } Parser;
 
 static Token get(Parser *P) {
@@ -167,6 +172,36 @@ static AST *exprvar(Parser *P, VarTableEntry *v) {
 		newusedef(P, v, a);
 	}
 	
+	if(P->loopScope) {
+		// XXX: O(n)!!!!!!!!!
+		
+		int inloop = 0;
+		for(VarTable *vt = v->owner; vt; vt = vt->parent) {
+			if(vt->parent == P->loopScope) {
+				inloop = 1;
+				break;
+			}
+		}
+		
+		if(!inloop) {
+			int alreadyAdded = 0;
+			for(size_t i = 0; i < P->guardedVarCount; i++) {
+				if(P->guardedVars[i]->thing == v) {
+					alreadyAdded = 1;
+					break;
+				}
+			}
+			
+			if(!alreadyAdded) {
+				ASTExprVar *ev = malloc(sizeof(*ev));
+				memcpy(ev, a, sizeof(*ev));
+				
+				P->guardedVars = realloc(P->guardedVars, sizeof(*P->guardedVars) * (P->guardedVarCount + 1));
+				P->guardedVars[P->guardedVarCount++] = ev;
+			}
+		}
+	}
+	
 	return a;
 }
 
@@ -614,6 +649,7 @@ static AST *parse_declaration(Parser *P) {
 		
 		ASTStmtAssign *assign = malloc(sizeof(*assign));
 		assign->nodeKind = AST_STMT_ASSIGN;
+		assign->next = NULL;
 		
 		entry->data.var.reachingDefs = reachingdefs_push(NULL);
 		reachingdefs_set(entry->data.var.reachingDefs, (AST*) assign);
@@ -765,10 +801,29 @@ void nct_parse_statement(Parser *P) {
 		ret->nodeKind = AST_STMT_LOOP;
 		ret->next = NULL;
 		
+		int isFirstLoop = P->loopScope == NULL;
+		
+		if(isFirstLoop) {
+			P->loopScope = P->scope;
+		}
+		
 		expect(P, TOKEN_SQUIGGLY_L);
 		ret->body = (AST*) nct_parse_chunk(P, 0, 1);
 		expect(P, TOKEN_SQUIGGLY_R);
 		
+		if(isFirstLoop) {
+			P->loopScope = NULL;
+			
+			for(size_t i = 0; i < P->guardedVarCount; i++) {
+				ASTExprVar *ev = P->guardedVars[i];
+				
+				newusedef(P, ev->thing, ev);
+			}
+			P->guardedVarCount = 0;
+			free(P->guardedVars);
+			P->guardedVars = NULL;
+		}
+		
 		loop_second_pass(ret->body);
 		
 		pushstat(P, ret);
@@ -958,6 +1013,7 @@ ASTChunk *nct_parse_chunk(Parser *P, int isTopLevel, int varPrioritize) {
 	for(size_t i = 0; i < P->scope->count; i++) {
 		if(P->scope->data[i]->kind == VARTABLEENTRY_VAR) {
 			P->topLevel->vars[P->topLevel->varCount++] = P->scope->data[i];
+			P->scope->data[i]->owner = P->topLevel;
 		}
 	}
 	
diff --git a/src/vartable.c b/src/vartable.c
index fab5e34..b6be829 100644
--- a/src/vartable.c
+++ b/src/vartable.c
@@ -65,6 +65,7 @@ VarTableEntry *vartable_set(VarTable *this, const char *name, VarTableEntry *e)
 	this->data[this->count] = e;
 	this->count++;
 	if(e->kind == VARTABLEENTRY_VAR) e->data.var.name = name;
+	e->owner = this;
 	return e;
 }
 
diff --git a/src/vartable.h b/src/vartable.h
index 30dd4a0..a97f7fb 100644
--- a/src/vartable.h
+++ b/src/vartable.h
@@ -27,9 +27,12 @@ struct ReachingDefs *reachingdefs_push(struct ReachingDefs*);
 struct ReachingDefs *reachingdefs_coalesce(struct ReachingDefs*);
 void reachingdefs_set(struct ReachingDefs*, union AST*);
 
+struct VarTable;
 typedef struct VarTableEntry {
 	Type *type;
 	
+	struct VarTable *owner;
+	
 	VarTableEntryKind kind;
 	struct {
 		union {
diff --git a/tests/loopregalloc.nct b/tests/loopregalloc.nct
new file mode 100644
index 0000000..332c793
--- /dev/null
+++ b/tests/loopregalloc.nct
@@ -0,0 +1,7 @@
+u8 a;
+
+loop {
+	a = 3;
+	
+	u8 b = 1;
+}
\ No newline at end of file