From 773115ab2a9c4b96f804311b95b20e9771f0190a Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Sun, 4 Aug 2019 17:24:03 +0900
Subject: [PATCH] Add "while" statement

---
 chibicc.h  |  2 +-
 codegen.c  |  3 ++-
 parse.c    | 10 ++++++++++
 test.sh    |  2 ++
 tokenize.c |  2 +-
 5 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/chibicc.h b/chibicc.h
index 9f7acd65..c0e38b5f 100644
--- a/chibicc.h
+++ b/chibicc.h
@@ -72,7 +72,7 @@ typedef enum {
   ND_ASSIGN,    // =
   ND_RETURN,    // "return"
   ND_IF,        // "if"
-  ND_FOR,       // "for"
+  ND_FOR,       // "for" or "while"
   ND_BLOCK,     // { ... }
   ND_EXPR_STMT, // Expression statement
   ND_VAR,       // Variable
diff --git a/codegen.c b/codegen.c
index 944ff272..827a6322 100644
--- a/codegen.c
+++ b/codegen.c
@@ -115,7 +115,8 @@ static void gen_stmt(Node *node) {
   }
   case ND_FOR: {
     int c = count();
-    gen_stmt(node->init);
+    if (node->init)
+      gen_stmt(node->init);
     printf(".L.begin.%d:\n", c);
     if (node->cond) {
       gen_expr(node->cond);
diff --git a/parse.c b/parse.c
index 334d58cb..fde080af 100644
--- a/parse.c
+++ b/parse.c
@@ -65,6 +65,7 @@ static Obj *new_lvar(char *name) {
 // stmt = "return" expr ";"
 //      | "if" "(" expr ")" stmt ("else" stmt)?
 //      | "for" "(" expr-stmt expr? ";" expr? ")" stmt
+//      | "while" "(" expr ")" stmt
 //      | "{" compound-stmt
 //      | expr-stmt
 static Node *stmt(Token **rest, Token *tok) {
@@ -104,6 +105,15 @@ static Node *stmt(Token **rest, Token *tok) {
     return node;
   }
 
+  if (equal(tok, "while")) {
+    Node *node = new_node(ND_FOR);
+    tok = skip(tok->next, "(");
+    node->cond = expr(&tok, tok);
+    tok = skip(tok, ")");
+    node->then = stmt(rest, tok);
+    return node;
+  }
+
   if (equal(tok, "{"))
     return compound_stmt(rest, tok->next);
 
diff --git a/test.sh b/test.sh
index 1140e96f..f2d5d01d 100755
--- a/test.sh
+++ b/test.sh
@@ -72,4 +72,6 @@ assert 3 '{ if (1) { 1; 2; return 3; } else { return 4; } }'
 assert 55 '{ i=0; j=0; for (i=0; i<=10; i=i+1) j=i+j; return j; }'
 assert 3 '{ for (;;) {return 3;} return 5; }'
 
+assert 10 '{ i=0; while(i<10) { i=i+1; } return i; }'
+
 echo OK
diff --git a/tokenize.c b/tokenize.c
index 5e66d84c..daf98b81 100644
--- a/tokenize.c
+++ b/tokenize.c
@@ -72,7 +72,7 @@ static bool is_ident2(char c) {
 }
 
 static bool is_keyword(Token *tok) {
-  static char *kw[] = {"return", "if", "else", "for"};
+  static char *kw[] = {"return", "if", "else", "for", "while"};
 
   for (int i = 0; i < sizeof(kw) / sizeof(*kw); i++)
     if (equal(tok, kw[i]))