mark the nearest CallExpr as implicitly async; not just the nearest ApplyExpr.

Author: kavon

lib/Sema/TypeCheckConcurrency.cpp

@@ -621,6 +621,38 @@ namespace {
       contextStack.push_back(dc);
     }
 
+    /// Searches the applyStack from back to front for the inner-most CallExpr
+    /// and marks that CallExpr as implicitly async. 
+    ///
+    /// NOTE: Crashes if no CallExpr was found.
+    /// 
+    /// For example, for global actor function `curryAdd`, if we have:
+    ///     ((curryAdd 1) 2)
+    /// then we want to mark the inner-most CallExpr, `(curryAdd 1)`.
+    ///
+    /// The same goes for calls to member functions, such as calc.add(1, 2),
+    /// aka ((add calc) 1 2), looks like this:
+    /// 
+    ///  (call_expr
+    ///    (dot_syntax_call_expr
+    ///      (declref_expr add)
+    ///      (declref_expr calc))
+    ///    (tuple_expr 
+    ///      ...))
+    ///
+    /// and we reach up to mark the CallExpr.
+    void markNearestCallAsImplicitlyAsync() {
+      assert(applyStack.size() > 0 && "not contained within an Apply?");
+
+      const auto End = applyStack.rend();
+      for (auto I = applyStack.rbegin(); I != End; ++I)
+        if (auto call = dyn_cast<CallExpr>(*I)) {
+          call->setImplicitlyAsync(true);
+          return;
+        }
+      llvm_unreachable("expected a CallExpr in applyStack!");
+    }
+
     bool shouldWalkCaptureInitializerExpressions() override { return true; }
 
     bool shouldWalkIntoTapExpression() override { return false; }
@@ -661,10 +693,9 @@ namespace {
           if (auto memberRef = findMemberReference(partialApply->fn)) {
             // NOTE: partially-applied thunks are never annotated as 
             // implicitly async, regardless of whether they are escaping.
-            // So, we do not pass the ApplyExpr along to checkMemberReference.
             checkMemberReference(
                 partialApply->base, memberRef->first, memberRef->second,
-                partialApply->isEscaping);
+                partialApply->isEscaping, /*maybeImplicitAsync=*/false);
 
             partialApply->base->walk(*this);
 
@@ -683,7 +714,7 @@ namespace {
         if (auto memberRef = findMemberReference(fn)) {
           checkMemberReference(
               call->getArg(), memberRef->first, memberRef->second, 
-              /*isEscapingPartialApply=*/false, call);
+              /*isEscapingPartialApply=*/false, /*maybeImplicitAsync=*/true);
 
           call->getArg()->walk(*this);
 
@@ -903,7 +934,7 @@ namespace {
           auto concDecl = memberRef->first;
           if (value == concDecl.getDecl() && !apply->implicitlyAsync()) {
             // then this ValueDecl appears as the called value of the ApplyExpr.
-            apply->setImplicitlyAsync(true);
+            markNearestCallAsImplicitlyAsync();
             return true;
           }
         }
@@ -1012,7 +1043,7 @@ namespace {
     bool checkMemberReference(
         Expr *base, ConcreteDeclRef memberRef, SourceLoc memberLoc,
         bool isEscapingPartialApply = false, 
-        ApplyExpr *maybeImplicitAsync = nullptr) {
+        bool maybeImplicitAsync = false) {
       if (!base || !memberRef)
         return false;
 
@@ -1028,7 +1059,7 @@ namespace {
         if (!selfVar) {
           // actor-isolated non-self calls are implicitly async and thus OK.
           if (maybeImplicitAsync && isa<AbstractFunctionDecl>(member)) {
-            maybeImplicitAsync->setImplicitlyAsync(true);
+            markNearestCallAsImplicitlyAsync();
             return false;
           }
           ctx.Diags.diagnose(

test/Concurrency/actor_call_implicitly_async.swift

@@ -184,4 +184,34 @@ func blender(_ peeler : () -> Void) {
 
 @OrangeActor func quinoa() async {
   rice() // expected-error {{call is 'async' but is not marked with 'await'}}
+}
+
+///////////
+// check various curried applications to ensure we mark the right expression.
+
+actor class Calculator {
+  func addCurried(_ x : Int) -> ((Int) -> Int) { 
+    return { (_ y : Int) in x + y }
+  }
+
+  func add(_ x : Int, _ y : Int) -> Int {
+    return x + y
+  }
+}
+
+@BananaActor func bananaAdd(_ x : Int) -> ((Int) -> Int) { 
+  return { (_ y : Int) in x + y }
+}
+
+@OrangeActor func doSomething() async {
+  let _ = (await bananaAdd(1))(2)
+  let _ = await (await bananaAdd(1))(2) // expected-warning{{no calls to 'async' functions occur within 'await' expression}}
+
+  let calc = Calculator()
+  
+  let _ = (await calc.addCurried(1))(2)
+  let _ = await (await calc.addCurried(1))(2) // expected-warning{{no calls to 'async' functions occur within 'await' expression}}
+
+  let plusOne = await calc.addCurried(await calc.add(0, 1))
+  let _ = plusOne(2)
 }