loop and break

1 files changed, 44 insertions, 6 deletions

diff --git a/src/main.c b/src/main.c
index de2f81f..966b004 100644
--- a/src/main.c
+++ b/src/main.c
@@ -159,9 +159,11 @@ typedef enum {
 	NOD_BLOCK_STMT,
 	NOD_EXPR_STMT,
 	NOD_RETURN_STMT,
+	NOD_BREAK_STMT,
 	NOD_WRT_STMT,
 	NOD_ASSIGN_STMT,
 	NOD_IFELSE_STMT,
+	NOD_LOOP_STMT,
 	NOD_FUN,
 	NOD_WORD,
 	NOD_FNCALL,
@@ -271,9 +273,11 @@ static Node *expr(Token **rest, Token *tok);
 static Node *block_stmt(Token **rest, Token *tok);
 static Node *expr_stmt(Token **rest, Token *tok);
 static Node *return_stmt(Token **rest, Token *tok);
+static Node *break_stmt(Token **rest, Token *tok);
 static Node *wrt_stmt(Token **rest, Token *tok);
 static Node *assign_stmt(Token **rest, Token *tok);
 static Node *ifelse_stmt(Token **rest, Token *tok);
+static Node *loop_stmt(Token **rest, Token *tok);
 static Node *equality(Token **rest, Token *tok);
 static Node *relational(Token **rest, Token *tok);
 static Node *add(Token **rest, Token *tok);
@@ -495,6 +499,10 @@ stmt(Token **rest, Token *tok)
 		return ifelse_stmt(rest, tok);
 	if (equal(tok, "wrt"))
 		return wrt_stmt(rest, tok);
+	if (equal(tok, "loop"))
+		return loop_stmt(rest, tok);
+	if (equal(tok, "break"))
+		return break_stmt(rest, tok);
 	if (equal(tok, "return"))
 		return return_stmt(rest, tok);
 	if (equal(tok->next, "="))
@@ -538,6 +546,14 @@ return_stmt(Token **rest, Token *tok)
 }
 
 static Node *
+break_stmt(Token **rest, Token *tok)
+{
+	tok = skip(tok, "break");
+	*rest = skip(tok, ";");
+	return new_node(NOD_BREAK_STMT);
+}
+
+static Node *
 wrt_stmt(Token **rest, Token *tok)
 {
 	tok = skip(tok, "wrt");
@@ -574,6 +590,15 @@ ifelse_stmt(Token **rest, Token *tok)
 	return node;
 }
 
+static Node *
+loop_stmt(Token **rest, Token *tok)
+{
+	tok = skip(tok, "loop");
+	Node *const node = new_unary(NOD_LOOP_STMT, stmt(&tok, tok));
+	*rest = tok;
+	return node;
+}
+
 // code generator
 
 static int depth;
@@ -676,7 +701,7 @@ gen_variableget(Node *node)
 }
 
 static void
-gen_stmt(Node *node, Node *fname)
+gen_stmt(Node *node, Node *fname, int break_lbl)
 {
 	static int label = 0;
 	int lbl;
@@ -684,7 +709,7 @@ gen_stmt(Node *node, Node *fname)
 	case NOD_BLOCK_STMT:
 		node = node->lhs;
 		while (node != NULL) {
-			gen_stmt(node, fname);
+			gen_stmt(node, fname, break_lbl);
 			node = node->next;
 		}
 		break;
@@ -698,6 +723,11 @@ gen_stmt(Node *node, Node *fname)
 		printf("\tJMP ,__fnret_%.*s\n", fname->len, fname->loc);
 		depth -= 1;
 		break;
+	case NOD_BREAK_STMT:
+		if (break_lbl == -1)
+			error("break statement outside of loop");
+		printf("\tJMP ,__lbl_%x\n", break_lbl);
+		break;
 	case NOD_WRT_STMT:
 		gen_expr(node->lhs);
 		printf("\tWRT\n");
@@ -721,19 +751,27 @@ gen_stmt(Node *node, Node *fname)
 		break;
 	case NOD_IFELSE_STMT:
 		lbl = label;
-		label = lbl + 1 + (node->rhs->next != NULL);
+		label += 1 + (node->rhs->next != NULL);
 		gen_expr(node->lhs);
 		printf("\tJEZ ,__lbl_%x\n", lbl);
 		depth -= 1;
-		gen_stmt(node->rhs, fname);
+		gen_stmt(node->rhs, fname, break_lbl);
 		if (node->rhs->next != NULL) {
 			printf("\tJMP ,__lbl_%x\n", lbl + 1);
 			printf("@__lbl_%x\n", lbl);
-			gen_stmt(node->rhs->next, fname);
+			gen_stmt(node->rhs->next, fname, break_lbl);
 			lbl += 1;
 		}
 		printf("@__lbl_%x\n", lbl);
 		break;
+	case NOD_LOOP_STMT:
+		lbl = label;
+		label += 2;
+		printf("@__lbl_%x\n", lbl);
+		gen_stmt(node->lhs, fname, lbl + 1);
+		printf("\tJMP ,__lbl_%x\n", lbl);
+		printf("@__lbl_%x\n", lbl + 1);
+		break;
 	default:
 		error("unexpected %d", node->type);
 	}
@@ -779,7 +817,7 @@ gen_function(Node *node)
 	}
 
 	while (node != NULL) {
-		gen_stmt(node, ognode->lhs);
+		gen_stmt(node, ognode->lhs, -1);
 		node = node->next;
 	}