Diagnose captured of non-sendable metatypes crossing isolation boundaries

Keep track of all of the type parameters and archetypes that are captured
by a local function or closure. Use that information to diagnose cases
where a non-Sendable metatype crosses an isolation boundary.

Author: DougGregor

include/swift/AST/CaptureInfo.h

@@ -17,6 +17,7 @@
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/OptionSet.h"
 #include "swift/Basic/SourceLoc.h"
+#include "swift/AST/Type.h"
 #include "swift/AST/TypeAlignments.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/PointerIntPair.h"
@@ -129,6 +130,18 @@ class CapturedValue {
   unsigned getFlags() const { return Value.getInt(); }
 };
 
+/// Describes a type that has been captured by a closure or local function.
+class CapturedType {
+  Type type;
+  SourceLoc loc;
+
+public:
+  CapturedType(Type type, SourceLoc loc) : type(type), loc(loc) { }
+  
+  Type getType() const { return type; }
+  SourceLoc getLoc() const { return loc; }
+};
+
 } // end swift namespace
 
 namespace swift {
@@ -140,26 +153,32 @@ class CaptureInfo {
   class CaptureInfoStorage final
       : public llvm::TrailingObjects<CaptureInfoStorage,
                                      CapturedValue,
-                                     GenericEnvironment *> {
+                                     GenericEnvironment *,
+                                     CapturedType> {
 
     DynamicSelfType *DynamicSelf;
     OpaqueValueExpr *OpaqueValue;
     unsigned NumCapturedValues;
     unsigned NumGenericEnvironments;
+    unsigned NumCapturedTypes;
 
   public:
     explicit CaptureInfoStorage(DynamicSelfType *dynamicSelf,
                                 OpaqueValueExpr *opaqueValue,
                                 unsigned numCapturedValues,
-                                unsigned numGenericEnvironments)
+                                unsigned numGenericEnvironments,
+                                unsigned numCapturedTypes)
       : DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue),
         NumCapturedValues(numCapturedValues),
-        NumGenericEnvironments(numGenericEnvironments) { }
+        NumGenericEnvironments(numGenericEnvironments),
+        NumCapturedTypes(numCapturedTypes) { }
 
     ArrayRef<CapturedValue> getCaptures() const;
 
     ArrayRef<GenericEnvironment *> getGenericEnvironments() const;
 
+    ArrayRef<CapturedType> getCapturedTypes() const;
+
     DynamicSelfType *getDynamicSelfType() const {
       return DynamicSelf;
     }
@@ -171,6 +190,14 @@ class CaptureInfo {
     unsigned numTrailingObjects(OverloadToken<CapturedValue>) const {
       return NumCapturedValues;
     }
+
+    unsigned numTrailingObjects(OverloadToken<GenericEnvironment *>) const {
+      return NumGenericEnvironments;
+    }
+
+    unsigned numTrailingObjects(OverloadToken<CapturedType>) const {
+      return NumCapturedTypes;
+    }
   };
 
   enum class Flags : unsigned {
@@ -187,7 +214,8 @@ class CaptureInfo {
               ArrayRef<CapturedValue> captures,
               DynamicSelfType *dynamicSelf, OpaqueValueExpr *opaqueValue,
               bool genericParamCaptures,
-              ArrayRef<GenericEnvironment *> genericEnv=ArrayRef<GenericEnvironment*>());
+              ArrayRef<GenericEnvironment *> genericEnv=ArrayRef<GenericEnvironment*>(),
+              ArrayRef<CapturedType> capturedTypes = ArrayRef<CapturedType>());
 
   /// A CaptureInfo representing no captures at all.
   static CaptureInfo empty();
@@ -196,11 +224,7 @@ class CaptureInfo {
     return StorageAndFlags.getPointer();
   }
 
-  bool isTrivial() const {
-    assert(hasBeenComputed());
-    return getCaptures().empty() && !hasGenericParamCaptures() &&
-           !hasDynamicSelfCapture() && !hasOpaqueValueCapture();
-  }
+  bool isTrivial() const;
 
   /// Returns all captured values and opaque expressions.
   ArrayRef<CapturedValue> getCaptures() const {
@@ -214,6 +238,12 @@ class CaptureInfo {
     return StorageAndFlags.getPointer()->getGenericEnvironments();
   }
 
+  /// Returns all captured values and opaque expressions.
+  ArrayRef<CapturedType> getCapturedTypes() const {
+    assert(hasBeenComputed());
+    return StorageAndFlags.getPointer()->getCapturedTypes();
+  }
+
   /// \returns true if the function captures the primary generic environment
   /// from its innermost declaration context.
   bool hasGenericParamCaptures() const {

include/swift/AST/DiagnosticsSema.def

@@ -5607,6 +5607,10 @@ ERROR(non_sendable_isolated_capture,none,
       "capture of %1 with non-sendable type %0 in an isolated "
       "%select{local function|closure}2",
       (Type, DeclName, bool))
+ERROR(non_sendable_metatype_capture,none,
+      "capture of non-sendable type %0 in an isolated "
+      "%select{local function|closure}1",
+      (Type, bool))
 ERROR(self_capture_deinit_task,none,
       "capture of 'self' in a closure that outlives deinit",
       ())

lib/AST/CaptureInfo.cpp

@@ -54,6 +54,11 @@ CaptureInfo::CaptureInfoStorage::getGenericEnvironments() const {
   return llvm::ArrayRef(this->getTrailingObjects<GenericEnvironment *>(), NumGenericEnvironments);
 }
 
+ArrayRef<CapturedType>
+CaptureInfo::CaptureInfoStorage::getCapturedTypes() const {
+  return llvm::ArrayRef(this->getTrailingObjects<CapturedType>(), NumCapturedTypes);
+}
+
 //===----------------------------------------------------------------------===//
 //                             MARK: CaptureInfo
 //===----------------------------------------------------------------------===//
@@ -62,7 +67,8 @@ CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
                          DynamicSelfType *dynamicSelf,
                          OpaqueValueExpr *opaqueValue,
                          bool genericParamCaptures,
-                         ArrayRef<GenericEnvironment *> genericEnv) {
+                         ArrayRef<GenericEnvironment *> genericEnv,
+                         ArrayRef<CapturedType> capturedTypes) {
   static_assert(IsTriviallyDestructible<CapturedValue>::value,
                 "Capture info is alloc'd on the ASTContext and not destroyed");
   static_assert(IsTriviallyDestructible<CaptureInfo::CaptureInfoStorage>::value,
@@ -79,37 +85,49 @@ CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
   if (genericParamCaptures)
     flags |= Flags::HasGenericParamCaptures;
 
-  if (captures.empty() && genericEnv.empty() && !dynamicSelf && !opaqueValue) {
+  if (captures.empty() && genericEnv.empty() && !dynamicSelf && !opaqueValue &&
+      capturedTypes.empty()) {
     *this = CaptureInfo::empty();
     StorageAndFlags.setInt(flags);
     return;
   }
 
   size_t storageToAlloc =
       CaptureInfoStorage::totalSizeToAlloc<CapturedValue,
-                                           GenericEnvironment *>(captures.size(),
-                                                                 genericEnv.size());
+                                           GenericEnvironment *,
+                                           CapturedType>(captures.size(),
+                                                         genericEnv.size(),
+                                                         capturedTypes.size());
   void *storageBuf = ctx.Allocate(storageToAlloc, alignof(CaptureInfoStorage));
   auto *storage = new (storageBuf) CaptureInfoStorage(dynamicSelf,
                                                       opaqueValue,
                                                       captures.size(),
-                                                      genericEnv.size());
+                                                      genericEnv.size(),
+                                                      capturedTypes.size());
   StorageAndFlags.setPointerAndInt(storage, flags);
   std::uninitialized_copy(captures.begin(), captures.end(),
                           storage->getTrailingObjects<CapturedValue>());
   std::uninitialized_copy(genericEnv.begin(), genericEnv.end(),
                           storage->getTrailingObjects<GenericEnvironment *>());
+  std::uninitialized_copy(capturedTypes.begin(), capturedTypes.end(),
+                          storage->getTrailingObjects<CapturedType>());
 }
 
 CaptureInfo CaptureInfo::empty() {
   static const CaptureInfoStorage empty{/*dynamicSelf*/nullptr,
                                         /*opaqueValue*/nullptr,
-                                        0, 0};
+                                        0, 0, 0};
   CaptureInfo result;
   result.StorageAndFlags.setPointer(&empty);
   return result;
 }
 
+bool CaptureInfo::isTrivial() const {
+  assert(hasBeenComputed());
+  return getCaptures().empty() && !hasGenericParamCaptures() &&
+         !hasDynamicSelfCapture() && !hasOpaqueValueCapture();
+}
+
 VarDecl *CaptureInfo::getIsolatedParamCapture() const {
   for (const auto &capture : getCaptures()) {
     // NOTE: isLocalCapture() returns false if we have dynamic self metadata
@@ -178,6 +196,12 @@ void CaptureInfo::print(raw_ostream &OS) const {
                OS << genericEnv->getOpenedElementShapeClass();
              },
              [&] { OS << ","; });
+
+  interleave(getCapturedTypes(),
+             [&](CapturedType type) {
+               OS << " type=" << type.getType().getString();
+             },
+             [&] { OS << ","; });
   OS << ')';
 }
 

lib/SIL/IR/TypeLowering.cpp

@@ -4425,6 +4425,10 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   // Captured pack element environments.
   llvm::SetVector<GenericEnvironment *> genericEnv;
 
+  // Captured types.
+  SmallVector<CapturedType, 4> capturedTypes;
+  llvm::SmallDenseSet<CanType, 4> alreadyCapturedTypes;
+
   bool capturesGenericParams = false;
   DynamicSelfType *capturesDynamicSelf = nullptr;
   OpaqueValueExpr *capturesOpaqueValue = nullptr;
@@ -4620,6 +4624,13 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
       // Collect non-function captures.
       recordCapture(capture);
     }
+
+    for (const auto &capturedType : captureInfo.getCapturedTypes()) {
+      if (alreadyCapturedTypes.insert(capturedType.getType()->getCanonicalType())
+              .second) {
+        capturedTypes.push_back(capturedType);
+      }
+    }
   };
 
   collectFunctionCaptures = [&](AnyFunctionRef curFn) {
@@ -4697,7 +4708,8 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   // Cache the result.
   CaptureInfo info(Context, resultingCaptures,
                    capturesDynamicSelf, capturesOpaqueValue,
-                   capturesGenericParams, genericEnv.getArrayRef());
+                   capturesGenericParams, genericEnv.getArrayRef(),
+                   capturedTypes);
   auto inserted = LoweredCaptures.insert({fn, info});
   assert(inserted.second && "already in map?!");
   (void)inserted;

lib/Sema/TypeCheckCaptures.cpp

@@ -54,6 +54,10 @@ class FindCapturedVars : public ASTWalker {
   /// can go here too.
   llvm::SetVector<GenericEnvironment *> CapturedEnvironments;
 
+  /// The captured types.
+  SmallVector<CapturedType, 4> CapturedTypes;
+  llvm::SmallDenseMap<CanType, unsigned, 4> CapturedTypeEntryNumber;
+
   SourceLoc GenericParamCaptureLoc;
   SourceLoc DynamicSelfCaptureLoc;
   DynamicSelfType *DynamicSelf = nullptr;
@@ -83,7 +87,8 @@ class FindCapturedVars : public ASTWalker {
 
     return CaptureInfo(Context, Captures, dynamicSelfToRecord,
                        OpaqueValue, HasGenericParamCaptures,
-                       CapturedEnvironments.getArrayRef());
+                       CapturedEnvironments.getArrayRef(),
+                       CapturedTypes);
   }
 
   bool hasGenericParamCaptures() const {
@@ -156,6 +161,22 @@ class FindCapturedVars : public ASTWalker {
       }));
     }
 
+    // Note that we're using a generic type.
+    auto recordUseOfGenericType = [&](Type type) {
+      if (!HasGenericParamCaptures) {
+        GenericParamCaptureLoc = loc;
+        HasGenericParamCaptures = true;
+      }
+
+      auto [insertionPos, inserted] = CapturedTypeEntryNumber.insert(
+          {type->getCanonicalType(), CapturedTypes.size()});
+      if (inserted) {
+        CapturedTypes.push_back(CapturedType(type, loc));
+      } else if (CapturedTypes[insertionPos->second].getLoc().isInvalid()) {
+        CapturedTypes[insertionPos->second] = CapturedType(type, loc);
+      }
+    };
+
     // Similar to dynamic 'Self', IRGen doesn't really need type metadata
     // for class-bound archetypes in nearly as many cases as with opaque
     // archetypes.
@@ -174,23 +195,18 @@ class FindCapturedVars : public ASTWalker {
             CapturedEnvironments.insert(env);
         }
 
-        if ((t->is<PrimaryArchetypeType>() ||
-             t->is<PackArchetypeType>() ||
-             t->is<GenericTypeParamType>()) &&
-            !HasGenericParamCaptures) {
-          GenericParamCaptureLoc = loc;
-          HasGenericParamCaptures = true;
+        if (t->is<PrimaryArchetypeType>() ||
+            t->is<PackArchetypeType>() ||
+            t->is<GenericTypeParamType>()) {
+          recordUseOfGenericType(t);
         }
       }));
     }
 
     if (auto *gft = type->getAs<GenericFunctionType>()) {
       TypeCaptureWalker walker(ObjC, [&](Type t) {
-        if (t->is<GenericTypeParamType>() &&
-            !HasGenericParamCaptures) {
-          GenericParamCaptureLoc = loc;
-          HasGenericParamCaptures = true;
-        }
+        if (t->is<GenericTypeParamType>())
+          recordUseOfGenericType(t);
       });
 
       for (const auto &param : gft->getParams())