i'm a sheeple

1 files changed, 161 insertions, 22 deletions

diff --git a/src/main.c b/src/main.c
index ba67715..a38a66a 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <ctype.h>
 #include <stdarg.h>
 #include <stdbool.h>
@@ -79,6 +80,12 @@ new_token(TokenType type, char *start, char *end)
 	return tok;
 }
 
+static bool
+is_punct(int c)
+{
+	return (strchr("+-/*();", c) != NULL);
+}
+
 static Token *
 tokenize(char *p)
 {
@@ -99,7 +106,7 @@ tokenize(char *p)
 			continue;
 		}
 
-		if (strchr("+-/*()", *p) != NULL) {
+		if (is_punct(*p)) {
 			cur = cur->next = new_token(TOK_PUNCT, p, p + 1);
 			p += 1;
 			continue;
@@ -120,6 +127,153 @@ tokenize(char *p)
 	return head.next;
 }
 
+typedef enum {
+	NOD_ADD,
+	NOD_SUB,
+	NOD_MUL,
+	NOD_DIV,
+	NOD_NUM,
+} NodeType;
+
+typedef struct Node Node;
+struct Node {
+	NodeType type;
+	Node *lhs;
+	Node *rhs;
+	int val;
+};
+
+static Node *
+new_node(NodeType type)
+{
+	Node *const node = calloc(1, sizeof(Node));
+	if (node == NULL)
+		error("calloc failed");
+	node->type = type;
+	return node;
+}
+
+static Node *
+new_binary(NodeType type, Node *lhs, Node *rhs)
+{
+	Node *const node = new_node(type);
+	node->lhs = lhs;
+	node->rhs = rhs;
+	return node;
+}
+
+static Node *
+new_num(int val)
+{
+	Node *const node = new_node(NOD_NUM);
+	node->val = val;
+	return node;
+}
+
+static Node *expr(Token **rest, Token *tok);
+static Node *mul(Token **rest, Token *tok);
+static Node *primary(Token **rest, Token *tok);
+
+// expr = mul ("+" mul | "-" mul)*
+static Node *
+expr(Token **rest, Token *tok)
+{
+	Node *node = mul(&tok, tok);
+
+	for (;;) {
+		if (equal(tok, "+")) {
+			node = new_binary(NOD_ADD, node, mul(&tok, tok->next));
+			continue;
+		}
+
+		if (equal(tok, "-")) {
+			node = new_binary(NOD_SUB, node, mul(&tok, tok->next));
+			continue;
+		}
+
+		*rest = tok;
+		return node;
+	}
+}
+
+// mul = primary ("*" primary | "/" primary)*
+static Node *
+mul(Token **rest, Token *tok)
+{
+	Node *node = primary(&tok, tok);
+
+	for (;;) {
+		if (equal(tok, "*")) {
+			node = new_binary(NOD_MUL, node, primary(&tok, tok->next));
+			continue;
+		}
+
+		if (equal(tok, "/")) {
+			node = new_binary(NOD_DIV, node, primary(&tok, tok->next));
+			continue;
+		}
+
+		*rest = tok;
+		return node;
+	}
+}
+
+static Node *
+primary(Token **rest, Token *tok)
+{
+	if (equal(tok, "(")) {
+		Node *node = expr(&tok, tok->next);
+		*rest = skip(tok, ")");
+		return node;
+	}
+
+	if (tok->type == TOK_NUM) {
+		Node *node = new_num(tok->val);
+		*rest = tok->next;
+		return node;
+	}
+
+	error("expected an expression");
+	return NULL;
+}
+
+// code generator
+
+static int depth;
+
+static void
+gen_expr(Node *node)
+{
+	if (node->type == NOD_NUM) {
+		printf("\tLIT %04d\n", node->val);
+		depth += 1;
+		return;
+	}
+
+	gen_expr(node->lhs);
+	gen_expr(node->rhs);
+
+	depth -= 1;
+	switch (node->type) {
+	case NOD_ADD:
+		printf("\tADD\n");
+		return;
+	case NOD_SUB:
+		printf("\tSUB\n");
+		return;
+	case NOD_MUL:
+		printf("\tMUL\n");
+		return;
+	case NOD_DIV:
+		printf("\tDIV\n");
+		return;
+	default:
+		break;
+	}
+
+	error("invalid expression %d", node->type);
+}
+
 int
 main(int argc, char **argv)
 {
@@ -128,30 +282,15 @@ main(int argc, char **argv)
 
 	Token *tok = tokenize(argv[1]);
 
-	/* first token is a number */
-	printf("main:\n");
-	printf("\tLIT %04d\n", get_number(tok));
-	tok = tok->next;
-
-	/* followed by punct then another number */
-	while (tok->type != TOK_EOF) {
-		const char *punct;
-		expect(tok, TOK_PUNCT);
-		switch (tok->loc[0]) {
-		case '+': punct = "ADD"; break;
-		case '-': punct = "SUB"; break;
-		case '*': punct = "MUL"; break;
-		case '/': punct = "DIV"; break;
-		default: break;
-		}
-		tok = tok->next;
+	Node *node = expr(&tok, tok);
 
-		printf("\tLIT %04d %s\n", get_number(tok), punct);
-		tok = tok->next;
-	}
-	expect(tok, TOK_EOF);
+	if (tok->type != TOK_EOF)
+		error("extra token");
 
+	printf("main:\n");
+	gen_expr(node);
 	printf("\tRET\n");
 
+	assert(depth == 1);
 	return 0;
 }