Merge pull request #65735 from slavapestov/param-pack-captures

SILGen: Wrap captures of parameter packs inside tuples

Author: slavapestov

lib/SIL/IR/SILFunctionType.cpp

@@ -1929,6 +1929,16 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
     auto type = VD->getInterfaceType();
     auto canType = type->getReducedType(origGenericSig);
 
+    // If we're capturing a parameter pack, wrap it in a tuple.
+    if (isa<PackExpansionType>(canType)) {
+      assert(!cast<ParamDecl>(VD)->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);

lib/SIL/Verifier/SILVerifier.cpp

@@ -675,6 +675,7 @@ struct ImmutableAddressUseVerifier {
       case SILInstructionKind::TailAddrInst:
       case SILInstructionKind::IndexRawPointerInst:
       case SILInstructionKind::MarkMustCheckInst:
+      case SILInstructionKind::PackElementGetInst:
         // Add these to our worklist.
         for (auto result : inst->getResults()) {
           llvm::copy(result->getUses(), std::back_inserter(worklist));
@@ -693,6 +694,25 @@ struct ImmutableAddressUseVerifier {
         llvm_unreachable("invoking standard assertion failure");
         break;
       }
+      case SILInstructionKind::TuplePackElementAddrInst: {
+        if (&cast<TuplePackElementAddrInst>(inst)->getOperandRef(
+              TuplePackElementAddrInst::TupleOperand) == use) {
+          for (auto result : inst->getResults()) {
+            llvm::copy(result->getUses(), std::back_inserter(worklist));
+          }
+
+          break;
+        }
+
+        return false;
+      }
+      case SILInstructionKind::PackElementSetInst: {
+        if (&cast<PackElementSetInst>(inst)->getOperandRef(
+              PackElementSetInst::PackOperand) == use)
+          return true;
+
+        return false;
+      }
       default:
         llvm::errs() << "Unhandled, unexpected instruction: " << *inst;
         llvm_unreachable("invoking standard assertion failure");

lib/SILGen/SILGenFunction.cpp

@@ -505,10 +505,24 @@ void SILGenFunction::emitCaptures(SILLocation loc,
     }
 
     auto *vd = cast<VarDecl>(capture.getDecl());
-    auto type = FunctionDC->mapTypeIntoContext(
-      vd->getInterfaceType());
+
+    auto interfaceType = vd->getInterfaceType();
+
+    bool isPack = false;
+    if (interfaceType->is<PackExpansionType>()) {
+      assert(!vd->supportsMutation() &&
+             "Cannot capture a pack as an lvalue");
+
+      SmallVector<TupleTypeElt, 1> elts;
+      elts.push_back(interfaceType);
+      interfaceType = TupleType::get(elts, getASTContext());
+
+      isPack = true;
+    }
+
+    auto type = FunctionDC->mapTypeIntoContext(interfaceType);
     auto valueType = FunctionDC->mapTypeIntoContext(
-      vd->getValueInterfaceType());
+      interfaceType->getReferenceStorageReferent());
 
     //
     // If we haven't emitted the captured value yet, we're forming a closure
@@ -571,8 +585,7 @@ void SILGenFunction::emitCaptures(SILLocation loc,
 
     // Get an address value for a SILValue if it is address only in an type
     // expansion context without opaque archetype substitution.
-    auto getAddressValue = [&](VarLoc entryVarLoc) -> SILValue {
-      SILValue entryValue = entryVarLoc.value;
+    auto getAddressValue = [&](SILValue entryValue, bool forceCopy) -> SILValue {
       if (SGM.M.useLoweredAddresses()
           && SGM.Types
                  .getTypeLowering(
@@ -582,32 +595,62 @@ void SILGenFunction::emitCaptures(SILLocation loc,
                  .isAddressOnly()
           && !entryValue->getType().isAddress()) {
 
+        assert(!isPack);
+
         auto addr = emitTemporaryAllocation(vd, entryValue->getType(), false,
                                             false, /*generateDebugInfo*/ false);
         auto val = B.emitCopyValueOperation(loc, entryValue);
         auto &lowering = getTypeLowering(entryValue->getType());
         lowering.emitStore(B, loc, val, addr, StoreOwnershipQualifier::Init);
-        entryValue = addr;
-        enterDestroyCleanup(addr);
+
+        if (!forceCopy)
+          enterDestroyCleanup(addr);
+        return addr;
+
+      } else if (isPack) {
+        SILType ty = getLoweredType(valueType).getObjectType();
+        auto addr = B.createAllocStack(loc, ty);
+        enterDeallocStackCleanup(addr);
+
+        auto formalPackType = cast<TupleType>(valueType->getCanonicalType())
+            .getInducedPackType();
+        copyPackElementsToTuple(loc, addr, entryValue, formalPackType);
+
+        if (!forceCopy)
+          enterDestroyCleanup(addr);
+        return addr;
+
+      } else if (forceCopy) {
+        // We cannot pass a valid SILDebugVariable while creating the temp here
+        // See rdar://60425582
+        auto addr = B.createAllocStack(loc, entryValue->getType().getObjectType());
+        enterDeallocStackCleanup(addr);
+        B.createCopyAddr(loc, entryValue, addr, IsNotTake, IsInitialization);
+        return addr;
+
+      } else {
+        return entryValue;
       }
-      return entryValue;
     };
 
     auto Entry = found->second;
+    auto val = Entry.value;
+
     switch (SGM.Types.getDeclCaptureKind(capture, expansion)) {
     case CaptureKind::Constant: {
+      assert(!isPack);
+
       // let declarations.
       auto &tl = getTypeLowering(valueType);
-      SILValue Val = Entry.value;
       bool eliminateMoveOnlyWrapper =
-          Val->getType().isMoveOnlyWrapped() &&
-          !vd->getInterfaceType()->is<SILMoveOnlyWrappedType>();
+          val->getType().isMoveOnlyWrapped() &&
+          !interfaceType->is<SILMoveOnlyWrappedType>();
 
-      if (!Val->getType().isAddress()) {
+      if (!val->getType().isAddress()) {
         // Our 'let' binding can guarantee the lifetime for the callee,
         // if we don't need to do anything more to it.
         if (canGuarantee && !vd->getInterfaceType()->is<ReferenceStorageType>()) {
-          auto guaranteed = ManagedValue::forUnmanaged(Val).borrow(*this, loc);
+          auto guaranteed = ManagedValue::forUnmanaged(val).borrow(*this, loc);
           if (eliminateMoveOnlyWrapper)
             guaranteed = B.createGuaranteedMoveOnlyWrapperToCopyableValue(
                 loc, guaranteed);
@@ -616,64 +659,63 @@ void SILGenFunction::emitCaptures(SILLocation loc,
         }
 
         // Just copy a by-val let.
-        Val = B.emitCopyValueOperation(loc, Val);
+        val = B.emitCopyValueOperation(loc, val);
         // If we need to unwrap a moveonlywrapped value, do so now but in an
         // owned way to ensure that the partial apply is viewed as a semantic
         // use of the value.
         if (eliminateMoveOnlyWrapper)
-          Val = B.createOwnedMoveOnlyWrapperToCopyableValue(loc, Val);
+          val = B.createOwnedMoveOnlyWrapperToCopyableValue(loc, val);
       } else {
         // If we have a mutable binding for a 'let', such as 'self' in an
         // 'init' method, load it.
-        if (Val->getType().isMoveOnly()) {
-          Val = B.createMarkMustCheckInst(
-              loc, Val,
+        if (val->getType().isMoveOnly()) {
+          val = B.createMarkMustCheckInst(
+              loc, val,
               MarkMustCheckInst::CheckKind::AssignableButNotConsumable);
         }
-        Val = emitLoad(loc, Val, tl, SGFContext(), IsNotTake).forward(*this);
+        val = emitLoad(loc, val, tl, SGFContext(), IsNotTake).forward(*this);
       }
 
       // If we're capturing an unowned pointer by value, we will have just
       // loaded it into a normal retained class pointer, but we capture it as
       // an unowned pointer.  Convert back now.
-      if (vd->getInterfaceType()->is<ReferenceStorageType>())
-        Val = emitConversionFromSemanticValue(loc, Val, getLoweredType(type));
+      if (interfaceType->is<ReferenceStorageType>())
+        val = emitConversionFromSemanticValue(loc, val, getLoweredType(type));
 
-      capturedArgs.push_back(emitManagedRValueWithCleanup(Val));
+      capturedArgs.push_back(emitManagedRValueWithCleanup(val));
       break;
     }
     case CaptureKind::Immutable: {
       if (canGuarantee) {
         // No-escaping stored declarations are captured as the
         // address of the value.
-        auto entryValue = getAddressValue(Entry);
-        capturedArgs.push_back(ManagedValue::forBorrowedRValue(entryValue));
+        auto addr = getAddressValue(val, /*forceCopy=*/false);
+        capturedArgs.push_back(ManagedValue::forBorrowedRValue(addr));
       }
       else if (!silConv.useLoweredAddresses()) {
         capturedArgs.push_back(
-          B.createCopyValue(loc, ManagedValue::forUnmanaged(Entry.value)));
+          B.createCopyValue(loc, ManagedValue::forUnmanaged(val)));
       } else {
-        auto entryValue = getAddressValue(Entry);
-        // We cannot pass a valid SILDebugVariable while creating the temp here
-        // See rdar://60425582
-        auto addr = B.createAllocStack(loc, entryValue->getType().getObjectType());
-        enterDeallocStackCleanup(addr);
-        B.createCopyAddr(loc, entryValue, addr, IsNotTake, IsInitialization);
+        auto addr = getAddressValue(val, /*forceCopy=*/true);
         capturedArgs.push_back(ManagedValue::forLValue(addr));
       }
       break;
     }
     case CaptureKind::StorageAddress: {
-      auto entryValue = getAddressValue(Entry);
+      assert(!isPack);
+
+      auto addr = getAddressValue(val, /*forceCopy=*/false);
       // No-escaping stored declarations are captured as the
       // address of the value.
-      assert(entryValue->getType().isAddress() && "no address for captured var!");
-      capturedArgs.push_back(ManagedValue::forLValue(entryValue));
+      assert(addr->getType().isAddress() && "no address for captured var!");
+      capturedArgs.push_back(ManagedValue::forLValue(addr));
       break;
     }
 
     case CaptureKind::Box: {
-      auto entryValue = getAddressValue(Entry);
+      assert(!isPack);
+
+      auto entryValue = getAddressValue(val, /*forceCopy=*/false);
       // LValues are captured as both the box owning the value and the
       // address of the value.
       assert(entryValue->getType().isAddress() && "no address for captured var!");
@@ -722,7 +764,9 @@ void SILGenFunction::emitCaptures(SILLocation loc,
       break;
     }
     case CaptureKind::ImmutableBox: {
-      auto entryValue = getAddressValue(Entry);
+      assert(!isPack);
+
+      auto entryValue = getAddressValue(val, /*forceCopy=*/false);
       // LValues are captured as both the box owning the value and the
       // address of the value.
       assert(entryValue->getType().isAddress() &&

lib/SILGen/SILGenFunction.h

@@ -438,24 +438,6 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
       Result.box = box;
       return Result;
     }
-
-    static VarLoc getForBox(SILValue box) {
-      VarLoc Result;
-      Result.value = SILValue();
-      Result.box = box;
-      return Result;
-    }
-
-    /// Return either the value if we have one or if we only have a box, project
-    /// our a new box address and return that.
-    SILValue getValueOrBoxedValue(
-        SILGenFunction &SGF,
-        SILLocation loc = RegularLocation::getAutoGeneratedLocation()) {
-      if (value)
-        return value;
-      assert(box);
-      return SGF.B.createProjectBox(loc, box, 0);
-    }
   };
 
   /// VarLocs - Entries in this map are generated when a PatternBindingDecl is
@@ -2503,6 +2485,16 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                                            unsigned componentIndex,
                                            SILValue currentIndexWithinComponent);
 
+  /// Copy the elements of a pack, which must consist of a single pack expansion,
+  /// into a tuple value having the same pack expansion and its sole element type.
+  void copyPackElementsToTuple(SILLocation loc, SILValue tupleAddr, SILValue pack,
+                               CanPackType formalPackType);
+
+  /// Initialize a pack with the addresses of the elements of a tuple, which must
+  /// consist of a single pack expansion.
+  void projectTupleElementsToPack(SILLocation loc, SILValue tupleAddr, SILValue pack,
+                                  CanPackType formalPackType);
+
   /// Return an owned managed value for \p value that is cleaned up using an end_lifetime instruction.
   ///
   /// The end_lifetime cleanup is not placed into the ManagedValue itself and

lib/SILGen/SILGenPack.cpp

@@ -521,6 +521,49 @@ void SILGenFunction::emitPartialDestroyRemainingTuple(SILLocation loc,
   });
 }
 
+void SILGenFunction::copyPackElementsToTuple(SILLocation loc,
+                                             SILValue tupleAddr,
+                                             SILValue pack,
+                                             CanPackType formalPackType) {
+  auto pair = createOpenedElementValueEnvironment(
+    tupleAddr->getType().getTupleElementType(/*componentIndex=*/0));
+  auto elementEnv = pair.first;
+  auto elementTy = pair.second;
+
+  emitDynamicPackLoop(
+    loc, formalPackType, /*componentIndex=*/0, elementEnv,
+    [&](SILValue indexWithinComponent,
+        SILValue packExpansionIndex,
+        SILValue packIndex) {
+      auto packEltAddr = B.createPackElementGet(
+          loc, packIndex, pack, elementTy);
+      auto tupleEltAddr = B.createTuplePackElementAddr(
+          loc, packIndex, tupleAddr, elementTy);
+      B.createCopyAddr(loc, packEltAddr, tupleEltAddr,
+                       IsNotTake, IsInitialization);
+  });
+}
+
+void SILGenFunction::projectTupleElementsToPack(SILLocation loc,
+                                                SILValue tupleAddr,
+                                                SILValue pack,
+                                                CanPackType formalPackType) {
+  auto pair = createOpenedElementValueEnvironment(
+    tupleAddr->getType().getTupleElementType(/*componentIndex=*/0));
+  auto elementEnv = pair.first;
+  auto elementTy = pair.second;
+
+  emitDynamicPackLoop(
+    loc, formalPackType, /*componentIndex=*/0, elementEnv,
+    [&](SILValue indexWithinComponent,
+        SILValue packExpansionIndex,
+        SILValue packIndex) {
+      auto tupleEltAddr = B.createTuplePackElementAddr(
+          loc, packIndex, tupleAddr, elementTy);
+      B.createPackElementSet(loc, tupleEltAddr, packIndex, pack);
+  });
+}
+
 void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
                                          CanPackType formalPackType,
                                          unsigned componentIndex,