block statement

1 files changed, 108 insertions, 58 deletions

diff --git a/src/main.c b/src/main.c
index 4bbe051..ccd7643 100644
--- a/src/main.c
+++ b/src/main.c
@@ -156,10 +156,12 @@ typedef enum {
 	NOD_LT,  // <
 	NOD_LE,  // <=
 	NOD_NUM, // integer
+	NOD_BLOCK_STMT,
 	NOD_EXPR_STMT,
 	NOD_RETURN_STMT,
 	NOD_WRT_STMT,
 	NOD_ASSIGN_STMT,
+	NOD_IFELSE_STMT,
 	NOD_FUN,
 	NOD_WORD,
 	NOD_FNCALL,
@@ -266,10 +268,12 @@ static Node *function(Token **rest, Token *tok);
 static Node *global(Token **rest, Token *tok);
 static Node *local(Token **rest, Token *tok);
 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 *wrt_stmt(Token **rest, Token *tok);
 static Node *assign_stmt(Token **rest, Token *tok);
+static Node *ifelse_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);
@@ -297,19 +301,8 @@ function(Token **rest, Token *tok)
 	}
 	node->lhs->next = head.next;
 	tok = skip(tok, ")");
-	tok = skip(tok, "{");
 
-	head.next = NULL;
-	cur = &head;
-
-	while (tok->type != TOK_EOF && equal(tok, "local"))
-		cur = cur->next = local(&tok, tok);
-
-	while (tok->type != TOK_EOF && !equal(tok, "}"))
-		cur = cur->next = stmt(&tok, tok);
-	tok = skip(tok, "}");
-
-	node->rhs = head.next;
+	node->rhs = stmt(&tok, tok);
 
 	*rest = tok;
 	return node;
@@ -496,25 +489,47 @@ primary(Token **rest, Token *tok)
 static Node *
 stmt(Token **rest, Token *tok)
 {
-	if (equal(tok, "return"))
-		return return_stmt(rest, tok);
+	if (equal(tok, "{"))
+		return block_stmt(rest, tok);
+	if (equal(tok, "if"))
+		return ifelse_stmt(rest, tok);
 	if (equal(tok, "wrt"))
 		return wrt_stmt(rest, tok);
+	if (equal(tok, "return"))
+		return return_stmt(rest, tok);
 	if (equal(tok->next, "="))
 		return assign_stmt(rest, tok);
 	return expr_stmt(rest, tok);
 }
 
-static Node
-*expr_stmt(Token **rest, Token *tok)
+static Node *
+block_stmt(Token **rest, Token *tok)
+{
+	tok = skip(tok, "{");
+	Node *const node = new_node(NOD_BLOCK_STMT);
+
+	Node head = {0};
+	Node *cur = &head;
+	while (tok->type != TOK_EOF && equal(tok, "local"))
+		cur = cur->next = local(&tok, tok);
+	while (tok->type != TOK_EOF && !equal(tok, "}"))
+		cur = cur->next = stmt(&tok, tok);
+	node->lhs = head.next;
+
+	*rest = skip(tok, "}");
+	return node;
+}
+
+static Node *
+expr_stmt(Token **rest, Token *tok)
 {
 	Node *const node = new_unary(NOD_EXPR_STMT, expr(&tok, tok));
 	*rest = skip(tok, ";");
 	return node;
 }
 
-static Node
-*return_stmt(Token **rest, Token *tok)
+static Node *
+return_stmt(Token **rest, Token *tok)
 {
 	tok = skip(tok, "return");
 	Node *const node = new_unary(NOD_RETURN_STMT, expr(&tok, tok));
@@ -522,8 +537,8 @@ static Node
 	return node;
 }
 
-static Node
-*wrt_stmt(Token **rest, Token *tok)
+static Node *
+wrt_stmt(Token **rest, Token *tok)
 {
 	tok = skip(tok, "wrt");
 	Node *const node = new_unary(NOD_WRT_STMT, expr(&tok, tok));
@@ -531,8 +546,8 @@ static Node
 	return node;
 }
 
-static Node
-*assign_stmt(Token **rest, Token *tok)
+static Node *
+assign_stmt(Token **rest, Token *tok)
 {
 	expect(tok, TOK_WORD);
 	skip(tok->next, "=");
@@ -542,6 +557,21 @@ static Node
 	return node;
 }
 
+static Node *
+ifelse_stmt(Token **rest, Token *tok)
+{
+	tok = skip(tok, "if");
+	Node *const node = new_unary(NOD_IFELSE_STMT, stmt(&tok, tok));
+
+	if (equal(tok, "else")) {
+		tok = tok->next;
+		node->rhs = stmt(&tok, tok);
+	}
+
+	*rest = skip(tok, ";");
+	return node;
+}
+
 // code generator
 
 static int depth;
@@ -644,12 +674,63 @@ gen_variableget(Node *node)
 }
 
 static void
+gen_stmt(Node *node, Node *fname)
+{
+	switch (node->type) {
+	case NOD_BLOCK_STMT:
+		node = node->lhs;
+		while (node != NULL) {
+			gen_stmt(node, fname);
+			node = node->next;
+		}
+		break;
+	case NOD_EXPR_STMT:
+		gen_expr(node->lhs);
+		printf("\tPOP ( discard stack )\n");
+		depth -= 1;
+		break;
+	case NOD_RETURN_STMT:
+		gen_expr(node->lhs);
+		printf("\tJMP ,__fnret_%.*s\n", fname->len, fname->loc);
+		depth -= 1;
+		break;
+	case NOD_WRT_STMT:
+		gen_expr(node->lhs);
+		printf("\tWRT\n");
+		depth -= 1;
+		break;
+	case NOD_ASSIGN_STMT:
+		gen_expr(node->rhs);
+		const int found = node_find(locals, node->lhs);
+		if (found != -1) {
+			printf("\tLIT ,__stack_ptr LDA\n");
+			if (locals_size - found - 1)
+				printf("\tLIT %04x SUB\n",
+				       locals_size - found - 1);
+			printf("\tSTA\n");
+		} else {
+			printf("\tLIT ,__gl_%.*s\n",
+			       node->lhs->len, node->lhs->loc);
+			printf("\tSTA\n");
+		}
+		depth -= 1;
+		break;
+	case NOD_IFELSE_STMT:
+		gen_stmt(node->lhs, fname);
+		if (node->rhs != NULL)
+			gen_stmt(node->rhs, fname);
+		break;
+	default:
+		error("unexpected %d", node->type);
+	}
+}
+
+static void
 gen_function(Node *node)
 {
 	Node *const ognode = node;
 	printf("@__fn_%.*s\n", node->lhs->len, node->lhs->loc);
 
-	/* GOTO */
 	Node head = {0};
 	if (node->lhs->next)
 		head.next = node->lhs->next;
@@ -657,6 +738,9 @@ gen_function(Node *node)
 
 	Node *cur = node_tail(&head);
 	node = node->rhs;
+	if (node->type == NOD_BLOCK_STMT)
+		node = node->lhs;
+
 	while (node->type == NOD_LOCAL) {
 		cur = cur->next = node->lhs;
 		node = node->next;
@@ -681,41 +765,7 @@ gen_function(Node *node)
 	}
 
 	while (node != NULL) {
-		switch (node->type) {
-		case NOD_EXPR_STMT:
-			gen_expr(node->lhs);
-			printf("\tPOP ( discard stack )\n");
-			depth -= 1;
-			break;
-		case NOD_RETURN_STMT:
-			gen_expr(node->lhs);
-			printf("\tJMP ,__fnret_%.*s\n", ognode->lhs->len, ognode->lhs->loc);
-			depth -= 1;
-			break;
-		case NOD_WRT_STMT:
-			gen_expr(node->lhs);
-			printf("\tWRT\n");
-			depth -= 1;
-			break;
-		case NOD_ASSIGN_STMT:
-			gen_expr(node->rhs);
-			const int found = node_find(locals, node->lhs);
-			if (found != -1) {
-				printf("\tLIT ,__stack_ptr LDA\n");
-				if (locals_size - found - 1)
-					printf("\tLIT %04x SUB\n",
-					       locals_size - found - 1);
-				printf("\tSTA\n");
-			} else {
-				printf("\tLIT ,__gl_%.*s\n",
-				       node->lhs->len, node->lhs->loc);
-				printf("\tSTA\n");
-			}
-			depth -= 1;
-			break;
-		default:
-			error("unexpected %d", node->type);
-		}
+		gen_stmt(node, ognode->lhs);
 		node = node->next;
 	}