Merge pull request #74020 from slavapestov/pack-expansion-closures-part-7

Pack expansion closures part 7

Author: slavapestov

include/swift/AST/CaptureInfo.h

@@ -41,9 +41,12 @@ template <> struct DenseMapInfo<swift::CapturedValue>;
 namespace swift {
 class ValueDecl;
 class FuncDecl;
+class Expr;
 class OpaqueValueExpr;
+class PackElementExpr;
 class VarDecl;
 class GenericEnvironment;
+class Type;
 
 /// CapturedValue includes both the declaration being captured, along with flags
 /// that indicate how it is captured.
@@ -52,7 +55,7 @@ class CapturedValue {
 
 public:
   using Storage =
-      llvm::PointerIntPair<llvm::PointerUnion<ValueDecl*, OpaqueValueExpr*>, 2,
+      llvm::PointerIntPair<llvm::PointerUnion<ValueDecl *, Expr *>, 2,
                            unsigned>;
 
 private:
@@ -78,15 +81,7 @@ class CapturedValue {
   CapturedValue(ValueDecl *Val, unsigned Flags, SourceLoc Loc)
       : Value(Val, Flags), Loc(Loc) {}
 
-private:
-  // This is only used in TypeLowering when forming Lowered Capture
-  // Info. OpaqueValueExpr captured value should never show up in the AST
-  // itself.
-  //
-  // NOTE: AbstractClosureExpr::getIsolationCrossing relies upon this and
-  // asserts that it never sees one of these.
-  explicit CapturedValue(OpaqueValueExpr *Val, unsigned Flags)
-      : Value(Val, Flags), Loc(SourceLoc()) {}
+  CapturedValue(Expr *Val, unsigned Flags);
 
 public:
   static CapturedValue getDynamicSelfMetadata() {
@@ -97,36 +92,38 @@ class CapturedValue {
   bool isNoEscape() const { return Value.getInt() & IsNoEscape; }
 
   bool isDynamicSelfMetadata() const { return !Value.getPointer(); }
-  bool isOpaqueValue() const {
-    return Value.getPointer().is<OpaqueValueExpr *>();
+
+  bool isExpr() const {
+    return Value.getPointer().dyn_cast<Expr *>();
   }
 
+  bool isPackElement() const;
+  bool isOpaqueValue() const;
+
   /// Returns true if this captured value is a local capture.
   ///
   /// NOTE: This implies that the value is not dynamic self metadata, since
   /// values with decls are the only values that are able to be local captures.
   bool isLocalCapture() const;
 
-  CapturedValue mergeFlags(CapturedValue cv) {
-    assert(Value.getPointer() == cv.Value.getPointer() &&
-           "merging flags on two different value decls");
-    return CapturedValue(
-        Storage(Value.getPointer(), getFlags() & cv.getFlags()),
-        Loc);
+  CapturedValue mergeFlags(unsigned flags) const {
+    return CapturedValue(Storage(Value.getPointer(), getFlags() & flags), Loc);
   }
 
   ValueDecl *getDecl() const {
-    assert(Value.getPointer() && "dynamic Self metadata capture does not "
-           "have a value");
     return Value.getPointer().dyn_cast<ValueDecl *>();
   }
 
-  OpaqueValueExpr *getOpaqueValue() const {
-    assert(Value.getPointer() && "dynamic Self metadata capture does not "
-           "have a value");
-    return Value.getPointer().dyn_cast<OpaqueValueExpr *>();
+  Expr *getExpr() const {
+    return Value.getPointer().dyn_cast<Expr *>();
   }
 
+  OpaqueValueExpr *getOpaqueValue() const;
+
+  PackElementExpr *getPackElement() const;
+
+  Type getPackElementType() const;
+
   SourceLoc getLoc() const { return Loc; }
 
   unsigned getFlags() const { return Value.getInt(); }

lib/AST/CaptureInfo.cpp

@@ -13,11 +13,36 @@
 #include "swift/AST/CaptureInfo.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
 
+CapturedValue::CapturedValue(Expr *Val, unsigned Flags)
+    : Value(Val, Flags), Loc(SourceLoc()) {
+  assert(isa<OpaqueValueExpr>(Val) || isa<PackElementExpr>(Val));
+}
+
+bool CapturedValue::isPackElement() const {
+  return isExpr() && isa<PackElementExpr>(getExpr());
+}
+bool CapturedValue::isOpaqueValue() const {
+  return isExpr() && isa<OpaqueValueExpr>(getExpr());
+}
+
+OpaqueValueExpr *CapturedValue::getOpaqueValue() const {
+  return dyn_cast_or_null<OpaqueValueExpr>(getExpr());
+}
+
+PackElementExpr *CapturedValue::getPackElement() const {
+  return dyn_cast_or_null<PackElementExpr>(getExpr());
+}
+
+Type CapturedValue::getPackElementType() const {
+  return getPackElement()->getType();
+}
+
 ArrayRef<CapturedValue>
 CaptureInfo::CaptureInfoStorage::getCaptures() const {
   return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), NumCapturedValues);
@@ -126,7 +151,18 @@ void CaptureInfo::print(raw_ostream &OS) const {
 
   interleave(getCaptures(),
              [&](const CapturedValue &capture) {
-               OS << capture.getDecl()->getBaseName();
+               if (capture.getDecl())
+                 OS << capture.getDecl()->getBaseName();
+               else if (capture.isPackElement()) {
+                 OS << "[pack element] ";
+                 capture.getPackElement()->dump(OS);
+               } else if (capture.isOpaqueValue()) {
+                 OS << "[opaque] ";
+                 capture.getOpaqueValue()->dump(OS);
+               } else {
+                 OS << "[unknown] ";
+                 assert(false);
+               }
 
                if (capture.isDirect())
                  OS << "<direct>";

lib/SIL/IR/SILFunctionType.cpp

@@ -2007,33 +2007,39 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
       continue;
     }
 
-    auto *varDecl = cast<VarDecl>(capture.getDecl());
+    auto options = SILParameterInfo::Options();
+
+    Type type;
+    VarDecl *varDecl = nullptr;
+    if (auto *expr = capture.getPackElement()) {
+      type = expr->getType();
+    } else {
+      varDecl = cast<VarDecl>(capture.getDecl());
+      type = varDecl->getTypeInContext();
+
+      // If we're capturing a parameter pack, wrap it in a tuple.
+      if (type->is<PackExpansionType>()) {
+        assert(!cast<ParamDecl>(varDecl)->supportsMutation() &&
+               "Cannot capture a pack as an lvalue");
+
+        SmallVector<TupleTypeElt, 1> elts;
+        elts.push_back(type);
+        type = TupleType::get(elts, TC.Context);
+      }
+
+      if (isolatedParam == varDecl) {
+        options |= SILParameterInfo::Isolated;
+        isolatedParam = nullptr;
+      }
+    }
 
-    auto type = varDecl->getTypeInContext();
     assert(!type->hasLocalArchetype() ||
            (genericSig && origGenericSig &&
             !genericSig->isEqual(origGenericSig)));
     type = mapTypeOutOfContext(type);
 
     auto canType = type->getReducedType(
         genericSig ? genericSig : origGenericSig);
-
-    auto options = SILParameterInfo::Options();
-    if (isolatedParam == varDecl) {
-      options |= SILParameterInfo::Isolated;
-      isolatedParam = nullptr;
-    }
-
-    // If we're capturing a parameter pack, wrap it in a tuple.
-    if (isa<PackExpansionType>(canType)) {
-      assert(!cast<ParamDecl>(varDecl)->supportsMutation() &&
-             "Cannot capture a pack as an lvalue");
-
-      SmallVector<TupleTypeElt, 1> elts;
-      elts.push_back(canType);
-      canType = CanTupleType(TupleType::get(elts, TC.Context));
-    }
-
     auto &loweredTL =
         TC.getTypeLowering(AbstractionPattern(genericSig, canType), canType,
                            expansion);
@@ -2055,6 +2061,8 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
       break;
     }
     case CaptureKind::Box: {
+      assert(varDecl);
+
       // The type in the box is lowered in the minimal context.
       auto minimalLoweredTy =
           TC.getTypeLowering(AbstractionPattern(genericSig, canType), canType,
@@ -2072,6 +2080,8 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
       break;
     }
     case CaptureKind::ImmutableBox: {
+      assert(varDecl);
+
       // The type in the box is lowered in the minimal context.
       auto minimalLoweredTy =
           TC.getTypeLowering(AbstractionPattern(genericSig, canType), canType,

lib/SIL/IR/TypeLowering.cpp

@@ -118,6 +118,19 @@ static bool hasSingletonMetatype(CanType instanceType) {
 
 CaptureKind TypeConverter::getDeclCaptureKind(CapturedValue capture,
                                               TypeExpansionContext expansion) {
+  if (auto *expr = capture.getPackElement()) {
+    auto contextTy = expr->getType();
+    auto &lowering = getTypeLowering(
+        contextTy, TypeExpansionContext::noOpaqueTypeArchetypesSubstitution(
+                            expansion.getResilienceExpansion()));
+
+    assert(!contextTy->isNoncopyable() && "Not implemented");
+    if (!lowering.isAddressOnly())
+      return CaptureKind::Constant;
+
+    return CaptureKind::Immutable;
+  }
+
   auto decl = capture.getDecl();
   auto *var = cast<VarDecl>(decl);
   assert(var->hasStorage() &&
@@ -4241,7 +4254,10 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
 
   // Recursively collect transitive captures from captured local functions.
   llvm::DenseSet<AnyFunctionRef> visitedFunctions;
-  llvm::MapVector<ValueDecl*,CapturedValue> captures;
+
+  // FIXME: CapturedValue should just be a hash key
+  llvm::MapVector<VarDecl *, CapturedValue> varCaptures;
+  llvm::MapVector<PackElementExpr *, CapturedValue> packElementCaptures;
 
   // If there is a capture of 'self' with dynamic 'Self' type, it goes last so
   // that IRGen can pass dynamic 'Self' metadata.
@@ -4259,12 +4275,23 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   std::function<void (SILDeclRef)> collectConstantCaptures;
 
   auto recordCapture = [&](CapturedValue capture) {
-    ValueDecl *value = capture.getDecl();
-    auto existing = captures.find(value);
-    if (existing != captures.end()) {
-      existing->second = existing->second.mergeFlags(capture);
+    if (auto *expr = capture.getPackElement()) {
+      auto existing = packElementCaptures.find(expr);
+      if (existing != packElementCaptures.end()) {
+        existing->second = existing->second.mergeFlags(capture.getFlags());
+      } else {
+        packElementCaptures.insert(std::pair<PackElementExpr *, CapturedValue>(
+          expr, capture));
+      }
     } else {
-      captures.insert(std::pair<ValueDecl *, CapturedValue>(value, capture));
+      VarDecl *value = cast<VarDecl>(capture.getDecl());
+      auto existing = varCaptures.find(value);
+      if (existing != varCaptures.end()) {
+        existing->second = existing->second.mergeFlags(capture.getFlags());
+      } else {
+        varCaptures.insert(std::pair<VarDecl *, CapturedValue>(
+          value, capture));
+      }
     }
   };
 
@@ -4284,6 +4311,11 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
     }
 
     for (auto capture : captureInfo.getCaptures()) {
+      if (capture.isPackElement()) {
+        recordCapture(capture);
+        continue;
+      }
+
       if (!capture.isLocalCapture())
         continue;
 
@@ -4398,7 +4430,7 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
             // If we've already captured the same value already, just merge
             // flags.
             if (selfCapture && selfCapture->getDecl() == capture.getDecl()) {
-              selfCapture = selfCapture->mergeFlags(capture);
+              selfCapture = selfCapture->mergeFlags(capture.getFlags());
               continue;
 
             // Otherwise, record the canonical self capture. It will appear
@@ -4492,7 +4524,10 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   collectConstantCaptures(fn);
 
   SmallVector<CapturedValue, 4> resultingCaptures;
-  for (auto capturePair : captures) {
+  for (auto capturePair : varCaptures) {
+    resultingCaptures.push_back(capturePair.second);
+  }
+  for (auto capturePair : packElementCaptures) {
     resultingCaptures.push_back(capturePair.second);
   }
 
@@ -4511,7 +4546,7 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
     resultingCaptures.push_back(*selfCapture);
   }
 
-  // Cache the uniqued set of transitive captures.
+  // Cache the result.
   CaptureInfo info(Context, resultingCaptures,
                    capturesDynamicSelf, capturesOpaqueValue,
                    capturesGenericParams, genericEnv.getArrayRef());

lib/SILGen/SILGenExpr.cpp

@@ -1667,6 +1667,14 @@ RValueEmitter::visitPackExpansionExpr(PackExpansionExpr *E,
 
 RValue
 RValueEmitter::visitPackElementExpr(PackElementExpr *E, SGFContext C) {
+  // If this is a captured pack element reference, just emit the parameter value
+  // that was passed to the closure.
+  auto found = SGF.OpaqueValues.find(E);
+  if (found != SGF.OpaqueValues.end())
+    return RValue(SGF, E, SGF.manageOpaqueValue(found->second, E, C));
+
+  // Otherwise, we're going to project the address of an element from the pack
+  // itself.
   FormalEvaluationScope scope(SGF);
 
   LValue lv = SGF.emitLValue(E, SGFAccessKind::OwnedObjectRead);
@@ -5972,9 +5980,9 @@ RValue RValueEmitter::visitMakeTemporarilyEscapableExpr(
 }
 
 RValue RValueEmitter::visitOpaqueValueExpr(OpaqueValueExpr *E, SGFContext C) {
-  assert(SGF.OpaqueValues.count(E) && "Didn't bind OpaqueValueExpr");
-  auto value = SGF.OpaqueValues[E];
-  return RValue(SGF, E, SGF.manageOpaqueValue(value, E, C));
+  auto found = SGF.OpaqueValues.find(E);
+  assert(found != SGF.OpaqueValues.end());
+  return RValue(SGF, E, SGF.manageOpaqueValue(found->second, E, C));
 }
 
 RValue RValueEmitter::visitPropertyWrapperValuePlaceholderExpr(

lib/SILGen/SILGenFunction.cpp

@@ -569,10 +569,9 @@ void SILGenFunction::emitCaptures(SILLocation loc,
       continue;
     }
 
-    if (capture.isOpaqueValue()) {
-      OpaqueValueExpr *opaqueValue = capture.getOpaqueValue();
+    if (capture.isOpaqueValue() || capture.isPackElement()) {
       capturedArgs.push_back(
-          emitRValueAsSingleValue(opaqueValue).ensurePlusOne(*this, loc));
+          emitRValueAsSingleValue(capture.getExpr()).ensurePlusOne(*this, loc));
       continue;
     }
 

lib/SILGen/SILGenFunction.h

@@ -2185,9 +2185,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
     emitOpenExistentialExprImpl(e, emitSubExpr);
   }
 
-  /// Mapping from active opaque value expressions to their values.
-  llvm::SmallDenseMap<OpaqueValueExpr *, ManagedValue>
-    OpaqueValues;
+  /// Mapping from OpaqueValueExpr/PackElementExpr to their values.
+  llvm::SmallDenseMap<Expr *, ManagedValue> OpaqueValues;
 
   /// A mapping from opaque value expressions to the open-existential
   /// expression that determines them, used while lowering lvalues.