From bf9ab52860c1cbbeeca40df515468f42300ff429 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Sat, 3 Aug 2019 15:41:40 +0900
Subject: [PATCH] Add unary plus and minus

---
 main.c  | 35 ++++++++++++++++++++++++++++++-----
 test.sh |  3 +++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/main.c b/main.c
index 791d481..ae300a0 100644
--- a/main.c
+++ b/main.c
@@ -134,6 +134,7 @@ typedef enum {
   ND_SUB, // -
   ND_MUL, // *
   ND_DIV, // /
+  ND_NEG, // unary -
   ND_NUM, // Integer
 } NodeKind;
 
@@ -159,6 +160,12 @@ static Node *new_binary(NodeKind kind, Node *lhs, Node *rhs) {
   return node;
 }
 
+static Node *new_unary(NodeKind kind, Node *expr) {
+  Node *node = new_node(kind);
+  node->lhs = expr;
+  return node;
+}
+
 static Node *new_num(int val) {
   Node *node = new_node(ND_NUM);
   node->val = val;
@@ -167,6 +174,7 @@ static Node *new_num(int val) {
 
 static Node *expr(Token **rest, Token *tok);
 static Node *mul(Token **rest, Token *tok);
+static Node *unary(Token **rest, Token *tok);
 static Node *primary(Token **rest, Token *tok);
 
 // expr = mul ("+" mul | "-" mul)*
@@ -189,18 +197,18 @@ static Node *expr(Token **rest, Token *tok) {
   }
 }
 
-// mul = primary ("*" primary | "/" primary)*
+// mul = unary ("*" unary | "/" unary)*
 static Node *mul(Token **rest, Token *tok) {
-  Node *node = primary(&tok, tok);
+  Node *node = unary(&tok, tok);
 
   for (;;) {
     if (equal(tok, "*")) {
-      node = new_binary(ND_MUL, node, primary(&tok, tok->next));
+      node = new_binary(ND_MUL, node, unary(&tok, tok->next));
       continue;
     }
 
     if (equal(tok, "/")) {
-      node = new_binary(ND_DIV, node, primary(&tok, tok->next));
+      node = new_binary(ND_DIV, node, unary(&tok, tok->next));
       continue;
     }
 
@@ -209,6 +217,18 @@ static Node *mul(Token **rest, Token *tok) {
   }
 }
 
+// unary = ("+" | "-") unary
+//       | primary
+static Node *unary(Token **rest, Token *tok) {
+  if (equal(tok, "+"))
+    return unary(rest, tok->next);
+
+  if (equal(tok, "-"))
+    return new_unary(ND_NEG, unary(rest, tok->next));
+
+  return primary(rest, tok);
+}
+
 // primary = "(" expr ")" | num
 static Node *primary(Token **rest, Token *tok) {
   if (equal(tok, "(")) {
@@ -243,9 +263,14 @@ static void pop(char *arg) {
 }
 
 static void gen_expr(Node *node) {
-  if (node->kind == ND_NUM) {
+  switch (node->kind) {
+  case ND_NUM:
     printf("  mov $%d, %%rax\n", node->val);
     return;
+  case ND_NEG:
+    gen_expr(node->lhs);
+    printf("  neg %%rax\n");
+    return;
   }
 
   gen_expr(node->rhs);
diff --git a/test.sh b/test.sh
index d882d77..3011b04 100755
--- a/test.sh
+++ b/test.sh
@@ -23,5 +23,8 @@ assert 41 ' 12 + 34 - 5 '
 assert 47 '5+6*7'
 assert 15 '5*(9-6)'
 assert 4 '(3+5)/2'
+assert 10 '-10+20'
+assert 10 '- -10'
+assert 10 '- - +10'
 
 echo OK