AST: Record captured generic environments in CaptureInfo

slavapestov · slavapestov · commit 97d95218f81d · 2024-04-30T21:38:15.000-04:00
diff --git a/include/swift/AST/CaptureInfo.h b/include/swift/AST/CaptureInfo.h
@@ -43,6 +43,7 @@ class ValueDecl;
 class FuncDecl;
 class OpaqueValueExpr;
 class VarDecl;
+class GenericEnvironment;
 
 /// CapturedValue includes both the declaration being captured, along with flags
 /// that indicate how it is captured.
@@ -140,19 +141,27 @@ class DynamicSelfType;
 /// Stores information about captured variables.
 class CaptureInfo {
   class CaptureInfoStorage final
-      : public llvm::TrailingObjects<CaptureInfoStorage, CapturedValue> {
+      : public llvm::TrailingObjects<CaptureInfoStorage,
+                                     CapturedValue,
+                                     GenericEnvironment *> {
 
     DynamicSelfType *DynamicSelf;
     OpaqueValueExpr *OpaqueValue;
-    unsigned Count;
+    unsigned NumCapturedValues;
+    unsigned NumGenericEnvironments;
+
   public:
-    explicit CaptureInfoStorage(unsigned count, DynamicSelfType *dynamicSelf,
-                                OpaqueValueExpr *opaqueValue)
-      : DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue), Count(count) { }
+    explicit CaptureInfoStorage(DynamicSelfType *dynamicSelf,
+                                OpaqueValueExpr *opaqueValue,
+                                unsigned numCapturedValues,
+                                unsigned numGenericEnvironments)
+      : DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue),
+        NumCapturedValues(numCapturedValues),
+        NumGenericEnvironments(numGenericEnvironments) { }
 
-    ArrayRef<CapturedValue> getCaptures() const {
-      return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), Count);
-    }
+    ArrayRef<CapturedValue> getCaptures() const;
+
+    ArrayRef<GenericEnvironment *> getGenericEnvironments() const;
 
     DynamicSelfType *getDynamicSelfType() const {
       return DynamicSelf;
@@ -161,6 +170,10 @@ class CaptureInfo {
     OpaqueValueExpr *getOpaqueValue() const {
       return OpaqueValue;
     }
+
+    unsigned numTrailingObjects(OverloadToken<CapturedValue>) const {
+      return NumCapturedValues;
+    }
   };
 
   enum class Flags : unsigned {
@@ -173,9 +186,11 @@ class CaptureInfo {
 public:
   /// The default-constructed CaptureInfo is "not yet computed".
   CaptureInfo() = default;
-  CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
+  CaptureInfo(ASTContext &ctx,
+              ArrayRef<CapturedValue> captures,
               DynamicSelfType *dynamicSelf, OpaqueValueExpr *opaqueValue,
-              bool genericParamCaptures);
+              bool genericParamCaptures,
+              ArrayRef<GenericEnvironment *> genericEnv=ArrayRef<GenericEnvironment*>());
 
   /// A CaptureInfo representing no captures at all.
   static CaptureInfo empty();
@@ -190,12 +205,20 @@ class CaptureInfo {
            !hasDynamicSelfCapture() && !hasOpaqueValueCapture();
   }
 
+  /// Returns all captured values and opaque expressions.
   ArrayRef<CapturedValue> getCaptures() const {
     assert(hasBeenComputed());
     return StorageAndFlags.getPointer()->getCaptures();
   }
 
-  /// \returns true if the function captures any generic type parameters.
+  /// Returns all captured pack element environments.
+  ArrayRef<GenericEnvironment *> getGenericEnvironments() const {
+    assert(hasBeenComputed());
+    return StorageAndFlags.getPointer()->getGenericEnvironments();
+  }
+
+  /// \returns true if the function captures the primary generic environment
+  /// from its innermost declaration context.
   bool hasGenericParamCaptures() const {
     assert(hasBeenComputed());
     return StorageAndFlags.getInt().contains(Flags::HasGenericParamCaptures);
diff --git a/lib/AST/CaptureInfo.cpp b/lib/AST/CaptureInfo.cpp
@@ -13,47 +13,72 @@
 #include "swift/AST/CaptureInfo.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
 
+ArrayRef<CapturedValue>
+CaptureInfo::CaptureInfoStorage::getCaptures() const {
+  return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), NumCapturedValues);
+}
+
+ArrayRef<GenericEnvironment *>
+CaptureInfo::CaptureInfoStorage::getGenericEnvironments() const {
+  return llvm::ArrayRef(this->getTrailingObjects<GenericEnvironment *>(), NumGenericEnvironments);
+}
+
 //===----------------------------------------------------------------------===//
 //                             MARK: CaptureInfo
 //===----------------------------------------------------------------------===//
 
 CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
                          DynamicSelfType *dynamicSelf,
                          OpaqueValueExpr *opaqueValue,
-                         bool genericParamCaptures) {
+                         bool genericParamCaptures,
+                         ArrayRef<GenericEnvironment *> genericEnv) {
   static_assert(IsTriviallyDestructible<CapturedValue>::value,
                 "Capture info is alloc'd on the ASTContext and not destroyed");
   static_assert(IsTriviallyDestructible<CaptureInfo::CaptureInfoStorage>::value,
                 "Capture info is alloc'd on the ASTContext and not destroyed");
 
+  // This is the only kind of local generic environment we can capture right now.
+#ifndef NDEBUG
+  for (auto *env : genericEnv) {
+    assert(env->getKind() == GenericEnvironment::Kind::OpenedElement);
+  }
+#endif
+
   OptionSet<Flags> flags;
   if (genericParamCaptures)
     flags |= Flags::HasGenericParamCaptures;
 
-  if (captures.empty() && !dynamicSelf && !opaqueValue) {
+  if (captures.empty() && genericEnv.empty() && !dynamicSelf && !opaqueValue) {
     *this = CaptureInfo::empty();
     StorageAndFlags.setInt(flags);
     return;
   }
 
   size_t storageToAlloc =
-      CaptureInfoStorage::totalSizeToAlloc<CapturedValue>(captures.size());
+      CaptureInfoStorage::totalSizeToAlloc<CapturedValue,
+                                           GenericEnvironment *>(captures.size(),
+                                                                 genericEnv.size());
   void *storageBuf = ctx.Allocate(storageToAlloc, alignof(CaptureInfoStorage));
-  auto *storage = new (storageBuf) CaptureInfoStorage(captures.size(),
-                                                      dynamicSelf,
-                                                      opaqueValue);
+  auto *storage = new (storageBuf) CaptureInfoStorage(dynamicSelf,
+                                                      opaqueValue,
+                                                      captures.size(),
+                                                      genericEnv.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 *>());
 }
 
 CaptureInfo CaptureInfo::empty() {
-  static const CaptureInfoStorage empty{0, /*dynamicSelf*/nullptr,
-                                        /*opaqueValue*/nullptr};
+  static const CaptureInfoStorage empty{/*dynamicSelf*/nullptr,
+                                        /*opaqueValue*/nullptr,
+                                        0, 0};
   CaptureInfo result;
   result.StorageAndFlags.setPointer(&empty);
   return result;
@@ -109,6 +134,13 @@ void CaptureInfo::print(raw_ostream &OS) const {
                  OS << "<noescape>";
              },
              [&] { OS << ", "; });
+
+  interleave(getGenericEnvironments(),
+             [&](GenericEnvironment *genericEnv) {
+               OS << " shape_class=";
+               OS << genericEnv->getOpenedElementShapeClass();
+             },
+             [&] { OS << ","; });
   OS << ')';
 }
 
diff --git a/lib/SIL/IR/TypeLowering.cpp b/lib/SIL/IR/TypeLowering.cpp
@@ -4150,6 +4150,9 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   // that IRGen can pass dynamic 'Self' metadata.
   std::optional<CapturedValue> selfCapture;
 
+  // Captured pack element environments.
+  llvm::SetVector<GenericEnvironment *> genericEnv;
+
   bool capturesGenericParams = false;
   DynamicSelfType *capturesDynamicSelf = nullptr;
   OpaqueValueExpr *capturesOpaqueValue = nullptr;
@@ -4176,6 +4179,13 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
     if (captureInfo.hasOpaqueValueCapture())
       capturesOpaqueValue = captureInfo.getOpaqueValue();
 
+    // Any set of mutually-recursive local functions will capture the same
+    // element environments, because we can't "cross" a pack expansion
+    // expression.
+    for (auto *env : captureInfo.getGenericEnvironments()) {
+      genericEnv.insert(env);
+    }
+
     for (auto capture : captureInfo.getCaptures()) {
       if (!capture.isLocalCapture())
         continue;
@@ -4399,8 +4409,9 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   }
 
   // Cache the uniqued set of transitive captures.
-  CaptureInfo info{Context, resultingCaptures, capturesDynamicSelf,
-                   capturesOpaqueValue, capturesGenericParams};
+  CaptureInfo info(Context, resultingCaptures,
+                   capturesDynamicSelf, capturesOpaqueValue,
+                   capturesGenericParams, genericEnv.getArrayRef());
   auto inserted = LoweredCaptures.insert({fn, info});
   assert(inserted.second && "already in map?!");
   (void)inserted;
diff --git a/lib/Sema/TypeCheckCaptures.cpp b/lib/Sema/TypeCheckCaptures.cpp
@@ -38,6 +38,16 @@ class FindCapturedVars : public ASTWalker {
   ASTContext &Context;
   SmallVector<CapturedValue, 4> Captures;
   llvm::SmallDenseMap<ValueDecl*, unsigned, 4> captureEntryNumber;
+
+  /// A stack of pack element environments we're currently walking into.
+  /// A reference to an element archetype defined by one of these is not
+  /// a capture.
+  llvm::SetVector<GenericEnvironment *> VisitingEnvironments;
+
+  /// A set of pack element environments we've encountered that were not
+  /// in the above stack; those are the captures.
+  llvm::SetVector<GenericEnvironment *> CapturedEnvironments;
+
   SourceLoc GenericParamCaptureLoc;
   SourceLoc DynamicSelfCaptureLoc;
   DynamicSelfType *DynamicSelf = nullptr;
@@ -65,8 +75,9 @@ class FindCapturedVars : public ASTWalker {
         dynamicSelfToRecord = DynamicSelf;
     }
 
-    return CaptureInfo(Context, Captures, dynamicSelfToRecord, OpaqueValue,
-                       HasGenericParamCaptures);
+    return CaptureInfo(Context, Captures, dynamicSelfToRecord,
+                       OpaqueValue, HasGenericParamCaptures,
+                       CapturedEnvironments.getArrayRef());
   }
 
   bool hasGenericParamCaptures() const {
@@ -148,9 +159,17 @@ class FindCapturedVars : public ASTWalker {
     // perform it accurately.
     if (type->hasArchetype() || type->hasTypeParameter()) {
       type.walk(TypeCaptureWalker(ObjC, [&](Type t) {
-        if ((t->is<ArchetypeType>() ||
+        // Record references to element archetypes that were bound
+        // outside the body of the current closure.
+        if (auto *element = t->getAs<ElementArchetypeType>()) {
+          auto *env = element->getGenericEnvironment();
+          if (VisitingEnvironments.count(env) == 0)
+            CapturedEnvironments.insert(env);
+        }
+
+        if ((t->is<PrimaryArchetypeType>() ||
+             t->is<PackArchetypeType>() ||
              t->is<GenericTypeParamType>()) &&
-            !t->isOpenedExistential() &&
             !HasGenericParamCaptures) {
           GenericParamCaptureLoc = loc;
           HasGenericParamCaptures = true;
@@ -619,6 +638,25 @@ class FindCapturedVars : public ASTWalker {
       }
     }
 
+    if (auto expansion = dyn_cast<PackExpansionExpr>(E)) {
+      if (auto *env = expansion->getGenericEnvironment()) {
+        assert(VisitingEnvironments.count(env) == 0);
+        VisitingEnvironments.insert(env);
+      }
+    }
+
+    return Action::Continue(E);
+  }
+
+  PostWalkResult<Expr *> walkToExprPost(Expr *E) override {
+    if (auto expansion = dyn_cast<PackExpansionExpr>(E)) {
+      if (auto *env = expansion->getGenericEnvironment()) {
+        assert(env == VisitingEnvironments.back());
+        VisitingEnvironments.pop_back();
+        (void) env;
+      }
+    }
+
     return Action::Continue(E);
   }
 };
