Merge pull request #83555 from rjmccall/variadic-witness-fixes

Fix some major SILGen bugs with pack handling:

Author: rjmccall

lib/SILGen/SILGenApply.cpp

@@ -523,7 +523,8 @@ class Callee {
 
     auto &ci = SGF.getConstantInfo(SGF.getTypeExpansionContext(), c);
     return Callee(
-        Kind::WitnessMethod, SGF, c, ci.FormalPattern, ci.FormalType,
+        Kind::WitnessMethod, SGF, c,
+        ci.FormalPattern.withSubstitutions(subs), ci.FormalType,
         substOpaqueTypesWithUnderlyingTypes(subs, SGF.getTypeExpansionContext()), l);
   }
   static Callee forDynamic(SILGenFunction &SGF,

lib/SILGen/SILGenPack.cpp

@@ -1131,8 +1131,9 @@ SILGenFunction::emitPackTransform(SILLocation loc,
       outputEltAddr = outputElt.forward(*this);
 
     // Otherwise, if the value is not already in the temporary, put it there.
-    } else if (!outputElt.isInContext()) {
-      outputElt.forwardInto(*this, loc, outputEltInit.get());
+    } else {
+      if (!outputElt.isInContext())
+        outputElt.forwardInto(*this, loc, outputEltInit.get());
       outputEltInit->getManagedAddress().forward(*this);
     }
 

lib/SILGen/SILGenPoly.cpp

@@ -2416,8 +2416,15 @@ class TranslateArguments : public ExpanderBase<TranslateArguments, ParamInfo> {
     return innerPackParam.getConvention();
   }
 
-  ParamInfo getInnerPackExpansionSlot(SILValue packAddr) {
-    return ParamInfo(IndirectSlot(packAddr), getInnerPackConvention());
+  ParamInfo getInnerPackExpansionSlot(ManagedValue packAddr) {
+    // Ignore the ownership of the pack address --- it'll be an
+    // l-value or r-value depending on what kind of argument we're
+    // producing, but there's no meaningful ownership yet, and in
+    // any case it won't have any cleanups attached.
+    assert(!packAddr.hasCleanup());
+
+    return ParamInfo(IndirectSlot(packAddr.getValue()),
+                     getInnerPackConvention());
   }
 
   /// Given an element of an inner pack that we're emitting into,
@@ -2770,15 +2777,29 @@ ManagedValue TranslateArguments::expandPackInnerParam(
         // cleanup onto the elements.
         auto outerComponent = outerParam.projectPackComponent(SGF, Loc);
 
-        // We can only do direct forwarding of of the pack elements in
-        // one very specific case right now.  That isn't great, but we
-        // have to live with it.
-        bool forwardouterToinner =
-          (outerExpansionTy.getPatternType()
-             == innerExpansionTy.getPatternType());
-
-        // The result of the transformation will be +1 unless we do that.
-        bool innerIsPlusOne = !forwardouterToinner;
+        bool patternTypesMatch =
+          outerExpansionTy.getPatternType() == innerExpansionTy.getPatternType();
+
+        auto innerPatternTypeIsTrivial =
+          SGF.getTypeProperties(innerExpansionTy.getPatternType()).isTrivial();
+        bool outerIsConsumed = outerComponent.hasCleanup();
+        bool innerIsConsumed =
+          isConsumedParameterInCaller(innerPackParam.getConvention());
+
+        // We can only do direct forwarding of of the pack elements (without
+        // needing temporary memory) if they have exactly the same type and
+        // we're not turning a borrowed parameter into a consuming one.
+        bool forwardOuterToInner =
+          (patternTypesMatch &&
+           (innerPatternTypeIsTrivial || outerIsConsumed || !innerIsConsumed));
+
+        // The result of the transformation function below will be +1 unless
+        // we directly forward that (or if direct forwarding produces an owned
+        // or trivial value).
+        bool innerIsPlusOne =
+          (!forwardOuterToInner ||
+           innerPatternTypeIsTrivial ||
+           outerIsConsumed);
 
         ManagedValue inner =
           SGF.emitPackTransform(Loc, outerComponent,
@@ -2787,13 +2808,18 @@ ManagedValue TranslateArguments::expandPackInnerParam(
                                 innerPackAddr,
                                 innerFormalPackType,
                                 innerComponentIndex,
-                                /*is trivial*/ forwardouterToinner,
+                                /*is simple projection*/ forwardOuterToInner,
                                 innerIsPlusOne,
             [&](ManagedValue outerEltAddr, SILType innerEltTy,
                 SGFContext ctxt) {
           // If we decided to just forward, we can do that now.
-          if (forwardouterToinner)
+          if (forwardOuterToInner) {
+            // It's okay to return an owned value here even if we're
+            // producing this for a borrowed parameter. We'll end up with
+            // an argument with cleanups, which in the end we just won't
+            // forward.
             return outerEltAddr;
+          }
 
           // Otherwise, map the subst pattern types into element context.
           CanType innerSubstEltType =
@@ -3454,8 +3480,8 @@ class ResultPlanner : public ExpanderBase<ResultPlanner, IndirectSlot> {
     return temporary;
   }
 
-  IndirectSlot getInnerPackExpansionSlot(SILValue packAddr) {
-    return IndirectSlot(packAddr);
+  IndirectSlot getInnerPackExpansionSlot(ManagedValue packAddr) {
+    return IndirectSlot(packAddr.getLValueAddress());
   }
 
   IndirectSlot getInnerPackElementSlot(SILType elementTy) {
@@ -3809,8 +3835,8 @@ void ExpanderBase<Impl, InnerSlotType>::expandParallelTuples(
       ManagedValue outerPackComponent =
         outerElt.projectPackComponent(SGF, Loc);
 
-      auto innerEltSlot = asImpl().getInnerPackExpansionSlot(
-        innerElt.getPackValue().getLValueAddress());
+      auto innerEltSlot =
+        asImpl().getInnerPackExpansionSlot(innerElt.getPackValue());
       ManagedValue innerPackComponent = asImpl().expandPackExpansion(
                                 innerElt.getOrigType(),
         cast<PackExpansionType>(innerElt.getSubstType()),
@@ -4590,8 +4616,8 @@ void ExpanderBase<Impl, InnerSlotType>::expandParallelTuplesOuterIndirect(
       continue;
     }
 
-    auto innerExpansionSlot = asImpl().getInnerPackExpansionSlot(
-      innerElt.getPackValue().getLValueAddress());
+    auto innerExpansionSlot =
+      asImpl().getInnerPackExpansionSlot(innerElt.getPackValue());
     asImpl().expandPackExpansion(innerElt.getOrigType(),
          cast<PackExpansionType>(innerElt.getSubstType()),
                                  outerElt.getOrigType(),

lib/SILGen/SILGenProlog.cpp

@@ -902,6 +902,8 @@ class ArgumentInitHelper {
     ManagedValue argrv = makeArgument(loc, pd);
     if (pd->isInOut()) {
       assert(argrv.getType().isAddress() && "expected inout to be address");
+    } else if (argrv.getType().is<SILPackType>()) {
+      assert(argrv.getType().isAddress() && "expected pack to be address");
     } else if (!pd->isImmutableInFunctionBody()) {
       // If it's a locally mutable parameter, then we need to move the argument
       // value into a local box to hold the mutated value.