Lift restriction to uncurried functions.

Use the existing convention in the printer to determine a conventional meaning to the result.
`x => y => 3` is formatted as `(x,y) => 3`
so `@async x => y => 3` adds a promise to the result `3`.

Author: cristianoc

example-async/lib/js/src/AA.js

@@ -173,6 +173,14 @@ let rec runAllTests =
 
 runAllTests(. 0)->ignore
 
+//
+//
+// Curried functions
+
+let bb = @async x => @await x
+
+let cc = @async (x, ~y=x, z) => @await x + @await y + @await z
+
 //
 //
 // Errors
@@ -184,4 +192,4 @@ runAllTests(. 0)->ignore
 //     cb
 //   }
 
-// let bb = @async x => x // Error: Async can only be applied to uncurried function
+// let _ = @async (_, . x) => @await x // Error: Await on expression not in an async context

example-async/src/AA.res

@@ -19,3 +19,17 @@ let add_async_attribute ~async (body : Parsetree.expression) =
         :: body.pexp_attributes;
     }
   else body
+
+let rec add_promise_to_result (e : Parsetree.expression) =
+  match e.pexp_desc with
+  | Pexp_fun (label, eo, pat, body) ->
+      let body = add_promise_to_result body in
+      { e with pexp_desc = Pexp_fun (label, eo, pat, body) }
+  | _ -> add_promise_type ~async:true e
+
+let make_function_async ~async (e : Parsetree.expression) =
+  if async then
+    match e.pexp_desc with
+    | Pexp_fun _ -> add_async_attribute ~async (add_promise_to_result e)
+    | _ -> assert false
+  else e

jscomp/frontend/ast_async.ml

@@ -76,7 +76,9 @@ let pat_mapper (self : mapper) (e : Parsetree.pattern) =
   | Ppat_constant (Pconst_integer (s, Some 'l')) ->
     {e with ppat_desc = Ppat_constant (Pconst_integer(s,None))}
   | _ -> default_pat_mapper  self e 
-let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
+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;
   match e.pexp_desc with
   (* Its output should not be rewritten anymore *)
   | Pexp_extension extension ->
@@ -122,13 +124,13 @@ let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
 
   | Pexp_fun (label, _, pat , body)
     ->
-    async_context := false;
     let async = Ast_attributes.has_async_payload e.pexp_attributes <> None in
     begin match Ast_attributes.process_attributes_rev e.pexp_attributes with
       | Nothing, _ ->
-        if async then
-          Location.raise_errorf ~loc:e.pexp_loc "Async can only be applied to uncurried function";
-        default_expr_mapper self e
+        (* Handle @async x => y => ... is in async context *)
+        async_context := (old_in_function_def && !async_context) || async;
+        in_function_def := true;
+        Ast_async.make_function_async ~async (default_expr_mapper self e)
       | Uncurry _, pexp_attributes
         ->
         async_context := async;
@@ -140,6 +142,7 @@ let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
       | Meth_callback _, pexp_attributes
         ->
         (* FIXME: does it make sense to have a label for [this] ? *)
+        async_context := false;
         {e with pexp_desc = Ast_uncurry_gen.to_method_callback e.pexp_loc  self label pat body ;
                 pexp_attributes }
     end
@@ -200,9 +203,9 @@ let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
   *)
   | _ ->  default_expr_mapper self e
 
-let expr_mapper ~async_context (self : mapper) (e : Parsetree.expression) =
+let expr_mapper ~async_context ~in_function_def (self : mapper) (e : Parsetree.expression) =
   let async_saved = !async_context in
-  let result = expr_mapper ~async_context self e in
+  let result = expr_mapper ~async_context ~in_function_def self e in
   async_context := async_saved;
   match Ast_attributes.has_await_payload e.pexp_attributes with
   | None -> result
@@ -472,7 +475,7 @@ let rec
 
 let  mapper : mapper =
   { default_mapper with
-    expr = expr_mapper ~async_context:(ref false);
+    expr = expr_mapper ~async_context:(ref false) ~in_function_def:(ref false);
     pat = pat_mapper; 
     typ = typ_mapper ;
     class_type = class_type_mapper;