Reject await in a function that an async function returns directly (#9169)

Why:
    The built-in PPX let a function in the body of an async function
    inherit its async context, so `await` in a plain arrow returned by an
    async function passed the check, and the JavaScript had `await` in a
    function that is not async (#9168). The inheritance comes from the
    curried parsetree, where `async (x, y) => ...` was nested functions;
    in the n-ary parsetree a function in a function's body is a function
    of its own, so each function is now in an async context exactly when
    it is async.

Diagnostic added (await in a returned function):
Example:
    let f = async x => y => await y + x
Before (js): a SyntaxError when the module loads
    async function f(x) {
      return y => await y + x | 0;
    }
After (diagnostic):
    A.res:1:25-31
    Await on expression not in an async context

Review: Does a working program stop compiling? Only one whose `await` in
    such a function was removed by the optimizer, as in an unused
    definition or a branch under a constant `false`; the same code in any
    other plain function is already rejected. Programs that still compile
    give the same JavaScript, since only the await check reads the async
    context.

Signed-off-by: mununki <woonki.moon@gmail.com>

Author: mununki

changes/await-in-returned-function.fix.md

@@ -0,0 +1 @@
+- Reject `await` in a function that an async function returns directly, as in `async x => y => await y`, which produced JavaScript with `await` outside an async function.

compiler/frontend/bs_builtin_ppx.ml

@@ -113,10 +113,7 @@ let local_module_type_decl ~await_context
           (Mtd.mk ~loc {txt = name; loc} ~typ:(Mty.typeof_ ~loc me)));
     ])
 
-let expr_mapper ~async_context ~in_function_def (self : mapper)
-    (e : Parsetree.expression) =
-  let old_in_function_def = !in_function_def in
-  in_function_def := false;
+let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
   match e.pexp_desc with
   (* Its output should not be rewritten anymore *)
   | Pexp_regexp {pattern; flags} ->
@@ -168,8 +165,10 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
   | Pexp_fun {newtypes; params; body; async} -> (
     match Ast_attributes.process_attributes_rev e.pexp_attributes with
     | Nothing, _ ->
-      (* Handle @async x => y => ... is in async context *)
-      async_context := (old_in_function_def && !async_context) || async;
+      (* A [Pexp_fun] holds all parameters of one function, so a [Pexp_fun]
+         in its body is another function: in [async x => y => await y], the
+         inner function is not async and compiles to a plain arrow. *)
+      async_context := async;
       let attrs = self.attributes self e.pexp_attributes in
       let params =
         Ext_list.map params (fun (p : Parsetree.fun_param) ->
@@ -179,13 +178,7 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
               p_pat = self.pat self p.p_pat;
             })
       in
-      let saved_in_function_def = !in_function_def in
-      in_function_def := true;
-      (* Keep reporting the body as part of a function definition so nested
-         functions propagate async context exactly like the legacy curried
-         mapper did (GH #7974). *)
       let body = self.expr self body in
-      in_function_def := saved_in_function_def;
       let newtypes =
         Ext_list.map newtypes (fun (name, nt_attrs) ->
             (name, self.attributes self nt_attrs))
@@ -432,10 +425,9 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
     | Some (_, _, Some _) | None -> default_expr_mapper self e)
   | _ -> default_expr_mapper self e
 
-let expr_mapper ~async_context ~in_function_def (self : mapper)
-    (e : Parsetree.expression) =
+let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
   let async_saved = !async_context in
-  let result = expr_mapper ~async_context ~in_function_def self e in
+  let result = expr_mapper ~async_context self e in
   async_context := async_saved;
   let check_await () =
     if !async_context = false then
@@ -690,7 +682,6 @@ let structure_item_mapper (self : mapper) (str : Parsetree.structure_item) :
 
 type mapper_context = {
   async_context: bool ref;
-  in_function_def: bool ref;
   await_context: (string, string) Hashtbl.t ref;
 }
 
@@ -699,11 +690,7 @@ type mapper_context = {
    function body must not change whether an await is accepted here. *)
 let mapper_context_key =
   Request_local.new_key (fun () ->
-      {
-        async_context = ref true;
-        in_function_def = ref false;
-        await_context = ref (Hashtbl.create 10);
-      })
+      {async_context = ref true; await_context = ref (Hashtbl.create 10)})
 
 let mapper_context () = Domain.DLS.get mapper_context_key
 
@@ -897,8 +884,7 @@ let mapper : mapper =
     expr =
       (fun self expression ->
         let context = mapper_context () in
-        expr_mapper ~async_context:context.async_context
-          ~in_function_def:context.in_function_def self expression);
+        expr_mapper ~async_context:context.async_context self expression);
     pat = pat_mapper;
     typ = typ_mapper;
     signature_item = signature_item_mapper;

tests/build_tests/super_errors/expected/await_in_returned_function.res.expected

@@ -0,0 +1,10 @@
+
+  We've found a bug for you!
+  /.../fixtures/await_in_returned_function.res:3:28-36
+
+  1 │ let a = async () => 3
+  2 │ 
+  3 │ let foo = async x => y => (await a()) + x + y
+  4 │ 
+
+  Await on expression not in an async context

tests/build_tests/super_errors/fixtures/await_in_returned_function.res

@@ -0,0 +1,3 @@
+let a = async () => 3
+
+let foo = async x => y => (await a()) + x + y

tests/tests/src/async_returned_function_test.mjs

@@ -0,0 +1,38 @@
+// Generated by ReScript, PLEASE EDIT WITH CARE
+
+import * as Mocha from "mocha";
+import * as Test_utils from "./test_utils.mjs";
+
+async function addAsync(x) {
+  return async y => await y + x | 0;
+}
+
+async function addPlain(x) {
+  return y => y + x | 0;
+}
+
+async function addThree(x) {
+  return y => (async z => (await y + await z | 0) + x | 0);
+}
+
+Mocha.describe("Async_returned_function_test", () => {
+  Mocha.test("async function returning an async function", async () => {
+    let add = await addAsync(1);
+    return Test_utils.eq("File \"async_returned_function_test.res\", line 18, characters 7-14", 3, await add(Promise.resolve(2)));
+  });
+  Mocha.test("async function returning a plain function", async () => {
+    let add = await addPlain(1);
+    return Test_utils.eq("File \"async_returned_function_test.res\", line 23, characters 7-14", 3, add(2));
+  });
+  Mocha.test("plain function between two async functions", async () => {
+    let add = await addThree(1);
+    return Test_utils.eq("File \"async_returned_function_test.res\", line 28, characters 7-14", 6, await add(Promise.resolve(2))(Promise.resolve(3)));
+  });
+});
+
+export {
+  addAsync,
+  addPlain,
+  addThree,
+}
+/*  Not a pure module */

tests/tests/src/async_returned_function_test.res

@@ -0,0 +1,30 @@
+open Mocha
+open Test_utils
+
+@module("mocha")
+external testAsync: (string, unit => promise<unit>) => unit = "test"
+
+// Each arrow is a function of its own: the inner arrows below must compile to
+// JS functions that are async exactly when they are async in the source.
+let addAsync = async x => async y => (await y) + x
+
+let addPlain = async x => y => y + x
+
+let addThree = async x => y => async z => (await y) + await z + x
+
+describe(__MODULE__, () => {
+  testAsync("async function returning an async function", async () => {
+    let add = await addAsync(1)
+    eq(__LOC__, 3, await add(Promise.resolve(2)))
+  })
+
+  testAsync("async function returning a plain function", async () => {
+    let add = await addPlain(1)
+    eq(__LOC__, 3, add(2))
+  })
+
+  testAsync("plain function between two async functions", async () => {
+    let add = await addThree(1)
+    eq(__LOC__, 6, await add(Promise.resolve(2))(Promise.resolve(3)))
+  })
+})