fix stack (i hate myself)

1 files changed, 19 insertions, 10 deletions

diff --git a/src/main.c b/src/main.c
index 3dff1fe..978dd31 100644
--- a/src/main.c
+++ b/src/main.c
@@ -819,7 +819,10 @@ gen_stmt(Node *node, Node *fname, int break_lbl)
 		break;
 	case NOD_RETURN_STMT:
 		gen_expr(node->lhs);
-		printf("\tJMP ,__fnret_%.*s\n", fname->len, fname->loc);
+		if (locals_size)
+			printf("\tJMP ,__fnret_%.*s\n", fname->len, fname->loc);
+		else
+			printf("\tRET\n");
 		depth -= 1;
 		break;
 	case NOD_BREAK_STMT:
@@ -887,9 +890,11 @@ gen_function(Node *node)
 
 	locals_size = node_size(locals);
 
-	printf("\tLIT ,__stack_ptr LDA\n");
-	printf("\tLIT %04x ADD\n", locals_size);
-	printf("\tLIT ,__stack_ptr STA\n");
+	if (locals_size) {
+		printf("\tLIT ,__stack_ptr LDA\n");
+		printf("\tLIT %04x ADD\n", locals_size);
+		printf("\tLIT ,__stack_ptr STA\n");
+	}
 
 	/* TODO: optimize this using DUP */
 	cur = locals;
@@ -907,11 +912,13 @@ gen_function(Node *node)
 		node = node->next;
 	}
 
-	printf("\tLIT 0000 ( default return )\n");
-	printf("@__fnret_%.*s\n", ognode->lhs->len, ognode->lhs->loc);
-	printf("\tLIT ,__stack_ptr LDA\n");
-	printf("\tLIT %04x SUB\n", locals_size);
-	printf("\tLIT ,__stack_ptr STA\n");
+	printf("\tLIT 0000\n");
+	if (locals_size) {
+		printf("@__fnret_%.*s\n", ognode->lhs->len, ognode->lhs->loc);
+		printf("\tLIT ,__stack_ptr LDA\n");
+		printf("\tLIT %04x SUB\n", locals_size);
+		printf("\tLIT ,__stack_ptr STA\n");
+	}
 	printf("\tRET\n");
 }
 
@@ -975,10 +982,12 @@ main(int argc, char **argv)
 
 	Node *node = parse(tok);
 
-	printf("\tJRT ,__fn_main ( entry point )\n");
+	printf("\tJRT ,__fn_main\n");
 	printf("\tRET\n");
 	codegen(node);
 	printf("@__stack_ptr\n");
+	printf("\t,__stack\n");
+	printf("@__stack\n");
 
 	assert(depth == 0);
 	return 0;