[AST] Shrink CaptureInfo down to a PointerIntPair (swiftlang#27261)

If there are any actual captures, a CaptureInfoStorage will be
allocated in the ASTContext, matching the previous behavior of
allocating the array of captures there.

No functionality change. I /tried/ to change the functionality to
assert everywhere capture info was used but hadn't explicitly been
computed, but it turns out we do that /all over the place/ for any
declaration that's been synthesized, imported, or deserialized. I
left in some groundwork for someone to make this more explicit (or
requestify it) in the future.

Author: jrose-apple

include/swift/AST/CaptureInfo.h

@@ -14,11 +14,13 @@
 #define SWIFT_AST_CAPTURE_INFO_H
 
 #include "swift/Basic/LLVM.h"
+#include "swift/Basic/OptionSet.h"
 #include "swift/Basic/SourceLoc.h"
 #include "swift/AST/TypeAlignments.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/Support/TrailingObjects.h"
 #include <vector>
 
 namespace swift {
@@ -114,31 +116,63 @@ class DynamicSelfType;
 
 /// Stores information about captured variables.
 class CaptureInfo {
-  const CapturedValue *Captures;
-  DynamicSelfType *DynamicSelf;
-  OpaqueValueExpr *OpaqueValue;
-  unsigned Count = 0;
-  bool GenericParamCaptures : 1;
-  bool Computed : 1;
+  class CaptureInfoStorage final
+      : public llvm::TrailingObjects<CaptureInfoStorage, CapturedValue> {
+
+    DynamicSelfType *DynamicSelf;
+    OpaqueValueExpr *OpaqueValue;
+    unsigned Count;
+  public:
+    explicit CaptureInfoStorage(unsigned count, DynamicSelfType *dynamicSelf,
+                                OpaqueValueExpr *opaqueValue)
+      : DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue), Count(count) { }
+
+    ArrayRef<CapturedValue> getCaptures() const {
+      return llvm::makeArrayRef(this->getTrailingObjects<CapturedValue>(),
+                                Count);
+    }
+
+    DynamicSelfType *getDynamicSelfType() const {
+      return DynamicSelf;
+    }
+
+    OpaqueValueExpr *getOpaqueValue() const {
+      return OpaqueValue;
+    }
+  };
+
+  enum class Flags : unsigned {
+    HasGenericParamCaptures = 1 << 0
+  };
+
+  llvm::PointerIntPair<const CaptureInfoStorage *, 2, OptionSet<Flags>>
+      StorageAndFlags;
 
 public:
-  CaptureInfo()
-    : Captures(nullptr), DynamicSelf(nullptr), OpaqueValue(nullptr), Count(0),
-      GenericParamCaptures(0), Computed(0) { }
+  /// The default-constructed CaptureInfo is "not yet computed".
+  CaptureInfo() = default;
+  CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
+              DynamicSelfType *dynamicSelf, OpaqueValueExpr *opaqueValue,
+              bool genericParamCaptures);
 
-  bool hasBeenComputed() const { return Computed; }
+  /// A CaptureInfo representing no captures at all.
+  static CaptureInfo empty();
+
+  bool hasBeenComputed() const {
+    return StorageAndFlags.getPointer();
+  }
 
   bool isTrivial() const {
-    return Count == 0 && !GenericParamCaptures && !DynamicSelf && !OpaqueValue;
+    return getCaptures().empty() && !hasGenericParamCaptures() &&
+           !hasDynamicSelfCapture() && !hasOpaqueValueCapture();
   }
 
   ArrayRef<CapturedValue> getCaptures() const {
-    return llvm::makeArrayRef(Captures, Count);
-  }
-  void setCaptures(ArrayRef<CapturedValue> C) {
-    Captures = C.data();
-    Computed = true;
-    Count = C.size();
+    // FIXME: Ideally, everywhere that synthesizes a function should include
+    // its capture info.
+    if (!hasBeenComputed())
+      return None;
+    return StorageAndFlags.getPointer()->getCaptures();
   }
 
   /// Return a filtered list of the captures for this function,
@@ -152,37 +186,37 @@ class CaptureInfo {
 
   /// \returns true if the function captures any generic type parameters.
   bool hasGenericParamCaptures() const {
-    return GenericParamCaptures;
-  }
-
-  void setGenericParamCaptures(bool genericParamCaptures) {
-    GenericParamCaptures = genericParamCaptures;
+    // FIXME: Ideally, everywhere that synthesizes a function should include
+    // its capture info.
+    if (!hasBeenComputed())
+      return false;
+    return StorageAndFlags.getInt().contains(Flags::HasGenericParamCaptures);
   }
 
   /// \returns true if the function captures the dynamic Self type.
   bool hasDynamicSelfCapture() const {
-    return DynamicSelf != nullptr;
+    return getDynamicSelfType() != nullptr;
   }
 
   /// \returns the captured dynamic Self type, if any.
   DynamicSelfType *getDynamicSelfType() const {
-    return DynamicSelf;
-  }
-
-  void setDynamicSelfType(DynamicSelfType *dynamicSelf) {
-    DynamicSelf = dynamicSelf;
+    // FIXME: Ideally, everywhere that synthesizes a function should include
+    // its capture info.
+    if (!hasBeenComputed())
+      return nullptr;
+    return StorageAndFlags.getPointer()->getDynamicSelfType();
   }
 
   bool hasOpaqueValueCapture() const {
-    return OpaqueValue != nullptr;
+    return getOpaqueValue() != nullptr;
   }
 
   OpaqueValueExpr *getOpaqueValue() const {
-    return OpaqueValue;
-  }
-
-  void setOpaqueValue(OpaqueValueExpr *OVE) {
-    OpaqueValue = OVE;
+    // FIXME: Ideally, everywhere that synthesizes a function should include
+    // its capture info.
+    if (!hasBeenComputed())
+      return nullptr;
+    return StorageAndFlags.getPointer()->getOpaqueValue();
   }
 
   void dump() const;

lib/AST/CaptureInfo.cpp

@@ -11,11 +11,50 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/AST/CaptureInfo.h"
+#include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
 
+CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
+                         DynamicSelfType *dynamicSelf,
+                         OpaqueValueExpr *opaqueValue,
+                         bool genericParamCaptures) {
+  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");
+
+  OptionSet<Flags> flags;
+  if (genericParamCaptures)
+    flags |= Flags::HasGenericParamCaptures;
+
+  if (captures.empty() && !dynamicSelf && !opaqueValue) {
+    *this = CaptureInfo::empty();
+    StorageAndFlags.setInt(flags);
+    return;
+  }
+
+  size_t storageToAlloc =
+      CaptureInfoStorage::totalSizeToAlloc<CapturedValue>(captures.size());
+  void *storageBuf = ctx.Allocate(storageToAlloc, alignof(CaptureInfoStorage));
+  auto *storage = new (storageBuf) CaptureInfoStorage(captures.size(),
+                                                      dynamicSelf,
+                                                      opaqueValue);
+  StorageAndFlags.setPointerAndInt(storage, flags);
+  std::uninitialized_copy(captures.begin(), captures.end(),
+                          storage->getTrailingObjects<CapturedValue>());
+}
+
+CaptureInfo CaptureInfo::empty() {
+  static const CaptureInfoStorage empty{0, /*dynamicSelf*/nullptr,
+                                        /*opaqueValue*/nullptr};
+  CaptureInfo result;
+  result.StorageAndFlags.setPointer(&empty);
+  return result;
+}
+
 bool CaptureInfo::hasLocalCaptures() const {
   for (auto capture : getCaptures())
     if (capture.getDecl()->getDeclContext()->isLocalContext())
@@ -28,7 +67,7 @@ void CaptureInfo::
 getLocalCaptures(SmallVectorImpl<CapturedValue> &Result) const {
   if (!hasLocalCaptures()) return;
 
-  Result.reserve(Count);
+  Result.reserve(getCaptures().size());
 
   // Filter out global variables.
   for (auto capture : getCaptures()) {

lib/SIL/TypeLowering.cpp

@@ -2310,15 +2310,12 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   }
 
   // Cache the uniqued set of transitive captures.
-  auto inserted = LoweredCaptures.insert({fn, CaptureInfo()});
+  CaptureInfo info{Context, resultingCaptures, capturesDynamicSelf,
+                   capturesOpaqueValue, capturesGenericParams};
+  auto inserted = LoweredCaptures.insert({fn, info});
   assert(inserted.second && "already in map?!");
-  auto &cachedCaptures = inserted.first->second;
-  cachedCaptures.setGenericParamCaptures(capturesGenericParams);
-  cachedCaptures.setDynamicSelfType(capturesDynamicSelf);
-  cachedCaptures.setOpaqueValue(capturesOpaqueValue);
-  cachedCaptures.setCaptures(Context.AllocateCopy(resultingCaptures));
-  
-  return cachedCaptures;
+  (void)inserted;
+  return info;
 }
 
 /// Given that type1 is known to be a subtype of type2, check if the two

lib/Sema/TypeCheckCaptures.cpp

@@ -42,39 +42,33 @@ class FindCapturedVars : public ASTWalker {
   OpaqueValueExpr *OpaqueValue = nullptr;
   SourceLoc CaptureLoc;
   DeclContext *CurDC;
-  bool NoEscape, ObjC;
+  bool NoEscape, ObjC, IsGenericFunction;
 
 public:
   FindCapturedVars(ASTContext &Context,
                    SourceLoc CaptureLoc,
                    DeclContext *CurDC,
                    bool NoEscape,
-                   bool ObjC)
+                   bool ObjC,
+                   bool IsGenericFunction)
       : Context(Context), CaptureLoc(CaptureLoc), CurDC(CurDC),
-        NoEscape(NoEscape), ObjC(ObjC) {}
+        NoEscape(NoEscape), ObjC(ObjC), IsGenericFunction(IsGenericFunction) {}
 
   CaptureInfo getCaptureInfo() const {
-    CaptureInfo result;
-
-    // Anything can capture an opaque value placeholder.
-    if (OpaqueValue)
-      result.setOpaqueValue(OpaqueValue);
+    DynamicSelfType *dynamicSelfToRecord = nullptr;
+    bool hasGenericParamCaptures = IsGenericFunction;
 
     // Only local functions capture dynamic 'Self'.
     if (CurDC->getParent()->isLocalContext()) {
       if (GenericParamCaptureLoc.isValid())
-        result.setGenericParamCaptures(true);
+        hasGenericParamCaptures = true;
 
       if (DynamicSelfCaptureLoc.isValid())
-        result.setDynamicSelfType(DynamicSelf);
+        dynamicSelfToRecord = DynamicSelf;
     }
 
-    if (Captures.empty())
-      result.setCaptures(None);
-    else
-      result.setCaptures(Context.AllocateCopy(Captures));
-
-    return result;
+    return CaptureInfo(Context, Captures, dynamicSelfToRecord, OpaqueValue,
+                       hasGenericParamCaptures);
   }
 
   SourceLoc getGenericParamCaptureLoc() const {
@@ -590,28 +584,26 @@ void TypeChecker::computeCaptures(AnyFunctionRef AFR) {
 
   PrettyStackTraceAnyFunctionRef trace("computing captures for", AFR);
 
+  // A generic function always captures outer generic parameters.
+  bool isGeneric = false;
+  auto *AFD = AFR.getAbstractFunctionDecl();
+  if (AFD)
+    isGeneric = (AFD->getGenericParams() != nullptr);
+
   auto &Context = AFR.getAsDeclContext()->getASTContext();
   FindCapturedVars finder(Context,
                           AFR.getLoc(),
                           AFR.getAsDeclContext(),
                           AFR.isKnownNoEscape(),
-                          AFR.isObjC());
+                          AFR.isObjC(),
+                          isGeneric);
   AFR.getBody()->walk(finder);
 
   if (AFR.hasType() && !AFR.isObjC()) {
     finder.checkType(AFR.getType(), AFR.getLoc());
   }
 
-  // A generic function always captures outer generic parameters.
-  bool isGeneric = false;
-  auto *AFD = AFR.getAbstractFunctionDecl();
-  if (AFD)
-    isGeneric = AFD->isGeneric();
-
-  auto captures = finder.getCaptureInfo();
-  if (isGeneric)
-    captures.setGenericParamCaptures(true);
-  AFR.setCaptureInfo(captures);
+  AFR.setCaptureInfo(finder.getCaptureInfo());
 
   // Compute captures for default argument expressions.
   if (auto *AFD = AFR.getAbstractFunctionDecl()) {
@@ -621,7 +613,8 @@ void TypeChecker::computeCaptures(AnyFunctionRef AFR) {
                                 E->getLoc(),
                                 AFD,
                                 /*isNoEscape=*/false,
-                                /*isObjC=*/false);
+                                /*isObjC=*/false,
+                                /*IsGeneric*/isGeneric);
         E->walk(finder);
 
         if (!AFD->getDeclContext()->isLocalContext() &&
@@ -630,10 +623,7 @@ void TypeChecker::computeCaptures(AnyFunctionRef AFR) {
                                  diag::dynamic_self_default_arg);
         }
 
-        auto captures = finder.getCaptureInfo();
-        if (isGeneric)
-          captures.setGenericParamCaptures(true);
-        P->setDefaultArgumentCaptureInfo(captures);
+        P->setDefaultArgumentCaptureInfo(finder.getCaptureInfo());
       }
     }
   }
@@ -690,7 +680,8 @@ void TypeChecker::checkPatternBindingCaptures(NominalTypeDecl *typeDecl) {
                               init->getLoc(),
                               PBD->getInitContext(i),
                               /*NoEscape=*/false,
-                              /*ObjC=*/false);
+                              /*ObjC=*/false,
+                              /*IsGenericFunction*/false);
       init->walk(finder);
 
       if (finder.getDynamicSelfCaptureLoc().isValid() && !isLazy(PBD)) {