implemented while and for loops

lib/parser.ml

@@ -251,6 +251,63 @@ and if_then_else (state : state ref) : stmt_result =
   in
   make_if pos cond then_ else_ |> Result.ok
 
+and while_loop (state : state ref) : stmt_result =
+  let pos = cur_pos state in
+  let* _ = consume state While in
+  let* _ = consume state LeftParen in
+  let* cond = expression state in
+  let* _ = consume state RightParen in
+  let* body = statement state in
+  make_while pos cond body |> Result.ok
+
+and for_loop (state : state ref) : stmt_result =
+  let for_pos = cur_pos state in
+  let* _ = consume state For in
+  let* _ = consume state LeftParen in
+  let* init =
+    match peek_tt state with
+    | Semicolon ->
+        advance state;
+        Ok None
+    | Var -> var_declaration state |> Result.map Option.some
+    | _ -> expr_stmt state |> Result.map Option.some
+  in
+  let* cond =
+    match peek_tt state with
+    | Semicolon ->
+        let pos = cur_pos state in
+        Ok (make_bool pos true)
+    | _ -> expression state
+  in
+  (* expression has no final semicolon, so we need to consume it *)
+  let* _ = consume state Semicolon in
+  let* update =
+    match peek_tt state with
+    | RightParen -> Ok None
+    | _ -> expression state |> Result.map Option.some
+  in
+  let* _ = consume state RightParen in
+  let* body = statement state in
+  let body =
+    match update with
+    | Some update ->
+        let update_stmt = make_expr_stmt update.pos update in
+        make_block body.pos [ body; update_stmt ]
+    | None -> body
+  in
+  let loop = make_while for_pos cond body in
+  let outer_block =
+    match init with Some init -> make_block for_pos [ init; loop ] | None -> loop
+  in
+  Ok outer_block
+
+and expr_stmt (state : state ref) : stmt_result =
+  let pos = cur_pos state in
+  let* expr = expression state in
+  let* _ = consume state Semicolon in
+  let stmt = make_expr_stmt pos expr in
+  Ok stmt
+
 and statement (state : state ref) : stmt_result =
   let pos = cur_pos state in
   match peek_tt state with
@@ -262,11 +319,9 @@ and statement (state : state ref) : stmt_result =
       Ok stmt
   | LeftBrace -> block state
   | If -> if_then_else state
-  | _ ->
-      let* expr = expression state in
-      let* _ = consume state Semicolon in
-      let stmt = make_expr_stmt pos expr in
-      Ok stmt
+  | While -> while_loop state
+  | For -> for_loop state
+  | _ -> expr_stmt state
 
 and var_declaration (state : state ref) : stmt_result =
   let pos = cur_pos state in