improve n-ary continuation coercion

src/codegen/compile.ml

@@ -10590,8 +10590,7 @@ and compile_prim_invocation (env : E.t) ae p es at =
     let call_as_prim = match fun_sr, sort with
       | SR.Const (_, Const.Fun (mk_fi, Const.PrimWrapper prim)), _ ->
          begin match n_args, e2.it with
-         | 0, _ -> true
-         | 1, _ -> true
+         | (0 | 1), _ -> true
          | n, PrimE (TupPrim, es) when List.length es = n -> true
          | _, _ -> false
          end

src/ir_passes/async.ml

@@ -71,6 +71,23 @@ let new_async t =
   let fail = fresh_var "fail" (typ (projE call_new_async 2)) in
   (async, fulfill, fail), call_new_async
 
+let coerce_and_cont0T =
+  T.Func (
+    T.Local,
+    T.Returns,
+    [],
+    [t_async_fut unary T.unit],
+    [t_async_fut nary T.unit])
+
+let coerce_and_cont2T =
+  let t = T.(Tup [Var ("T", 0); Var ("U", 1)]) in
+  T.Func (
+    T.Local,
+    T.Returns,
+    [ { var = "T"; sort = T.Type; bound = T.Any }; { var = "U"; sort = T.Type; bound = T.Any } ],
+    [t_async_fut unary t],
+    [t_async_fut nary t])
+
 let new_nary_async_reply ts =
   (* The async implementation isn't n-ary *)
   let t = T.seq ts in
@@ -91,16 +108,22 @@ let new_nary_async_reply ts =
     in
     match ts with
     | [t1] ->
-      begin
-      match T.normalize t1 with
-      | T.Tup _ ->
-        (* TODO(#3740): find a better fix than PR #3741 *)
-        (* HACK *)
-        coerce t1
+      begin match T.normalize t1 with
+      | T.Tup [] ->
+        varE (var "@coerce_and_cont00" coerce_and_cont0T) -*- varE unary_async
+      | T.Tup ([_; _] as tu) ->
+        callE (varE (var "@coerce_and_cont2" coerce_and_cont2T)) tu (varE unary_async)
+      (* TODO(#3740): find a better fix than PR #3741 *)
+      (* HACK *)
+      | T.Tup _ -> coerce t1
       | _ ->
         varE unary_async
       end
-    | ts1 ->
+    | [] ->
+      varE (var "@coerce_and_cont0" coerce_and_cont0T) -*- varE unary_async
+    | [_; _] ->
+      callE (varE (var "@coerce_and_cont2" coerce_and_cont2T)) ts (varE unary_async)
+    | _ ->
       coerce t
   in
   (* construct the n-ary reply callback that take a *sequence* of values to fulfill the async *)

src/ir_passes/await.ml

@@ -594,7 +594,7 @@ and t_comp_unit context = function
         let e = fresh_var "e" T.catch in
         ProgU [
           funcD throw e (assertE (falseE ()));
-          expD (c_block context' ds (tupE []) (meta (T.unit) (fun v1 -> tupE [])))
+          expD (c_block context' ds (tupE []) (meta T.unit (fun v1 -> tupE [])))
         ]
     end
   | ActorU (as_opt, ds, ids, { meta = m; preupgrade; postupgrade; heartbeat; timer; inspect}, t) ->
@@ -619,7 +619,7 @@ and t_ignore_throw context exp =
      { (blockE [
           funcD throw e (tupE[]);
         ]
-        (c_exp context' exp (meta (T.unit) (fun v1 -> tupE []))))
+        (c_exp context' exp (meta T.unit (fun v1 -> tupE []))))
        (* timer logic requires us to preserve any source location,
           or timer won't be initialized in compile.ml *)
        with at = exp.at

src/prelude/internals.mo

@@ -307,6 +307,18 @@ func @getSystemRefund() : @Refund {
   return (prim "cyclesRefunded" : () -> Nat) ();
 };
 
+func @coerce_and_cont2<T <: Any, U <: Any>(a : @Async<(T, U)>) :
+  (k : (T, U) -> (), r : @Cont<Error>) -> {
+    #suspend;
+    #schedule : () -> ()
+  } = func(k, r) = a(func tup = k tup, r);
+
+func @coerce_and_cont0(a : @Async<()>) :
+  (k : () -> (), r : @Cont<Error>) -> {
+    #suspend;
+    #schedule : () -> ()
+  } = func(k, r) = a(func() = k(), r);
+
 func @new_async<T <: Any>() : (@Async<T>, @Cont<T>, @Cont<Error>) {
   let w_null = func(r : @Refund, t : T) { };
   let r_null = func(_ : Error) {};

test/run-drun/await.mo

@@ -66,6 +66,9 @@ actor a {
        Prim.debugPrint a;
        Prim.debugPrint b;
     };
+    Prim.debugPrint (debug_show (await p()));
+    Prim.debugPrint (debug_show (await async (1, "two", '3')));
+    await async ();
 
     ignore(await a);
     ignore(await b);

test/run-drun/ok/await.drun-run.ok

@@ -18,10 +18,12 @@ debug.print: cnt: 2 i: 2
 debug.print: cnt: 3 i: 4
 debug.print: cnt: 4 i: 5
 debug.print: cnt: 5 i: 10
+debug.print: ("fst", "snd")
 debug.print: .
 debug.print: cnt: 6 i: 3
 debug.print: cnt: 7 i: 6
 debug.print: cnt: 8 i: 11
+debug.print: (1, "two", '3')
 debug.print: .
 debug.print: cnt: 9 i: 7
 debug.print: cnt: 10 i: 12

test/run-drun/ok/await.run-ir.ok

@@ -16,10 +16,12 @@ cnt: 2 i: 2
 cnt: 3 i: 4
 cnt: 4 i: 5
 cnt: 5 i: 10
+("fst", "snd")
 .
 cnt: 6 i: 3
 cnt: 7 i: 6
 cnt: 8 i: 11
+(1, "two", '3')
 .
 cnt: 9 i: 7
 cnt: 10 i: 12

test/run-drun/ok/await.run-low.ok

@@ -16,10 +16,12 @@ cnt: 2 i: 2
 cnt: 3 i: 4
 cnt: 4 i: 5
 cnt: 5 i: 10
+("fst", "snd")
 .
 cnt: 6 i: 3
 cnt: 7 i: 6
 cnt: 8 i: 11
+(1, "two", '3')
 .
 cnt: 9 i: 7
 cnt: 10 i: 12

test/run-drun/ok/await.run.ok

@@ -16,10 +16,12 @@ cnt: 2 i: 2
 cnt: 3 i: 4
 cnt: 4 i: 5
 cnt: 5 i: 10
+("fst", "snd")
 .
 cnt: 6 i: 3
 cnt: 7 i: 6
 cnt: 8 i: 11
+(1, "two", '3')
 .
 cnt: 9 i: 7
 cnt: 10 i: 12

test/run-drun/ok/await.tc.ok

@@ -1,3 +1,2 @@
 await.mo:45.9-45.10: warning [M0194], unused identifier g (delete or rename to wildcard `_` or `_g`)
-await.mo:63.10-63.11: warning [M0194], unused identifier p (delete or rename to wildcard `_` or `_p`)
 await.mo:64.9-64.10: warning [M0194], unused identifier h (delete or rename to wildcard `_` or `_h`)