[Concurrency] Fix a few issues with actor-isolated 'inout' argument diagnostics.

ActorIsolationChecker's apply stack was not considering LookupExpr, which caused
bogus diagnostics for inout arguments that happen to be subexpressions of async
function calls. Similarly, inout arguments to async subscript calls were not
diagnosed because the subscript application was not tracked in the apply stack.

There's still an issue with async computed getter calls because the type-checked
AST does not have an InOutExpr for the base expression.

(cherry picked from commit 17ac5fe)

Author: hborla

lib/Sema/TypeCheckConcurrency.cpp

@@ -556,20 +556,30 @@ findReference(Expr *expr) {
 ///
 /// Note that this must be called after the implicitlyAsync flag has been set,
 /// or implicitly async calls will not return the correct value.
-static bool isAsyncCall(const ApplyExpr *call) {
-  if (call->isImplicitlyAsync())
+static bool isAsyncCall(
+    llvm::PointerUnion<ApplyExpr *, LookupExpr *> call) {
+
+  if (auto *apply = call.dyn_cast<ApplyExpr *>()) {
+    if (apply->isImplicitlyAsync())
+      return true;
+
+    // Effectively the same as doing a
+    // `cast_or_null<FunctionType>(call->getFn()->getType())`, check the
+    // result of that and then checking `isAsync` if it's defined.
+    Type funcTypeType = apply->getFn()->getType();
+    if (!funcTypeType)
+      return false;
+    AnyFunctionType *funcType = funcTypeType->getAs<AnyFunctionType>();
+    if (!funcType)
+      return false;
+    return funcType->isAsync();
+  }
+
+  auto *lookup = call.get<LookupExpr *>();
+  if (lookup->isImplicitlyAsync())
     return true;
 
-  // Effectively the same as doing a
-  // `cast_or_null<FunctionType>(call->getFn()->getType())`, check the
-  // result of that and then checking `isAsync` if it's defined.
-  Type funcTypeType = call->getFn()->getType();
-  if (!funcTypeType)
-    return false;
-  AnyFunctionType *funcType = funcTypeType->getAs<AnyFunctionType>();
-  if (!funcType)
-    return false;
-  return funcType->isAsync();
+  return isAsyncDecl(lookup->getDecl());
 }
 
 /// Determine whether we should diagnose data races within the current context.
@@ -1932,7 +1942,7 @@ namespace {
   class ActorIsolationChecker : public ASTWalker {
     ASTContext &ctx;
     SmallVector<const DeclContext *, 4> contextStack;
-    SmallVector<ApplyExpr*, 4> applyStack;
+    SmallVector<llvm::PointerUnion<ApplyExpr *, LookupExpr *>, 4> applyStack;
     SmallVector<std::pair<OpaqueValueExpr *, Expr *>, 4> opaqueValues;
     SmallVector<const PatternBindingDecl *, 2> patternBindingStack;
     llvm::function_ref<Type(Expr *)> getType;
@@ -1950,6 +1960,13 @@ namespace {
     using MutableVarParent
         = llvm::PointerUnion<InOutExpr *, LoadExpr *, AssignExpr *>;
 
+    ApplyExpr *getImmediateApply() const {
+      if (applyStack.empty())
+        return nullptr;
+
+      return applyStack.back().dyn_cast<ApplyExpr *>();
+    }
+
     const PatternBindingDecl *getTopPatternBindingDecl() const {
       return patternBindingStack.empty() ? nullptr : patternBindingStack.back();
     }
@@ -2207,7 +2224,7 @@ namespace {
 
       const auto End = applyStack.rend();
       for (auto I = applyStack.rbegin(); I != End; ++I)
-        if (auto call = dyn_cast<CallExpr>(*I)) {
+        if (auto call = dyn_cast<CallExpr>(I->dyn_cast<ApplyExpr *>())) {
           if (setAsync) {
             call->setImplicitlyAsync(*setAsync);
           }
@@ -2261,6 +2278,11 @@ namespace {
     }
 
     PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
+      // Skip expressions that didn't make it to solution application
+      // because the constraint system diagnosed an error.
+      if (!expr->getType())
+        return Action::SkipChildren(expr);
+
       if (auto *openExistential = dyn_cast<OpenExistentialExpr>(expr)) {
         opaqueValues.push_back({
             openExistential->getOpaqueValue(),
@@ -2295,6 +2317,7 @@ namespace {
         recordMutableVarParent(load, load->getSubExpr());
 
       if (auto lookup = dyn_cast<LookupExpr>(expr)) {
+        applyStack.push_back(lookup);
         checkReference(lookup->getBase(), lookup->getMember(), lookup->getLoc(),
                        /*partialApply*/ llvm::None, lookup);
         return Action::Continue(expr);
@@ -2337,7 +2360,7 @@ namespace {
           // Self applications are checked as part of the outer call.
           // However, we look for inout issues here.
           if (applyStack.size() >= 2) {
-            ApplyExpr *outerCall = applyStack[applyStack.size() - 2];
+            auto outerCall = applyStack[applyStack.size() - 2];
             if (isAsyncCall(outerCall)) {
               // This call is a partial application within an async call.
               // If the partial application take a value inout, it is bad.
@@ -2391,17 +2414,21 @@ namespace {
       }
 
       if (auto *apply = dyn_cast<ApplyExpr>(expr)) {
-        assert(applyStack.back() == apply);
+        assert(applyStack.back().get<ApplyExpr *>() == apply);
         applyStack.pop_back();
       }
 
       // Clear out the mutable local variable parent map on the way out.
-      if (auto *declRefExpr = dyn_cast<DeclRefExpr>(expr))
+      if (auto *declRefExpr = dyn_cast<DeclRefExpr>(expr)) {
         mutableLocalVarParent.erase(declRefExpr);
-      else if (auto *lookupExpr = dyn_cast<LookupExpr>(expr))
+      } else if (auto *lookupExpr = dyn_cast<LookupExpr>(expr)) {
         mutableLocalVarParent.erase(lookupExpr);
-      else if (auto *inoutExpr = dyn_cast<InOutExpr>(expr))
+
+        assert(applyStack.back().dyn_cast<LookupExpr *>() == lookupExpr);
+        applyStack.pop_back();
+      } else if (auto *inoutExpr = dyn_cast<InOutExpr>(expr)) {
         mutableLocalVarParent.erase(inoutExpr);
+      }
 
       // Remove the tracked capture contexts.
       if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
@@ -2616,28 +2643,31 @@ namespace {
     /// Diagnose an inout argument passed into an async call
     ///
     /// \returns true if we diagnosed the entity, \c false otherwise.
-    bool diagnoseInOutArg(const ApplyExpr *call, const InOutExpr *arg,
-                          bool isPartialApply) {
+    bool diagnoseInOutArg(
+        llvm::PointerUnion<ApplyExpr *, LookupExpr *> call,
+        const InOutExpr *arg,
+        bool isPartialApply) {
       // check that the call is actually async
       if (!isAsyncCall(call))
         return false;
 
       bool result = false;
-      auto checkDiagnostic = [this, call, isPartialApply, &result](
+      auto diagnoseIsolatedInoutState = [this, call, isPartialApply, &result](
           ConcreteDeclRef declRef, SourceLoc argLoc) {
         auto decl = declRef.getDecl();
         auto isolation = getActorIsolationForReference(decl, getDeclContext());
         if (!isolation.isActorIsolated())
           return;
 
         if (isPartialApply) {
+          auto *apply = call.get<ApplyExpr *>();
           // The partially applied InoutArg is a property of actor. This
           // can really only happen when the property is a struct with a
           // mutating async method.
-          if (auto partialApply = dyn_cast<ApplyExpr>(call->getFn())) {
+          if (auto partialApply = dyn_cast<ApplyExpr>(apply->getFn())) {
             if (auto declRef = dyn_cast<DeclRefExpr>(partialApply->getFn())) {
               ValueDecl *fnDecl = declRef->getDecl();
-              ctx.Diags.diagnose(call->getLoc(),
+              ctx.Diags.diagnose(apply->getLoc(),
                                  diag::actor_isolated_mutating_func,
                                  fnDecl->getName(), decl);
               result = true;
@@ -2646,29 +2676,36 @@ namespace {
           }
         }
 
+        bool isImplicitlyAsync;
+        if (auto *apply = call.dyn_cast<ApplyExpr *>()) {
+          isImplicitlyAsync = apply->isImplicitlyAsync().has_value();
+        } else {
+          auto *lookup = call.get<LookupExpr *>();
+          isImplicitlyAsync = lookup->isImplicitlyAsync().has_value();
+        }
+
         ctx.Diags.diagnose(argLoc, diag::actor_isolated_inout_state,
-                           decl, call->isImplicitlyAsync().has_value());
+                           decl, isImplicitlyAsync);
         decl->diagnose(diag::kind_declared_here, decl->getDescriptiveKind());
         result = true;
         return;
       };
-      auto expressionWalker = [baseArg = arg->getSubExpr(),
-                               checkDiagnostic](Expr *expr) -> Expr * {
-        if (isa<InOutExpr>(expr))
-          return nullptr; // AST walker will hit this again
+
+      auto findIsolatedState = [&](Expr *expr) -> Expr * {
         if (LookupExpr *lookup = dyn_cast<LookupExpr>(expr)) {
           if (isa<DeclRefExpr>(lookup->getBase())) {
-            checkDiagnostic(lookup->getMember().getDecl(), baseArg->getLoc());
+            diagnoseIsolatedInoutState(lookup->getMember().getDecl(),
+                                       expr->getLoc());
             return nullptr; // Diagnosed. Don't keep walking
           }
         }
         if (DeclRefExpr *declRef = dyn_cast<DeclRefExpr>(expr)) {
-          checkDiagnostic(declRef->getDecl(), baseArg->getLoc());
+          diagnoseIsolatedInoutState(declRef->getDecl(), expr->getLoc());
           return nullptr; // Diagnosed. Don't keep walking
         }
         return expr;
       };
-      arg->getSubExpr()->forEachChildExpr(expressionWalker);
+      arg->getSubExpr()->forEachChildExpr(findIsolatedState);
       return result;
     }
 
@@ -3172,9 +3209,9 @@ namespace {
 
       // If this declaration is a callee from the enclosing application,
       // it's already been checked via the call.
-      if (!applyStack.empty()) {
+      if (auto *apply = getImmediateApply()) {
         auto immediateCallee =
-            applyStack.back()->getCalledValue(/*skipFunctionConversions=*/true);
+            apply->getCalledValue(/*skipFunctionConversions=*/true);
         if (decl == immediateCallee)
           return false;
       }

test/Concurrency/actor_inout_isolation.swift

@@ -136,8 +136,8 @@ extension TestActor {
   func callMutatingFunctionOnStruct() async {
     // expected-targeted-complete-warning @+4 {{passing argument of non-sendable type 'inout Point' outside of actor-isolated context may introduce data races}}
     // expected-error@+3:20{{cannot call mutating async function 'setComponents(x:y:)' on actor-isolated property 'position'}}
-    // expected-error@+2:51{{actor-isolated property 'nextPosition' cannot be passed 'inout' to 'async' function call}}
-    // expected-error@+1:71{{actor-isolated property 'nextPosition' cannot be passed 'inout' to 'async' function call}}
+    // expected-error@+2:38{{actor-isolated property 'nextPosition' cannot be passed 'inout' to 'async' function call}}
+    // expected-error@+1:58{{actor-isolated property 'nextPosition' cannot be passed 'inout' to 'async' function call}}
     await position.setComponents(x: &nextPosition.x, y: &nextPosition.y)
 
     // expected-targeted-complete-warning @+4 {{passing argument of non-sendable type 'inout Point' outside of actor-isolated context may introduce data races}}
@@ -265,3 +265,45 @@ struct Dog {
     await cat?.meow()
   }
 }
+
+@available(SwiftStdlib 5.1, *)
+func passToAsync(_: Int) async {}
+
+@available(SwiftStdlib 5.1, *)
+func wrapInClosure(
+  @_inheritActorContext _ block: @Sendable () async throws -> Void
+) async {}
+
+@available(SwiftStdlib 5.1, *)
+extension Array {
+  var mutateAsynchronously: Int {
+    mutating get async { 0 }
+  }
+
+  subscript(mutateAsynchronously i: Int) -> Int {
+    mutating get async { 0 }
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+actor ProtectArray {
+  var array: [Int] = []
+  // expected-note@-1 {{property declared here}}
+
+  func test() async {
+    // FIXME: this is invalid too!
+    _ = await array.mutateAsynchronously
+    // expected-targeted-complete-sns-warning@-1 {{non-sendable type '@lvalue [Int]' exiting actor-isolated context in call to non-isolated property 'mutateAsynchronously' cannot cross actor boundary}}
+
+    _ = await array[mutateAsynchronously: 0]
+    // expected-error@-1 {{actor-isolated property 'array' cannot be passed 'inout' to 'async' function call}}
+    // expected-targeted-complete-sns-warning@-2 {{non-sendable type 'inout Array<Int>' exiting actor-isolated context in call to non-isolated subscript 'subscript(mutateAsynchronously:)' cannot cross actor boundary}}
+
+    await passToAsync(array[0])
+
+    await wrapInClosure {
+      _ = array[0]
+      array.append(1)
+    }
+  }
+}