Implemented transformation of for-loops to while-loops

main.ml

@@ -13,11 +13,11 @@ let compile args =
         Load.phase;
         (*Print.phase;*)
         Parse.phase;
-        (*Print.phase;*)
+        Print.phase;
         Desug.phase;
         Print.phase;
-        Context_analysis.phase;
         (*
+        Context_analysis.phase;
         Print.phase;
         Typecheck.phase;
         Extern_vars.phase;

phases/desug.ml

@@ -1,12 +1,81 @@
 open Ast
 open Util
 
+let block_body = function
+    | Block nodes -> nodes
+    | _ -> raise InvalidNode
+
 let rec flatten_blocks = function
     | [] -> []
     | Block nodes :: t -> (flatten_blocks nodes) @ (flatten_blocks t)
     | h :: t -> h :: (flatten_blocks t)
 
+let rec replace_var var replacement = function
+    | Var (name, loc) when name = var -> Var (replacement, loc)
+    | node -> transform_children (replace_var var replacement) node
+
+let for_to_while node =
+    let new_vars = ref [] in
+    let rec traverse = function
+        (* Do not traverse into local functions (already done by var_init) *)
+        | FunDef (_, _, _, _, _, _) as node -> node
+
+        (* Transform for-loops to while-loops *)
+        | For (counter, start, stop, step, body, loc) ->
+            let _i = fresh_var counter in
+            let _stop = fresh_var "stop" in
+            let _step = fresh_var "step" in
+            new_vars := !new_vars @ [_i; _stop; _step];
+
+            let vi = Var (_i, noloc) in
+            let vstop = Var (_stop, locof stop) in
+            let vstep = Var (_step, locof step) in
+            let cond = Cond (
+                Binop (Gt, vstep, IntConst (0, noloc), noloc),
+                Binop (Lt, vi, vstop, noloc),
+                Binop (Gt, vi, vstop, noloc),
+                noloc
+            ) in
+            Block [
+                Assign (_i, start, locof start);
+                Assign (_stop, stop, locof stop);
+                Assign (_step, step, locof step);
+                While (cond, traverse (Block (
+                    (* TODO: check for illegal assigments of counter in body*)
+                    block_body (replace_var counter _i body) @
+                    [Assign (_i, Binop (Add, vi, vstep, noloc), noloc)]
+                )), loc);
+            ]
+
+        | node -> transform_children traverse node
+    in
+    let node = traverse node in
+    (node, new_vars)
+
 let rec var_init = function
+    (* Move global initialisations to __init function *)
+    | Program (decls, loc) ->
+        let decls = flatten_blocks (List.map var_init decls) in
+        let rec trav assigns = function
+            | [] -> (assigns, [])
+            | (Assign (_, _, _) as h) :: t -> trav (assigns @ [h]) t
+            | h :: t ->
+                let (assigns, decls) = trav assigns t in
+                (assigns, (h :: decls))
+        in
+        let (assigns, decls) = trav [] decls in
+        (match assigns with
+            | [] -> Program (decls, loc)
+            | assigns ->
+                let init_func = FunDef (true, Void, "__init", [], Block assigns, noloc) in
+                Program (init_func :: decls, loc)
+            )
+
+    (* Move global variable initialisations to exported __init function *)
+    | GlobalDef (export, ctype, name, Some init, loc) ->
+        Block [GlobalDef (export, ctype, name, None, loc);
+               Assign (name, init, locof init)]
+
     (* Split local variable initialisations in declaration and assignment *)
     | FunDef (export, ret_type, name, params, Block body, loc) ->
         let move_inits body =
@@ -40,32 +109,14 @@ let rec var_init = function
 
                 (* rest of function body: recurse *)
                 | rest -> inits @ (List.map var_init rest)
-            in trav [] body
-        in
-        FunDef (export, ret_type, name, params, Block (move_inits body), loc)
-
-    (* Move global variable initialisations to exported __init function *)
-    | GlobalDef (export, ctype, name, Some init, loc) ->
-        Block [GlobalDef (export, ctype, name, None, loc);
-               Assign (name, init, locof init)]
-
-    (* Move global initialisations to __init function *)
-    | Program (decls, loc) ->
-        let decls = flatten_blocks (List.map var_init decls) in
-        let rec trav assigns = function
-            | [] -> (assigns, [])
-            | (Assign (_, _, _) as h) :: t -> trav (assigns @ [h]) t
-            | h :: t ->
-                let (assigns, decls) = trav assigns t in
-                (assigns, (h :: decls))
+            in
+            flatten_blocks (trav [] body)
         in
-        let (assigns, decls) = trav [] decls in
-        (match assigns with
-            | [] -> Program (decls, loc)
-            | assigns ->
-                let init_func = FunDef (true, Void, "__init", [], Block assigns, noloc) in
-                Program (init_func :: decls, loc)
-            )
+        let (body, new_vars) = for_to_while (Block (move_inits body)) in
+        let create_vardec name = VarDec (Int, name, None, noloc) in
+        let new_vardecs = List.map create_vardec !new_vars in
+        let stats = new_vardecs @ (flatten_blocks (block_body body)) in
+        FunDef (export, ret_type, name, params, Block stats, loc)
 
     | node -> transform_children var_init node