SILGen: Update for PackElementExpr captures

Author: slavapestov

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.

lib/SILGen/SILGenProlog.cpp

@@ -1052,6 +1052,49 @@ static void emitCaptureArguments(SILGenFunction &SGF,
                                  GenericSignature origGenericSig,
                                  CapturedValue capture,
                                  uint16_t ArgNo) {
+  if (auto *expr = capture.getPackElement()) {
+    SILLocation Loc(expr);
+    Loc.markAsPrologue();
+
+    auto interfaceType = expr->getType()->mapTypeOutOfContext();
+
+    auto type = SGF.F.mapTypeIntoContext(interfaceType);
+    auto &lowering = SGF.getTypeLowering(type);
+    SILType ty = lowering.getLoweredType();
+
+    SILValue arg;
+
+    auto expansion = SGF.getTypeExpansionContext();
+    auto captureKind = SGF.SGM.Types.getDeclCaptureKind(capture, expansion);
+    switch (captureKind) {
+    case CaptureKind::Constant:
+    case CaptureKind::StorageAddress:
+    case CaptureKind::Immutable: {
+      auto argIndex = SGF.F.begin()->getNumArguments();
+      // Non-escaping stored decls are captured as the address of the value.
+      auto param = SGF.F.getConventions().getParamInfoForSILArg(argIndex);
+      if (SGF.F.getConventions().isSILIndirect(param))
+        ty = ty.getAddressType();
+
+      auto *fArg = SGF.F.begin()->createFunctionArgument(ty, nullptr);
+      fArg->setClosureCapture(true);
+
+      arg = fArg;
+      break;
+    }
+
+    case CaptureKind::ImmutableBox:
+    case CaptureKind::Box:
+      llvm_unreachable("should be impossible");
+    }
+
+    ManagedValue mv = ManagedValue::forBorrowedRValue(arg);
+    auto inserted = SGF.OpaqueValues.insert(std::make_pair(expr, mv));
+    assert(inserted.second);
+    (void) inserted;
+
+    return;
+  }
 
   auto *VD = cast<VarDecl>(capture.getDecl());