feat: introduce await? to short-circuit fulfilled async futures

src/ir_def/arrange_ir.ml

@@ -80,8 +80,8 @@ and prim = function
   | GetLastArrayOffset -> Atom "GetLastArrayOffset"
   | BreakPrim i       -> "BreakPrim"  $$ [id i]
   | RetPrim           -> Atom "RetPrim"
-  | AwaitPrim Type.Fut -> Atom "AwaitPrim"
-  | AwaitPrim Type.Cmp -> Atom "AwaitPrim*"
+  | AwaitPrim (AwaitFut _) -> Atom "AwaitPrim"
+  | AwaitPrim AwaitCmp -> Atom "AwaitPrim*"
   | AssertPrim        -> Atom "AssertPrim"
   | ThrowPrim         -> Atom "ThrowPrim"
   | ShowPrim t        -> "ShowPrim" $$ [typ t]
@@ -108,8 +108,8 @@ and prim = function
   | SetCertifiedData  -> Atom "SetCertifiedData"
   | GetCertificate    -> Atom "GetCertificate"
   | OtherPrim s       -> Atom s
-  | CPSAwait (Type.Fut, t) -> "CPSAwait" $$ [typ t]
-  | CPSAwait (Type.Cmp, t) -> "CPSAwait*" $$ [typ t]
+  | CPSAwait (AwaitFut _, t) -> "CPSAwait" $$ [typ t]
+  | CPSAwait (AwaitCmp, t) -> "CPSAwait*" $$ [typ t]
   | CPSAsync (Type.Fut, t) -> "CPSAsync" $$ [typ t]
   | CPSAsync (Type.Cmp, t) -> "CPSAsync*" $$ [typ t]
   | ICArgDataPrim     -> Atom "ICArgDataPrim"

src/ir_def/check_ir.ml

@@ -538,7 +538,7 @@ let rec check_exp env (exp:Ir.exp) : unit =
       | Some c -> T.Con(c, [])
       | None -> error env exp.at "misplaced await" in
       let t1 = T.promote (typ exp1) in
-      let (t2, t3) = try T.as_async_sub s t0 t1
+      let (t2, t3) = try T.as_async_sub (to_async_sort s) t0 t1
              with Invalid_argument _ ->
                error env exp1.at "expected async type, but expression has type\n  %s"
                  (T.string_of_typ_expand t1)
@@ -566,8 +566,8 @@ let rec check_exp env (exp:Ir.exp) : unit =
       typ exp1 <: T.blob;
       T.Opt (T.seq ots) <: t
     | CPSAwait (s, cont_typ), [a; krb] ->
-      let (_, t1) =
-        try T.as_async_sub s T.Non (T.normalize (typ a))
+       let (_, t1) = (*assert (s = AwaitFut false);*)
+        try T.as_async_sub (to_async_sort s) T.Non (T.normalize (typ a))
         with _ -> error env exp.at "CPSAwait expect async arg, found %s" (T.string_of_typ (typ a))
       in
       (match cont_typ with

src/ir_def/construct.ml

@@ -162,8 +162,9 @@ let asyncE s typ_bind e typ1 =
                       eff = T.(if s = Fut then Await else Triv) }
   }
 
-let awaitE s e =
+let awaitE e =
   let (s, _ , typ) = T.as_async (T.normalize (typ e)) in
+  let s = match s with T.Cmp -> AwaitCmp | T.Fut -> AwaitFut false in
   { it = PrimE (AwaitPrim s, [e]);
     at = no_region;
     note = Note.{ def with typ; eff = T.Await }

src/ir_def/construct.mli

@@ -50,9 +50,9 @@ val primE : Ir.prim -> exp list -> exp
 val selfRefE : typ -> exp
 val assertE : exp -> exp
 val asyncE : async_sort -> typ_bind -> exp -> typ -> exp
-val awaitE : async_sort -> exp -> exp
+val awaitE : exp -> exp
 val cps_asyncE : async_sort -> typ -> typ -> exp -> exp
-val cps_awaitE : async_sort -> typ -> exp -> exp -> exp
+val cps_awaitE : await_sort -> typ -> exp -> exp -> exp
 val ic_replyE : typ list -> exp -> exp
 val ic_rejectE : exp -> exp
 val ic_callE : exp -> exp -> exp -> exp -> exp -> exp

src/ir_def/ir.ml

@@ -134,7 +134,7 @@ and prim =
   | IdxBlobPrim                       (* blob indexing *)
   | BreakPrim of id                   (* break *)
   | RetPrim                           (* return *)
-  | AwaitPrim of Type.async_sort       (* await/await* *)
+  | AwaitPrim of await_sort           (* await/await* *)
   | AssertPrim                        (* assertion *)
   | ThrowPrim                         (* throw *)
   | ShowPrim of Type.typ              (* debug_show *)
@@ -170,7 +170,7 @@ and prim =
 
   | OtherPrim of string               (* Other primitive operation, no custom typing rule *)
   (* backend stuff *)
-  | CPSAwait of Type.async_sort * Type.typ
+  | CPSAwait of await_sort * Type.typ
                                       (* typ is the current continuation type of cps translation *)
   | CPSAsync of Type.async_sort * Type.typ
   | ICPerformGC
@@ -186,6 +186,8 @@ and prim =
   | ICStableRead of Type.typ           (* deserialize value of stable type from stable memory *)
   | ICStableSize of Type.typ
 
+and await_sort = AwaitFut of bool | AwaitCmp
+
 (* Declarations *)
 
 and dec = dec' Source.phrase
@@ -335,3 +337,7 @@ let map_prim t_typ t_id p =
   | ICStableWrite t -> ICStableWrite (t_typ t)
   | ICStableRead t -> ICStableRead (t_typ t)
   | ICStableSize t -> ICStableSize (t_typ t)
+
+let to_async_sort = function
+  | AwaitCmp -> Type.Cmp
+  | AwaitFut _ -> Type.Fut

src/ir_interpreter/interpret_ir.ml

@@ -369,10 +369,10 @@ and interpret_exp_mut env exp (k : V.value V.cont) =
       | BreakPrim id, [v1] -> find id env.labs v1
       | RetPrim, [v1] -> Option.get env.rets v1
       | ThrowPrim, [v1] -> Option.get env.throws v1
-      | AwaitPrim Type.Fut, [v1] ->
+      | AwaitPrim (AwaitFut _), [v1] ->
         assert env.flavor.has_await;
         await env exp.at (V.as_async v1) k (Option.get env.throws)
-      | AwaitPrim Type.Cmp, [v1] ->
+      | AwaitPrim AwaitCmp, [v1] ->
         assert env.flavor.has_await;
         (V.as_comp v1) k (Option.get env.throws)
       | AssertPrim, [v1] ->

src/ir_passes/async.ml

@@ -256,7 +256,7 @@ let transform prog =
     | VarE (_, _) -> exp'
     | AssignE (exp1, exp2) ->
       AssignE (t_lexp exp1, t_exp exp2)
-    | PrimE (CPSAwait (Fut, cont_typ), [a; krb]) ->
+    | PrimE (CPSAwait (AwaitFut short, cont_typ), [a; krb]) ->
       begin match cont_typ with
         | Func(_, _, [], _, []) ->
           (* unit answer type, from await in `async {}` *)
@@ -265,19 +265,22 @@ let transform prog =
             switch_variantE (t_exp a -*- varE vkrb)
               [ ("suspend", wildP,
                   unitE()); (* suspend *)
-                ("schedule", varP schedule, (* resume later *)
-                  (* try await async (); schedule() catch e -> r(e) *)
-                 (let v = fresh_var "call" unit in
-                  letE v
-                    (selfcallE [] (ic_replyE [] (unitE())) (varE schedule) (projE (varE vkrb) 1)
-                       ([] -->* (projE (varE vkrb) 2 -*- unitE ())))
-                    (check_call_perform_status (varE v) (fun e -> projE (varE vkrb) 1 -*- e))))
+                ("schedule", varP schedule,
+                  (if short then (* resume immediately *)
+                     varE schedule -*- unitE ()
+                   else (* resume later *)
+                     (* try await async (); schedule() catch e -> r(e) *)
+                     let v = fresh_var "call" unit in
+                     letE v
+                       (selfcallE [] (ic_replyE [] (unitE())) (varE schedule) (projE (varE vkrb) 1)
+                          ([] -->* (projE (varE vkrb) 2 -*- unitE ())))
+                       (check_call_perform_status (varE v) (fun e -> projE (varE vkrb) 1 -*- e))))
               ]
               unit
           )).it
         | _ -> assert false
       end
-    | PrimE (CPSAwait (Cmp, cont_typ), [a; krb]) ->
+    | PrimE (CPSAwait (AwaitCmp, cont_typ), [a; krb]) ->
       begin match cont_typ with
       | Func(_, _, [], _, []) ->
          (t_exp a -*- t_exp krb).it
@@ -296,7 +299,7 @@ let transform prog =
           letP (tupP [varP nary_async; varP nary_reply; varP reject; varP clean]) def;
           let ic_reply = (* flatten v, here and below? *)
             let v = fresh_var "v" (T.seq ts1) in
-            v --> (ic_replyE ts1 (varE v)) in
+            v --> ic_replyE ts1 (varE v) in
           let ic_reject =
             let e = fresh_var "e" catch in
             e --> ic_rejectE (errorMessageE (varE e)) in
@@ -382,7 +385,7 @@ let transform prog =
     | FuncE (x, s, c, typbinds, args, ret_tys, exp) ->
       begin
         match s with
-        | Local  ->
+        | Local ->
           FuncE (x, s, c, t_typ_binds typbinds, t_args args, List.map t_typ ret_tys, t_exp exp)
         | Shared s' ->
           begin
@@ -403,7 +406,7 @@ let transform prog =
                 | t -> assert false in
               let k =
                 let v = fresh_var "v" t1 in
-                v --> (ic_replyE ret_tys (varE v)) in
+                v --> ic_replyE ret_tys (varE v) in
               let r =
                 let e = fresh_var "e" catch in
                 e --> ic_rejectE (errorMessageE (varE e)) in

src/ir_passes/await.ml

@@ -462,12 +462,8 @@ and c_exp' context exp k =
     in
     k' -@- cps_async
   | PrimE (AwaitPrim s, [exp1]) ->
-    let r = match LabelEnv.find_opt Throw context with
-      | Some (Cont r) -> r
-      | _ -> assert false
-    in
-    let b = match LabelEnv.find_opt Cleanup context with
-      | Some (Cont r) -> r
+    let r, b = match LabelEnv.find_opt Throw context, LabelEnv.find_opt Cleanup context with
+      | Some (Cont r), Some (Cont b) -> r, b
       | _ -> assert false
     in
     letcont k (fun k ->
@@ -477,7 +473,7 @@ and c_exp' context exp k =
         cps_awaitE s (typ_of_var k) (t_exp context exp1) krb
       | T.Await ->
         c_exp context exp1
-          (meta (typ exp1) (fun v1 -> (cps_awaitE s (typ_of_var k) (varE v1) krb)))
+          (meta (typ exp1) (fun v1 -> cps_awaitE s (typ_of_var k) (varE v1) krb))
     )
   | DeclareE (id, typ, exp1) ->
     unary context k (fun v1 -> e (DeclareE (id, typ, varE v1))) exp1

src/lowering/desugar.ml

@@ -253,7 +253,9 @@ and exp' at note = function
                        | T.Async (_, t, _) -> t
                        | _ -> assert false) in
     (blockE ds { at; note; it }).it
-  | S.AwaitE (s, e) -> I.PrimE (I.AwaitPrim s, [exp e])
+  | S.AwaitE (T.Cmp, e) when T.is_fut e.note.S.note_typ -> I.PrimE I.(AwaitPrim (AwaitFut true), [exp e])
+  | S.AwaitE (T.Cmp, e) -> I.PrimE I.(AwaitPrim AwaitCmp, [exp e])
+  | S.AwaitE (T.Fut, e) -> I.PrimE I.(AwaitPrim (AwaitFut false), [exp e])
   | S.AssertE (Runtime, e) -> I.PrimE (I.AssertPrim, [exp e])
   | S.AssertE (_, e) -> (unitE ()).it
   | S.AnnotE (e, _) -> assert false
@@ -454,7 +456,7 @@ and call_system_func_opt name es obj_typ =
                 (primE (Ir.OtherPrim "trap")
                   [textE "canister_inspect_message explicitly refused message"]))
         | "lowmemory" ->
-          awaitE T.Cmp 
+          awaitE
             (callE (varE (var id.it note)) [T.scope_bound] (unitE()))
         | name ->
            let inst = match name with
@@ -1193,7 +1195,7 @@ let import_compiled_class (lib : S.comp_unit) wasm : import_declaration =
     (asyncE T.Fut
       (typ_arg c' T.Scope T.scope_bound)
       (letE principal
-        (awaitE T.Cmp
+        (awaitE
           (callE (varE install_actor_helper) cs'
             (tupE [
               install_arg;

src/mo_frontend/typing.ml

@@ -1738,7 +1738,10 @@ and infer_exp'' env exp : T.typ =
     let t0 = check_AwaitCap env "await" exp.at in
     let t1 = infer_exp_promote env exp1 in
     (try
-       let (t2, t3) = T.as_async_sub s t0 t1 in
+       let s1 = match s, t1 with
+         | T.(Cmp, Async(Fut, _, _)) -> T.Fut (* we can await* an async *)
+         | _ -> s in
+       let (t2, t3) = T.as_async_sub s1 t0 t1 in
        if not (eq env exp.at t0 t2) then begin
           local_error env exp1.at "M0087"
             "ill-scoped await: expected async type from current scope %s, found async type from other scope %s%s%s"
@@ -1759,7 +1762,7 @@ and infer_exp'' env exp : T.typ =
             (if s = T.Fut then
               "\nUse keyword 'await*' (not 'await') to consume this type."
             else
-              "\nUse keyword 'await' (not 'await*') to consume this type.")
+              assert false)
           else "")
     )
   | AssertE (_, exp1) ->

src/mo_types/type.ml

@@ -602,6 +602,7 @@ let is_unit = function Tup [] -> true | _ -> false
 let is_pair = function Tup [_; _] -> true | _ -> false
 let is_func = function Func _ -> true | _ -> false
 let is_async = function Async _ -> true | _ -> false
+let is_fut = function Async (Fut, _, _) -> true | _ -> false
 let is_mut = function Mut _ -> true | _ -> false
 let is_typ = function Typ _ -> true | _ -> false
 let is_con = function Con _ -> true | _ -> false

src/mo_types/type.mli

@@ -136,6 +136,7 @@ val is_unit : typ -> bool
 val is_pair : typ -> bool
 val is_func : typ -> bool
 val is_async : typ -> bool
+val is_fut : typ -> bool
 val is_mut : typ -> bool
 val is_typ : typ -> bool
 val is_con : typ -> bool

src/prelude/internals.mo

@@ -364,10 +364,6 @@ func @new_async<T <: Any>() : (@Async<T>, @Cont<T>, @Cont<Error>, @CleanCont) {
 
   var cleanup : @BailCont = @cleanup;
 
-  func clean() {
-      cleanup();
-  };
-
   func enqueue(k : @Cont<T>, r : @Cont<Error>, b : @BailCont) : {
     #suspend;
     #schedule : () -> ();
@@ -391,16 +387,16 @@ func @new_async<T <: Any>() : (@Async<T>, @Cont<T>, @Cont<Error>, @CleanCont) {
         };
         #suspend
       };
-      case (? (#ok (r, t))) {
-        #schedule (func () { @refund := r; k(t) });
+      case (?#ok (r, t)) {
+        #schedule (func() { @refund := r; k(t) })
       };
-      case (? (#error e)) {
-        #schedule (func () { r(e) });
+      case (?#error e) {
+        #schedule (func _ = r(e))
       };
     };
   };
 
-  (enqueue, fulfill, fail, clean)
+  (enqueue, fulfill, fail, func() = cleanup())
 };
 
 // Subset of IC management canister interface required for our use

test/fail/bad-async-sort.mo

@@ -1,6 +1,6 @@
 actor {
   func f() : async* Nat {
-    await* g(); // reject
+    await* g(); // accept
   };
 
   func g() : async Nat {
@@ -12,7 +12,7 @@ actor {
   };
 
   func anon2() : async () {
-    await* async {}; // reject
+    await* async {}; // accept
   };
 
   func anon3() : async () {

test/fail/ok/bad-async-sort.tc.ok

@@ -1,15 +1,9 @@
-bad-async-sort.mo:3.12-3.15: type error [M0088], expected async* type, but expression has type
-  async<$f> Nat
-Use keyword 'await' (not 'await*') to consume this type.
 bad-async-sort.mo:7.10-7.13: type error [M0088], expected async type, but expression has type
   async*<$g> Nat
 Use keyword 'await*' (not 'await') to consume this type.
 bad-async-sort.mo:11.11-11.20: type error [M0088], expected async type, but expression has type
   async*<$anon1> ()
 Use keyword 'await*' (not 'await') to consume this type.
-bad-async-sort.mo:15.12-15.20: type error [M0088], expected async* type, but expression has type
-  async<$anon2> ()
-Use keyword 'await' (not 'await*') to consume this type.
 bad-async-sort.mo:19.14-19.22: type error [M0183], async expression cannot produce expected async type 
   async*<$anon3> ().
 Use keyword 'async*' (not 'async') to produce the expected type.

test/run-drun/await.mo

@@ -62,7 +62,7 @@ actor a {
 
     func p():async (Text,Text) { ("fst","snd"); };
     let h = async {
-       let (a,b) = ("a","b"); /* await p(a,b);*/
+       let (a,b) = await p();
        Prim.debugPrint a;
        Prim.debugPrint b;
     };

test/run-drun/blockless.mo

@@ -0,0 +1,6 @@
+actor A {
+    public func foo() : async Int { 42 };
+    public func go() : async Int { await* A.foo() }
+};
+
+A.go() //OR-CALL ingress go 0x4449444C0000

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

@@ -10,15 +10,15 @@ debug.print: holy
 debug.print: e-while
 debug.print: g-label
 debug.print: .
-debug.print: a
-debug.print: b
 debug.print: cnt: 0 i: 0
 debug.print: cnt: 1 i: 1
 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: .
+debug.print: fst
+debug.print: snd
 debug.print: cnt: 6 i: 3
 debug.print: cnt: 7 i: 6
 debug.print: cnt: 8 i: 11

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

@@ -8,15 +8,15 @@ holy
 e-while
 g-label
 .
-a
-b
 cnt: 0 i: 0
 cnt: 1 i: 1
 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