[Concurrency] Ensure we get the right context for captured variables.

The declaration context of an explicitly-captured value is the context
of the original, captured declaration itself... not the closure in which
the value is captured. Account for this in data race checking, by tracking
the effective capture context for such variables. This eliminates some
erroneous complains about accesses to explicitly-captured variables in
concurrent code.

Part of rdar://74281361.

Author: DougGregor

lib/Sema/TypeCheckConcurrency.cpp

@@ -886,6 +886,11 @@ namespace {
     SmallVector<const DeclContext *, 4> contextStack;
     SmallVector<ApplyExpr*, 4> applyStack;
 
+    /// Keeps track of the capture context of variables that have been
+    /// explicitly captured in closures.
+    llvm::SmallDenseMap<VarDecl *, TinyPtrVector<const DeclContext *>>
+      captureContexts;
+
     using MutableVarSource = llvm::PointerUnion<DeclRefExpr *, InOutExpr *>;
     using MutableVarParent = llvm::PointerUnion<InOutExpr *, LoadExpr *>;
 
@@ -1134,6 +1139,13 @@ namespace {
       if (isa<ObjCSelectorExpr>(expr))
         return { false, expr };
 
+      // Track the capture contexts for variables.
+      if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
+        for (const auto &entry : captureList->getCaptureList()) {
+          captureContexts[entry.Var].push_back(captureList->getClosureBody());
+        }
+      }
+
       return { true, expr };
     }
 
@@ -1156,6 +1168,17 @@ namespace {
         mutableLocalVarParent.erase(inoutExpr);
       }
 
+      // Remove the tracked capture contexts.
+      if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
+        for (const auto &entry : captureList->getCaptureList()) {
+          auto &contexts = captureContexts[entry.Var];
+          assert(contexts.back() == captureList->getClosureBody());
+          contexts.pop_back();
+          if (contexts.empty())
+            captureContexts.erase(entry.Var);
+        }
+      }
+
       return expr;
     }
 
@@ -1526,6 +1549,21 @@ namespace {
       llvm_unreachable("unhandled actor isolation kind!");
     }
 
+    /// Find the innermost context in which this declaration was explicitly
+    /// captured.
+    const DeclContext *findCapturedDeclContext(ValueDecl *value) {
+      assert(value->isLocalCapture());
+      auto var = dyn_cast<VarDecl>(value);
+      if (!var)
+        return value->getDeclContext();
+
+      auto knownContexts = captureContexts.find(var);
+      if (knownContexts == captureContexts.end())
+        return value->getDeclContext();
+
+      return knownContexts->second.back();
+    }
+
     /// Check a reference to a local or global.
     bool checkNonMemberReference(
         ConcreteDeclRef valueRef, SourceLoc loc, DeclRefExpr *declRefExpr) {
@@ -1552,7 +1590,7 @@ namespace {
         // Check whether we are in a context that will not execute concurrently
         // with the context of 'self'. If not, it's safe.
         if (!mayExecuteConcurrentlyWith(
-                getDeclContext(), isolation.getLocalContext()))
+                getDeclContext(), findCapturedDeclContext(value)))
           return false;
 
         // Check whether this is a local variable, in which case we can

test/Concurrency/actor_isolation.swift

@@ -139,12 +139,28 @@ extension MyActor {
     }
 
     // Concurrent closures might run... concurrently.
-    acceptConcurrentClosure {
+    var otherLocalVar = 12
+    acceptConcurrentClosure { [otherLocalVar] in
+      defer {
+        _ = otherLocalVar
+      }
+
       _ = self.text[0] // expected-error{{actor-isolated property 'text' cannot be referenced from a concurrent closure}}
       _ = self.synchronous() // expected-error{{actor-isolated instance method 'synchronous()' cannot be referenced from a concurrent closure}}
       _ = localVar // expected-error{{reference to captured var 'localVar' in concurrently-executing code}}
       localVar = 25 // expected-error{{mutation of captured var 'localVar' in concurrently-executing code}}
       _ = localConstant
+
+      _ = otherLocalVar
+    }
+    otherLocalVar = 17
+
+    acceptConcurrentClosure { [weak self, otherLocalVar] in
+      defer {
+        _ = self?.actorIndependentVar
+      }
+
+      _ = otherLocalVar
     }
 
     // Escaping closures might run concurrently.