fix store order

1 files changed, 39 insertions, 9 deletions

diff --git a/src/main.c b/src/main.c
index dbf39df..ae6b6cd 100644
--- a/src/main.c
+++ b/src/main.c
@@ -197,10 +197,24 @@ node_tail(Node *node)
 	return node;
 }
 
+/* doesn't check for \0 in op beware */
 static bool
 node_equal(const Node *node, const char *op)
 {
-	return memcmp(node->loc, op, node->len) == 0 && op[node->len] == '\0';
+	return memcmp(node->loc, op, node->len) == 0;
+}
+
+static int
+node_find(const Node *node, Node *op)
+{
+	int i = 0;
+	while (node != NULL) {
+		if (node->len == op->len && node_equal(node, op->loc))
+			return i;
+		i += 1;
+		node = node->next;
+	}
+	return -1;
 }
 
 static Node *
@@ -263,6 +277,7 @@ static Node *mul(Token **rest, Token *tok);
 static Node *primary(Token **rest, Token *tok);
 
 static Node *locals;
+static int locals_size;
 
 static Node *
 function(Token **rest, Token *tok)
@@ -627,17 +642,20 @@ gen_function(Node *node)
 	printf("@__fn_%.*s\n", node->lhs->len, node->lhs->loc);
 
 	/* GOTO */
-	locals = node->lhs->next;
-	const int argcount = node_size(locals);
+	Node head = {0};
+	if (node->lhs->next)
+		head.next = node->lhs->next;
+	const int argcount = node_size(head.next);
 
-	Node *cur = node_tail(locals);
+	Node *cur = node_tail(&head);
 	node = node->rhs;
 	while (node->type == NOD_LOCAL) {
 		cur = cur->next = node->lhs;
 		node = node->next;
 	}
+	locals = head.next;
 
-	const int locals_size = node_size(locals);
+	locals_size = node_size(locals);
 
 	printf("\tLIT ,__stack_ptr LDA\n");
 	printf("\tLIT %04x ADD\n", locals_size);
@@ -646,8 +664,10 @@ gen_function(Node *node)
 	/* TODO: optimize this using DUP */
 	cur = locals;
 	for (int i = 0; i < argcount; i++) {
-		printf("\tLIT ,__stack_ptr LDA ( %.*s )\n", cur->len, cur->loc);
-		printf("\tLIT %04x SUB\n", locals_size - i - 1);
+		printf("\tLIT ,__stack_ptr LDA\n");
+		const int offset = locals_size + i - argcount;
+		if (offset)
+			printf("\tLIT %04x SUB\n", offset);
 		printf("\tSWP STA\n");
 		cur = cur->next;
 	}
@@ -671,8 +691,18 @@ gen_function(Node *node)
 			break;
 		case NOD_ASSIGN_STMT:
 			gen_expr(node->rhs);
-			printf("\tLIT ,__gl_%.*s\n\tSTA\n",
-			       node->lhs->len, node->lhs->loc);
+			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: