From 1fd27bf0a9cb69c7e94934b79ed7f8cce86cbfea Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Thu, 27 Aug 2020 22:00:21 +0900
Subject: [PATCH] Add <<, >>, <<= and >>=

---
 chibicc.h    |  2 ++
 codegen.c    | 11 +++++++++++
 parse.c      | 42 ++++++++++++++++++++++++++++++++++++------
 test/arith.c | 13 +++++++++++++
 tokenize.c   | 12 ++++++++++--
 type.c       |  2 ++
 6 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/chibicc.h b/chibicc.h
index a0149887..aef92f11 100644
--- a/chibicc.h
+++ b/chibicc.h
@@ -94,6 +94,8 @@ typedef enum {
   ND_BITAND,    // &
   ND_BITOR,     // |
   ND_BITXOR,    // ^
+  ND_SHL,       // <<
+  ND_SHR,       // >>
   ND_EQ,        // ==
   ND_NE,        // !=
   ND_LT,        // <
diff --git a/codegen.c b/codegen.c
index 05f29a00..1119e8e8 100644
--- a/codegen.c
+++ b/codegen.c
@@ -328,6 +328,17 @@ static void gen_expr(Node *node) {
 
     println("  movzb %%al, %%rax");
     return;
+  case ND_SHL:
+    println("  mov %%rdi, %%rcx");
+    println("  shl %%cl, %s", ax);
+    return;
+  case ND_SHR:
+    println("  mov %%rdi, %%rcx");
+    if (node->ty->size == 8)
+      println("  sar %%cl, %s", ax);
+    else
+      println("  sar %%cl, %s", ax);
+    return;
   }
 
   error_tok(node->tok, "invalid expression");
diff --git a/parse.c b/parse.c
index 8e472171..89274a05 100644
--- a/parse.c
+++ b/parse.c
@@ -103,6 +103,7 @@ static Node *bitxor(Token **rest, Token *tok);
 static Node *bitand(Token **rest, Token *tok);
 static Node *equality(Token **rest, Token *tok);
 static Node *relational(Token **rest, Token *tok);
+static Node *shift(Token **rest, Token *tok);
 static Node *add(Token **rest, Token *tok);
 static Node *new_add(Node *lhs, Node *rhs, Token *tok);
 static Node *new_sub(Node *lhs, Node *rhs, Token *tok);
@@ -857,6 +858,7 @@ static Node *to_assign(Node *binary) {
 
 // assign    = logor (assign-op assign)?
 // assign-op = "=" | "+=" | "-=" | "*=" | "/=" | "%=" | "&=" | "|=" | "^="
+//           | "<<=" | ">>="
 static Node *assign(Token **rest, Token *tok) {
   Node *node = logor(&tok, tok);
 
@@ -887,6 +889,12 @@ static Node *assign(Token **rest, Token *tok) {
   if (equal(tok, "^="))
     return to_assign(new_binary(ND_BITXOR, node, assign(rest, tok->next), tok));
 
+  if (equal(tok, "<<="))
+    return to_assign(new_binary(ND_SHL, node, assign(rest, tok->next), tok));
+
+  if (equal(tok, ">>="))
+    return to_assign(new_binary(ND_SHR, node, assign(rest, tok->next), tok));
+
   *rest = tok;
   return node;
 }
@@ -968,30 +976,52 @@ static Node *equality(Token **rest, Token *tok) {
   }
 }
 
-// relational = add ("<" add | "<=" add | ">" add | ">=" add)*
+// relational = shift ("<" shift | "<=" shift | ">" shift | ">=" shift)*
 static Node *relational(Token **rest, Token *tok) {
-  Node *node = add(&tok, tok);
+  Node *node = shift(&tok, tok);
 
   for (;;) {
     Token *start = tok;
 
     if (equal(tok, "<")) {
-      node = new_binary(ND_LT, node, add(&tok, tok->next), start);
+      node = new_binary(ND_LT, node, shift(&tok, tok->next), start);
       continue;
     }
 
     if (equal(tok, "<=")) {
-      node = new_binary(ND_LE, node, add(&tok, tok->next), start);
+      node = new_binary(ND_LE, node, shift(&tok, tok->next), start);
       continue;
     }
 
     if (equal(tok, ">")) {
-      node = new_binary(ND_LT, add(&tok, tok->next), node, start);
+      node = new_binary(ND_LT, shift(&tok, tok->next), node, start);
       continue;
     }
 
     if (equal(tok, ">=")) {
-      node = new_binary(ND_LE, add(&tok, tok->next), node, start);
+      node = new_binary(ND_LE, shift(&tok, tok->next), node, start);
+      continue;
+    }
+
+    *rest = tok;
+    return node;
+  }
+}
+
+// shift = add ("<<" add | ">>" add)*
+static Node *shift(Token **rest, Token *tok) {
+  Node *node = add(&tok, tok);
+
+  for (;;) {
+    Token *start = tok;
+
+    if (equal(tok, "<<")) {
+      node = new_binary(ND_SHL, node, add(&tok, tok->next), start);
+      continue;
+    }
+
+    if (equal(tok, ">>")) {
+      node = new_binary(ND_SHR, node, add(&tok, tok->next), start);
       continue;
     }
 
diff --git a/test/arith.c b/test/arith.c
index 164627a4..e9721b0c 100644
--- a/test/arith.c
+++ b/test/arith.c
@@ -95,6 +95,19 @@ int main() {
   ASSERT(7, ({ int i=6; i|=3; i; }));
   ASSERT(10, ({ int i=15; i^=5; i; }));
 
+  ASSERT(1, 1<<0);
+  ASSERT(8, 1<<3);
+  ASSERT(10, 5<<1);
+  ASSERT(2, 5>>1);
+  ASSERT(-1, -1>>1);
+  ASSERT(1, ({ int i=1; i<<=0; i; }));
+  ASSERT(8, ({ int i=1; i<<=3; i; }));
+  ASSERT(10, ({ int i=5; i<<=1; i; }));
+  ASSERT(2, ({ int i=5; i>>=1; i; }));
+  ASSERT(-1, -1);
+  ASSERT(-1, ({ int i=-1; i; }));
+  ASSERT(-1, ({ int i=-1; i>>=1; i; }));
+
   printf("OK\n");
   return 0;
 }
diff --git a/tokenize.c b/tokenize.c
index 9bfe959c..4f5d834c 100644
--- a/tokenize.c
+++ b/tokenize.c
@@ -332,7 +332,14 @@ static Token *tokenize(char *filename, char *p) {
       continue;
     }
 
-    // Multi-letter punctuators
+    // Three-letter punctuators
+    if (startswith(p, "<<=") || startswith(p, ">>=")) {
+      cur = cur->next = new_token(TK_RESERVED, p, p + 3);
+      p += 3;
+      continue;
+    }
+
+    // Two-letter punctuators
     if (startswith(p, "==") || startswith(p, "!=") ||
         startswith(p, "<=") || startswith(p, ">=") ||
         startswith(p, "->") || startswith(p, "+=") ||
@@ -341,7 +348,8 @@ static Token *tokenize(char *filename, char *p) {
         startswith(p, "--") || startswith(p, "%=") ||
         startswith(p, "&=") || startswith(p, "|=") ||
         startswith(p, "^=") || startswith(p, "&&") ||
-        startswith(p, "||")) {
+        startswith(p, "||") || startswith(p, "<<") ||
+        startswith(p, ">>")) {
       cur = cur->next = new_token(TK_RESERVED, p, p + 2);
       p += 2;
       continue;
diff --git a/type.c b/type.c
index c482d15f..c87acd68 100644
--- a/type.c
+++ b/type.c
@@ -138,6 +138,8 @@ void add_type(Node *node) {
     node->ty = ty_int;
     return;
   case ND_BITNOT:
+  case ND_SHL:
+  case ND_SHR:
     node->ty = node->lhs->ty;
     return;
   case ND_VAR: