Accept private iterables in for..of

changes/for-of-private-iterables.feature.md

@@ -0,0 +1 @@
+- `for..of` and `for await..of` accept generators and other iterator types directly, without `asIterable`.

compiler/ml/typecore.ml

@@ -418,6 +418,27 @@ let unify_exp_types ~context loc env ty expected_ty =
   with Unify trace ->
     raise (Error (loc, env, Expr_type_clash {trace; context}))
 
+(* Like [unify_exp_types], but a private abbreviation such as [Generator.t]
+   also unifies through its expansion, as an explicit [:>] would. A failure
+   reports the type as written. *)
+let unify_exp_types_through_private ~context loc env ty expected_ty =
+  let snapshot = Btype.snapshot () and levels = Ctype.save_levels () in
+  let restore () =
+    Btype.backtrack snapshot;
+    Ctype.set_levels levels
+  in
+  try unify_exp_types ~context loc env ty expected_ty
+  with Error (_, _, Expr_type_clash _) as error -> (
+    restore ();
+    let expanded = Ctype.expand_head_opt env ty in
+    if Btype.repr expanded == Btype.repr (Ctype.expand_head env ty) then
+      raise error
+    else
+      try unify_exp_types ~context loc env expanded expected_ty
+      with Error (_, _, Expr_type_clash _) ->
+        restore ();
+        raise error)
+
 (* level at which to create the local type declarations *)
 let get_newtype_level () =
   match !(newtype_level_ref ()) with
@@ -3321,7 +3342,7 @@ and type_expect_ ?deprecated_context ~context ?(recarg = Rejected) env sexp
       let collection = type_exp ~context:None env scollection in
       let iterable_state = (Btype.snapshot (), Ctype.save_levels ()) in
       try
-        unify_exp_types ~context:None scollection.pexp_loc env
+        unify_exp_types_through_private ~context:None scollection.pexp_loc env
           collection.exp_type
           (Predef.type_iterable ty_elem);
         collection
@@ -3379,8 +3400,11 @@ and type_expect_ ?deprecated_context ~context ?(recarg = Rejected) env sexp
   | Pexp_for_await_of (param, scollection, sbody) ->
     let ty_elem = newvar () in
     let collection =
-      type_expect ~context:None env scollection
-        (Predef.type_async_iterable ty_elem)
+      let collection = type_exp ~context:None env scollection in
+      unify_exp_types_through_private ~context:None scollection.pexp_loc env
+        collection.exp_type
+        (Predef.type_async_iterable ty_elem);
+      collection
     in
     let id, new_env =
       match param.ppat_desc with

tests/tests/src/for_of_private_iterable_test.mjs

@@ -0,0 +1,69 @@
+// Generated by ReScript, PLEASE EDIT WITH CARE
+
+import * as Mocha from "mocha";
+import * as Test_utils from "./test_utils.mjs";
+
+let numbers = (function* () {
+  yield 1
+  yield 2
+  yield 3
+  return "done"
+});
+
+let words = (async function* () {
+  yield "a"
+  yield "b"
+});
+
+Mocha.describe("For_of_private_iterable_test", () => {
+  Mocha.test("for..of iterates a generator", () => {
+    let seen = [];
+    for (let x of numbers()) {
+      seen.push(x * 10 | 0);
+    }
+    Test_utils.eq("File \"for_of_private_iterable_test.res\", line 25, characters 7-14", [
+      10,
+      20,
+      30
+    ], seen);
+  });
+  Mocha.test("break in for..of over a generator", () => {
+    let seen = [];
+    for (let x of numbers()) {
+      if (x === 2) {
+        break;
+      }
+      seen.push(x);
+    }
+    Test_utils.eq("File \"for_of_private_iterable_test.res\", line 36, characters 7-14", [1], seen);
+  });
+  Mocha.test("for..of iterates an iterator object", () => {
+    let seen = [];
+    for (let x of [
+        4,
+        5
+      ].values()) {
+      seen.push(x + 1 | 0);
+    }
+    Test_utils.eq("File \"for_of_private_iterable_test.res\", line 44, characters 7-14", [
+      5,
+      6
+    ], seen);
+  });
+  Mocha.test("for await..of iterates an async generator", async () => {
+    let seen = [];
+    for await (let s of words()) {
+      seen.push(s + "!");
+    }
+    return Test_utils.eq("File \"for_of_private_iterable_test.res\", line 52, characters 7-14", [
+      "a!",
+      "b!"
+    ], seen);
+  });
+});
+
+export {
+  numbers,
+  words,
+}
+/* numbers Not a pure module */

tests/tests/src/for_of_private_iterable_test.res

@@ -0,0 +1,54 @@
+open Mocha
+open Test_utils
+
+@module("mocha")
+external testAsync: (string, unit => promise<unit>) => unit = "test"
+
+let numbers: unit => Generator.t<int, string, unit> = %raw(`function* () {
+  yield 1
+  yield 2
+  yield 3
+  return "done"
+}`)
+
+let words: unit => AsyncGenerator.t<string, unit, unit> = %raw(`async function* () {
+  yield "a"
+  yield "b"
+}`)
+
+describe(__MODULE__, () => {
+  test("for..of iterates a generator", () => {
+    let seen = []
+    for x of numbers() {
+      seen->Array.push(x * 10)
+    }
+    eq(__LOC__, [10, 20, 30], seen)
+  })
+
+  test("break in for..of over a generator", () => {
+    let seen = []
+    for x of numbers() {
+      if x == 2 {
+        break
+      }
+      seen->Array.push(x)
+    }
+    eq(__LOC__, [1], seen)
+  })
+
+  test("for..of iterates an iterator object", () => {
+    let seen = []
+    for x of [4, 5]->Array.values {
+      seen->Array.push(x + 1)
+    }
+    eq(__LOC__, [5, 6], seen)
+  })
+
+  testAsync("for await..of iterates an async generator", async () => {
+    let seen = []
+    for await s of words() {
+      seen->Array.push(s ++ "!")
+    }
+    eq(__LOC__, ["a!", "b!"], seen)
+  })
+})