assign to array

1 files changed, 48 insertions, 17 deletions

diff --git a/src/main.c b/src/main.c
index f7c8eee..2a15976 100644
--- a/src/main.c
+++ b/src/main.c
@@ -552,6 +552,19 @@ stmt(Token **rest, Token *tok)
 		return break_stmt(rest, tok);
 	if (equal(tok, "return"))
 		return return_stmt(rest, tok);
+	if (equal(tok, "[")) {
+		int opn = 1;
+		Token *cur = tok->next;
+		while (cur->type != TOK_EOF && opn) {
+			opn += equal(cur, "[");
+			opn -= equal(cur, "]");
+			cur = cur->next;
+		}
+		if (cur->type == TOK_EOF)
+			error("unexpected EOF");
+		if (equal(cur, "="))
+			return assign_stmt(rest, tok);
+	}
 	if (equal(tok->next, "="))
 		return assign_stmt(rest, tok);
 	return expr_stmt(rest, tok);
@@ -612,10 +625,18 @@ wrt_stmt(Token **rest, Token *tok)
 static Node *
 assign_stmt(Token **rest, Token *tok)
 {
-	expect(tok, TOK_WORD);
+	Node *lhs;
+	if (equal(tok, "[")) {
+		tok = tok->next;
+		lhs = new_node(NOD_DEREF);
+		lhs->lhs = expr(&tok, tok);
+		*rest = skip(tok, "]");
+	} else {
+		expect(tok, TOK_WORD);
+		lhs = new_word(tok);
+	}
 	skip(tok->next, "=");
-	Node *const word = new_word(tok);
-	Node *const node = new_binary(NOD_ASSIGN_STMT, word, expr(&tok, tok->next->next));
+	Node *const node = new_binary(NOD_ASSIGN_STMT, lhs, expr(&tok, tok->next->next));
 	*rest = skip(tok, ";");
 	return node;
 }
@@ -756,6 +777,29 @@ gen_deref(Node *node)
 }
 
 static void
+gen_assign_stmt(Node *node)
+{
+	gen_expr(node->rhs);
+	if (node->lhs->type == NOD_DEREF) {
+		gen_expr(node->lhs->lhs);
+		printf("\tSTA\n");
+		depth -= 2;
+	} else {
+		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);
+		} else
+			printf("\tLIT ,__gl_%.*s\n",
+			       node->lhs->len, node->lhs->loc);
+		printf("\tSTA\n");
+		depth -= 1;
+	}
+}
+
+static void
 gen_stmt(Node *node, Node *fname, int break_lbl)
 {
 	static int label = 0;
@@ -789,20 +833,7 @@ gen_stmt(Node *node, Node *fname, int break_lbl)
 		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;
+		gen_assign_stmt(node);
 		break;
 	case NOD_IFELSE_STMT:
 		lbl = label;