implemented while and for loops

lib/interpreter.ml

@@ -9,8 +9,9 @@ open Value
 let value_of_literal (literal : literal) : Value.lox_value =
   match literal with String s -> String s | Number x -> Number x | Bool b -> Bool b | Nil -> Nil
 
-let rec interpret_expr (env : environment) (expr : expr_node) : (lox_value, runtime_error) result =
-  let { pos; expr } = expr in
+let rec interpret_expr (env : environment) (expr_node : expr_node) :
+    (lox_value, runtime_error) result =
+  let { pos; expr } = expr_node in
   match expr with
   | Literal literal -> Ok (value_of_literal literal)
   | Variable name -> (
@@ -75,8 +76,8 @@ let rec interpret_expr (env : environment) (expr : expr_node) : (lox_value, runt
       | And, false | Or, true -> Ok left (* short circuit *)
       | _ -> interpret_expr env right)
 
-let rec interpret_stmt (env : environment) (stmt : stmt_node) : (unit, runtime_error) result =
-  let { pos; stmt } = stmt in
+let rec interpret_stmt (env : environment) (stmt_node : stmt_node) : (unit, runtime_error) result =
+  let { pos; stmt } = stmt_node in
   ignore pos;
   match stmt with
   | Expr expr ->
@@ -108,3 +109,10 @@ let rec interpret_stmt (env : environment) (stmt : stmt_node) : (unit, runtime_e
       let cond = lox_value_to_bool cond in
       if cond then interpret_stmt env then_
       else Option.map (interpret_stmt env) else_ |> Option.value ~default:(Ok ())
+  | While { cond; body } ->
+      let* cond = interpret_expr env cond in
+      let cond = lox_value_to_bool cond in
+      if cond then
+        let* _ = interpret_stmt env body in
+        interpret_stmt env stmt_node
+      else Ok ()

lib/lox.ml

@@ -27,7 +27,7 @@ let run ?(env : Environment.environment option) ?(debug = false) (source : strin
       print_tokens ()
     else ()
   in
-  let* stmts = Error.of_parser_error (Parser.parse tokens) in
+  let* stmts = Parser.parse tokens |> Error.of_parser_error in
   let () =
     if debug then
       let print_statements () =

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

lib/stmt.ml

@@ -6,13 +6,14 @@ type stmt =
   | VarDecl of { name : string; init : expr_node option }
   | Block of stmt_node list
   | If of { cond : expr_node; then_ : stmt_node; else_ : stmt_node option }
+  | While of { cond : expr_node; body : stmt_node }
 
 and stmt_node = { stmt : stmt; pos : Error.code_pos }
 
 let rec show_stmt ?(indent = 0) stmt =
   let indent_s = String.make indent ' ' in
-  let show_expr_ind ?(depth = 2) = show_expr ~indent:(indent + depth) in
-  let show_stmt_ind ?(depth = 2) = show_stmt ~indent:(indent + depth) in
+  let show_expr_ind ?(add = 2) = show_expr ~indent:(indent + add) in
+  let show_stmt_ind ?(add = 2) = show_stmt ~indent:(indent + add) in
   match stmt with
   | Expr expr -> indent_s ^ "Expr\n" ^ show_expr_ind expr.expr
   | Print expr -> indent_s ^ "Print\n" ^ show_expr_ind expr.expr
@@ -23,13 +24,17 @@ let rec show_stmt ?(indent = 0) stmt =
       let stmts_s =
         List.fold_left (fun acc stmt -> acc ^ show_stmt_ind stmt.stmt ^ "\n") "" stmts
       in
-      "Block" ^ stmts_s ^ "End"
+      indent_s ^ "Block\n" ^ stmts_s ^ indent_s ^ "End"
   | If { cond; then_; else_ } ->
       let cond_s = show_expr_ind cond.expr in
-      let then_s = show_stmt_ind ~depth:4 then_.stmt in
-      let else_s = Option.map (fun stmt -> show_stmt_ind ~depth:4 stmt.stmt) else_ in
-      indent_s ^ "If\n" ^ cond_s ^ "\n" ^ indent_s ^ "  Then\n" ^ then_s
+      let then_s = show_stmt_ind ~add:4 then_.stmt in
+      let else_s = Option.map (fun stmt -> show_stmt_ind ~add:4 stmt.stmt) else_ in
+      indent_s ^ "If\n" ^ indent_s ^ "  Cond" ^ cond_s ^ "\n" ^ indent_s ^ "  Then\n" ^ then_s
       ^ if Option.is_some else_s then "\n" ^ indent_s ^ "  Else\n" ^ Option.get else_s else ""
+  | While { cond; body } ->
+      let cond_s = show_expr_ind ~add:4 cond.expr in
+      let body_s = show_stmt_ind ~add:4 body.stmt in
+      indent_s ^ "While\n" ^ indent_s ^ "  Cond\n" ^ cond_s ^ "\n" ^ indent_s ^ "  Body\n" ^ body_s
 
 let show_stmt_node stmt_node = show_stmt stmt_node.stmt
 let make_stmt_node (pos : Error.code_pos) (stmt : stmt) : stmt_node = { stmt; pos }
@@ -49,3 +54,6 @@ let make_block (pos : Error.code_pos) (stmts : stmt_node list) : stmt_node =
 let make_if (pos : Error.code_pos) (cond : expr_node) (then_ : stmt_node) (else_ : stmt_node option)
     =
   If { cond; then_; else_ } |> make_stmt_node pos
+
+let make_while (pos : Error.code_pos) (cond : expr_node) (body : stmt_node) =
+  While { cond; body } |> make_stmt_node pos