Mapping to/from AST v0

Author: cknitt

compiler/ml/ast_mapper_from0.ml

@@ -310,6 +310,20 @@ module E = struct
         | _ -> true)
       attrs
 
+  let extract_for_of_attribute attrs =
+    List.find_map
+      (function
+        | {Location.txt = "res.for_of"}, Pt.PPat (_, Some expr) -> Some expr
+        | _ -> None)
+      attrs
+
+  let remove_for_of_attribute attrs =
+    List.filter
+      (function
+        | {Location.txt = "res.for_of"}, _ -> false
+        | _ -> true)
+      attrs
+
   let map_jsx_children sub (e : expression) : Pt.jsx_children =
     let rec visit (e : expression) : Pt.expression list =
       match e.pexp_desc with
@@ -524,9 +538,17 @@ module E = struct
       sequence ~loc ~attrs (sub.expr sub e1) (sub.expr sub e2)
     | Pexp_while (e1, e2) ->
       while_ ~loc ~attrs (sub.expr sub e1) (sub.expr sub e2)
-    | Pexp_for (p, e1, e2, d, e3) ->
-      for_ ~loc ~attrs (sub.pat sub p) (sub.expr sub e1) (sub.expr sub e2) d
-        (sub.expr sub e3)
+    | Pexp_for (p, e1, e2, d, e3) -> (
+      let array_expr = extract_for_of_attribute attrs in
+      let attrs = remove_for_of_attribute attrs in
+      match array_expr with
+      | Some array ->
+        (* This is actually a for...of loop, decode it *)
+        for_of ~loc ~attrs (sub.pat sub p) array (sub.expr sub e3)
+      | None ->
+        (* Regular for loop *)
+        for_ ~loc ~attrs (sub.pat sub p) (sub.expr sub e1) (sub.expr sub e2) d
+          (sub.expr sub e3))
     | Pexp_coerce (e, (), t2) ->
       coerce ~loc ~attrs (sub.expr sub e) (sub.typ sub t2)
     | Pexp_constraint (e, t) ->

compiler/ml/ast_mapper_to0.ml

@@ -450,9 +450,25 @@ module E = struct
     | Pexp_for (p, e1, e2, d, e3) ->
       for_ ~loc ~attrs (sub.pat sub p) (sub.expr sub e1) (sub.expr sub e2) d
         (sub.expr sub e3)
-    | Pexp_for_of (_pat, _array_expr, _body_expr) ->
-      (* Convert for...of to an extension since it doesn't exist in old AST *)
-      extension ~loc ~attrs (Location.mkloc "rescript.for_of" loc, PStr [])
+    | Pexp_for_of (pat, array_expr, body_expr) ->
+      (* Encode for...of as a for loop with attributes *)
+      let for_of_attr =
+        sub.attribute sub
+          ( Location.mkloc "res.for_of" loc,
+            PPat (Ast_helper.Pat.any (), Some array_expr) )
+      in
+      (* Use dummy values for start/end expressions since we store the array in attribute *)
+      let start_expr =
+        sub.expr sub
+          (Ast_helper.Exp.constant (Parsetree.Pconst_integer ("0", None)))
+      in
+      let end_expr =
+        sub.expr sub
+          (Ast_helper.Exp.constant (Parsetree.Pconst_integer ("0", None)))
+      in
+      (* Use Upto direction flag (arbitrary choice) *)
+      for_ ~loc ~attrs:(for_of_attr :: attrs) (sub.pat sub pat) start_expr
+        end_expr Asttypes.Upto (sub.expr sub body_expr)
     | Pexp_coerce (e, (), t2) ->
       coerce ~loc ~attrs (sub.expr sub e) (sub.typ sub t2)
     | Pexp_constraint (e, t) ->

tests/syntax_tests/data/ast-mapping/ForOfExpressions.res

@@ -0,0 +1,15 @@
+// Test for..of AST mapping
+
+let testForOf = () => {
+  let arr = [1, 2, 3]
+
+  // Basic for..of
+  for x of arr {
+    Console.log(x)
+  }
+
+  // For..of with complex expression
+  for item of Array.map(arr, x => x + 1) {
+    Console.log(item)
+  }
+}

tests/syntax_tests/data/ast-mapping/expected/ForOfExpressions.res.txt

@@ -0,0 +1,15 @@
+// Test for..of AST mapping
+
+let testForOf = () => {
+  let arr = [1, 2, 3]
+
+  // Basic for..of
+  for x of arr {
+    Console.log(x)
+  }
+
+  // For..of with complex expression
+  for item of Array.map(arr, x => x + 1) {
+    Console.log(item)
+  }
+}