Merge pull request #63557 from rjmccall/variadic-generic-sil-subst

rjmccall · web-flow · commit 912a4a8f6e31 · 2023-02-10T21:30:13.000-05:00
Test that SIL inlining works for the new variadic generics instructions
diff --git a/include/swift/AST/SubstitutionMap.h b/include/swift/AST/SubstitutionMap.h
@@ -169,7 +169,7 @@ class SubstitutionMap {
   bool isCanonical() const;
 
   /// Return the canonical form of this substitution map.
-  SubstitutionMap getCanonical() const;
+  SubstitutionMap getCanonical(bool canonicalizeSignature = true) const;
 
   /// Apply a substitution to all replacement types in the map. Does not
   /// change keys.
diff --git a/include/swift/SIL/SILCloner.h b/include/swift/SIL/SILCloner.h
@@ -204,7 +204,8 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
       }
     }
 
-    return asImpl().remapSubstitutionMap(Subs).getCanonical();
+    return asImpl().remapSubstitutionMap(Subs)
+                   .getCanonical(/*canonicalizeSignature*/false);
   }
 
   SILType getTypeInClonedContext(SILType Ty) {
@@ -261,14 +262,82 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
       if (auto origExpansionType =
             dyn_cast<PackExpansionType>(origComponentType)) {
         auto newShapeClass = getOpASTType(origExpansionType.getCountType());
-        newIndex += cast<PackType>(newShapeClass)->getNumElements();
+        if (auto newShapePack = dyn_cast<PackType>(newShapeClass))
+          newIndex += newShapePack->getNumElements();
+        else
+          newIndex++;
       } else {
         newIndex++;
       }
     }
     return newIndex;
   }
 
+  /// Does type substitution make the given tuple type no longer a tuple?
+  bool doesOpTupleDisappear(CanTupleType type) {
+    // Fast-path the empty tuple.
+    if (type->getNumElements() == 0) return false;
+
+    // Do a first pass over the tuple elements to check out the
+    // non-expansions.  If there's more than one of them, or any of them
+    // is labeled, we definitely stay a tuple and don't need to substitute
+    // any of the expansions.
+    unsigned numScalarElements = 0;
+    for (auto index : indices(type->getElements())) {
+      auto eltType = type.getElementType(index);
+      // Ignore pack expansions in this pass.
+      if (isa<PackExpansionType>(eltType)) continue;
+
+      // If there's a labeled scalar element, we'll stay a tuple.
+      if (type->getElement(index).hasName()) return false;
+
+      // If there are multiple scalar elements, we'll stay a tuple.
+      if (++numScalarElements > 1) return false;
+    }
+
+    assert(numScalarElements <= 1);
+
+    // We must have expansions if we got here: if all the elements were
+    // scalar, and none of them were labelled, and there wasn't more than
+    // one of them, and there was at least one of them, then somehow
+    // we had a tuple with a single unlabeled element.
+
+    // Okay, we need to substitute the count types for the expansions.
+    for (auto index : indices(type->getElements())) {
+      // Ignore non-expansions because we've already counted them.
+      auto expansion = dyn_cast<PackExpansionType>(type.getElementType(index));
+      if (!expansion) continue;
+
+      // Substitute the shape class of the expansion.
+      auto newShapeClass = getOpASTType(expansion.getCountType());
+      auto newShapePack = dyn_cast<PackType>(newShapeClass);
+
+      // If the element has a name, then the tuple sticks around unless
+      // the expansion disappears completely.
+      if (type->getElement(index).hasName()) {
+        if (newShapePack && newShapePack->getNumElements() == 0)
+          continue;
+        return false;
+      }
+
+      // Otherwise, walk the substituted shape components.
+      for (auto newShapeElement : newShapePack.getElementTypes()) {
+        // If there's an expansion in the shape, we'll have an expansion
+        // in the tuple elements, which forces the tuple structure to remain.
+        if (isa<PackExpansionType>(newShapeElement)) return false;
+
+        // Otherwise, add another scalar element.
+        if (++numScalarElements > 1) return false;
+      }
+    }
+
+    // All of the packs expanded to scalars.  We should've short-circuited
+    // if we ever saw a second or labeled scalar, so all we need to test
+    // is whether we have exactly one scalar.
+    assert(numScalarElements <= 1);
+    return numScalarElements == 1;
+  }
+
   void remapRootOpenedType(CanOpenedArchetypeType archetypeTy) {
     assert(archetypeTy->isRoot());
 
@@ -2492,13 +2561,15 @@ void SILCloner<ImplClass>::visitOpenPackElementInst(
   auto newContextSubs =
     getOpSubstitutionMap(origEnv->getPackElementContextSubstitutions());
 
-  // Substitute the shape class.
-  auto newShapeClass = getOpASTType(origEnv->getOpenedElementShapeClass());
+  // The opened shape class is a parameter of the original signature,
+  // which is unchanged.
+  auto openedShapeClass = origEnv->getOpenedElementShapeClass();
 
   // Build the new environment.
   auto newEnv =
     GenericEnvironment::forOpenedElement(origEnv->getGenericSignature(),
-                                         UUID::fromTime(), newShapeClass,
+                                         UUID::fromTime(),
+                                         openedShapeClass,
                                          newContextSubs);
 
   // Associate the old opened archetypes with the new ones.
@@ -2552,8 +2623,17 @@ void SILCloner<ImplClass>::visitTuplePackElementAddrInst(
   auto newIndex = getOpValue(Inst->getIndex());
   auto newTuple = getOpValue(Inst->getTuple());
   auto newElementType = getOpType(Inst->getElementType());
-  // FIXME: do we need to rewrite when substitution removes the
-  // tuple-ness of the type?  If so, what do we rewrite to?
+
+  // If the tuple-ness of the operand disappears due to substitution,
+  // replace this instruction with an unchecked_addr_cast.
+  // FIXME: use type_refine_addr instead
+  if (doesOpTupleDisappear(Inst->getTupleType())) {
+    recordClonedInstruction(
+        Inst, getBuilder().createUncheckedAddrCast(loc, newTuple,
+                                                   newElementType));
+    return;
+  }
+
   recordClonedInstruction(
       Inst, getBuilder().createTuplePackElementAddr(loc, newIndex, newTuple,
                                                     newElementType));
diff --git a/lib/AST/SubstitutionMap.cpp b/lib/AST/SubstitutionMap.cpp
@@ -164,10 +164,12 @@ bool SubstitutionMap::isCanonical() const {
   return true;
 }
 
-SubstitutionMap SubstitutionMap::getCanonical() const {
+SubstitutionMap SubstitutionMap::getCanonical(bool canonicalizeSignature) const {
   if (empty()) return *this;
 
-  auto canonicalSig = getGenericSignature().getCanonicalSignature();
+  auto sig = getGenericSignature();
+  if (canonicalizeSignature) sig = sig.getCanonicalSignature();
+
   SmallVector<Type, 4> replacementTypes;
   for (Type replacementType : getReplacementTypesBuffer()) {
     if (replacementType)
@@ -181,12 +183,11 @@ SubstitutionMap SubstitutionMap::getCanonical() const {
     conformances.push_back(conf.getCanonicalConformanceRef());
   }
 
-  return SubstitutionMap::get(canonicalSig,
+  return SubstitutionMap::get(sig,
                               ArrayRef<Type>(replacementTypes),
                               ArrayRef<ProtocolConformanceRef>(conformances));
 }
 
-
 SubstitutionMap SubstitutionMap::get(GenericSignature genericSig,
                                      SubstitutionMap substitutions) {
   if (!genericSig) {
diff --git a/lib/AST/Type.cpp b/lib/AST/Type.cpp
@@ -646,6 +646,21 @@ void TypeBase::getTypeVariables(
   }
 }
 
+static bool isLegalSILType(CanType type);
+
+static bool isLegalSILTypeOrPackExpansion(CanType type) {
+  // Pack expansions aren't legal in arbitrary positions in SIL;
+  // for example, we should never see a function parameter or result
+  // of pack-expansion type.  But they're allowed within SILPackTypes
+  // and SIL TupleTypes as long as their pattern types are legal.
+  // The count type should always be an AST type.
+  if (auto packExpansionType = dyn_cast<PackExpansionType>(type)) {
+    return isLegalSILType(packExpansionType.getPatternType());
+  }
+
+  return isLegalSILType(type);
+}
+
 static bool isLegalSILType(CanType type) {
   // L-values and inouts are not legal.
   if (!type->isMaterializable()) return false;
@@ -660,20 +675,14 @@ static bool isLegalSILType(CanType type) {
   // Tuples are legal if all their elements are legal.
   if (auto tupleType = dyn_cast<TupleType>(type)) {
     for (auto eltType : tupleType.getElementTypes()) {
-      if (!isLegalSILType(eltType)) return false;
+      if (!isLegalSILTypeOrPackExpansion(eltType)) return false;
     }
     return true;
   }
 
   // Packs must be lowered.
   if (isa<PackType>(type)) return false;
 
-  // Pack expansions are legal if all their pattern and count types are legal.
-  if (auto packExpansionType = dyn_cast<PackExpansionType>(type)) {
-    return (isLegalSILType(packExpansionType.getPatternType()) &&
-            isLegalSILType(packExpansionType.getCountType()));
-  }
-
   // Optionals are legal if their object type is legal.
   if (auto objectType = type.getOptionalObjectType()) {
     return isLegalSILType(objectType);
diff --git a/lib/SIL/IR/SILFunctionType.cpp b/lib/SIL/IR/SILFunctionType.cpp
@@ -3929,10 +3929,23 @@ class SILTypeSubstituter :
   CanGenericSignature Sig;
 
   struct PackExpansion {
-    CanType ExpandedShape;
+    /// The shape class of pack parameters that are expanded by this
+    /// expansion.  Set during construction and not changed.
+    CanType OrigShapeClass;
+
+    /// The count type of the pack expansion in the current lane of
+    /// expansion, if any.  Pack elements in this lane should be
+    /// expansions with this shape.
+    CanType SubstPackExpansionCount;
+
+    /// The index of the current lane of expansion.  Basic
+    /// substitution of pack parameters with the same shape as
+    /// OrigShapeClass should yield a pack, and lanewise
+    /// substitution should produce this element of that pack.
     unsigned Index;
 
-    PackExpansion(CanType shape) : ExpandedShape(shape), Index(0) {}
+    PackExpansion(CanType origShapeClass)
+      : OrigShapeClass(origShapeClass), Index(0) {}
   };
   SmallVector<PackExpansion> ActivePackExpansions;
 
@@ -4253,7 +4266,7 @@ class SILTypeSubstituter :
 
   CanType visitPackExpansionType(CanPackExpansionType origType) {
     CanType patternType = visit(origType.getPatternType());
-    CanType countType = visit(origType.getCountType());
+    CanType countType = substASTType(origType.getCountType());
 
     return CanType(PackExpansionType::get(patternType, countType));
   }
@@ -4264,21 +4277,24 @@ class SILTypeSubstituter :
     CanType origPatternType = origType.getPatternType();
 
     // Substitute the count type (as an AST type).
-    CanType substCountType = visitType(origCountType);
+    CanType substCountType = substASTType(origCountType);
 
     // If that produces a pack type, expand the pattern element-wise.
     if (auto substCountPackType = dyn_cast<PackType>(substCountType)) {
       // Set up for element-wise expansion.
       ActivePackExpansions.emplace_back(origCountType);
 
       for (CanType substCountEltType : substCountPackType.getElementTypes()) {
+        auto expansionType = dyn_cast<PackExpansionType>(substCountEltType);
+        ActivePackExpansions.back().SubstPackExpansionCount =
+          (expansionType ? expansionType.getCountType() : CanType());
+
         // Expand the pattern type in the element-wise context.
         CanType expandedType = visit(origPatternType);
 
         // Turn that into a pack expansion if appropriate for the
         // count element.
-        if (auto expansionType =
-              dyn_cast<PackExpansionType>(substCountEltType)) {
+        if (expansionType) {
           expandedType =
             CanPackExpansionType::get(expandedType,
                                       expansionType.getCountType());
@@ -4317,11 +4333,44 @@ class SILTypeSubstituter :
     SmallVector<TupleTypeElt, 8> substElts;
     substElts.reserve(origType->getNumElements());
     for (auto &origElt : origType->getElements()) {
-      auto substEltType = visit(CanType(origElt.getType()));
-      substElts.push_back(origElt.getWithType(substEltType));
+      CanType origEltType = CanType(origElt.getType());
+      if (auto origExpansion = dyn_cast<PackExpansionType>(origEltType)) {
+        bool first = true;
+        substPackExpansion(origExpansion, [&](CanType substEltType) {
+          auto substElt = origElt.getWithType(substEltType);
+          if (first) {
+            first = false;
+          } else {
+            substElt = substElt.getWithoutName();
+          }
+          substElts.push_back(substElt);
+        });
+      } else {
+        auto substEltType = visit(origEltType);
+        substElts.push_back(origElt.getWithType(substEltType));
+      }
     }
+
+    // Turn unlabeled singleton scalar tuples into their underlying types.
+    // The AST type substituter doesn't actually implement this rule yet,
+    // but we need to implement it in SIL in order to support testing,
+    // since the type parser can't parse a singleton tuple.
+    //
+    // For compatibility with previous behavior, don't do this if the
+    // original tuple type was singleton.  AutoDiff apparently really
+    // likes making singleton tuples.
+    if (isParenType(substElts) && !isParenType(origType->getElements()))
+      return CanType(substElts[0].getType());
+
     return CanType(TupleType::get(substElts, TC.Context));
   }
+
+  static bool isParenType(ArrayRef<TupleTypeElt> elts) {
+    return (elts.size() == 1 &&
+            !elts[0].hasName() &&
+            !isa<PackExpansionType>(CanType(elts[0].getType())));
+  }
+
   // Block storage types need to substitute their capture type by these same
   // rules.
   CanType visitSILBlockStorageType(CanSILBlockStorageType origType) {
@@ -4375,11 +4424,21 @@ class SILTypeSubstituter :
       [&](SubstitutableType *origType) -> Type {
         auto substType = Subst(origType);
         if (!substType) return substType;
-        auto substPackType = substType->getAs<PackType>();
+        auto substPackType = dyn_cast<PackType>(substType->getCanonicalType());
         if (!substPackType) return substType;
-        auto index = getPackExpansionIndex(CanType(origType));
-        if (!index) return substType;
-        return substPackType->getElementType(*index);
+        auto activeExpansion = getActivePackExpansion(CanType(origType));
+        if (!activeExpansion) return substType;
+        auto substEltType =
+          substPackType.getElementType(activeExpansion->Index);
+        auto substExpansion = dyn_cast<PackExpansionType>(substEltType);
+        assert((bool) substExpansion ==
+               (bool) activeExpansion->SubstPackExpansionCount);
+        if (substExpansion) {
+          assert(hasSameShape(substExpansion.getCountType(),
+                              activeExpansion->SubstPackExpansionCount));
+          return substExpansion.getPatternType();
+        }
+        return substEltType;
       },
       [&](CanType dependentType,
           Type conformingReplacementType,
@@ -4388,23 +4447,27 @@ class SILTypeSubstituter :
                                         conformingReplacementType,
                                         conformingProtocol);
         if (!conformance || !conformance.isPack()) return conformance;
-        auto index = getPackExpansionIndex(dependentType);
-        if (!index) return conformance;
+        auto activeExpansion = getActivePackExpansion(dependentType);
+        if (!activeExpansion) return conformance;
         auto pack = conformance.getPack();
-        return pack->getPatternConformances()[*index];
+        auto substEltConf =
+          pack->getPatternConformances()[activeExpansion->Index];
+        // There isn't currently a ProtocolConformanceExpansion that
+        // we would need to look through here.
+        return substEltConf;
       },
       substOptions)->getCanonicalType();
   }
 
-  Optional<size_t> getPackExpansionIndex(CanType dependentType) {
+  PackExpansion *getActivePackExpansion(CanType dependentType) {
     // We push new expansions onto the end of this vector, and we
     // want to honor the innermost expansion, so we have to traverse
     // in it reverse.
     for (auto &entry : reverse(ActivePackExpansions)) {
-      if (hasSameShape(dependentType, entry.ExpandedShape))
-        return entry.Index;
+      if (hasSameShape(dependentType, entry.OrigShapeClass))
+        return &entry;
     }
-    return None;
+    return nullptr;
   }
 
   bool hasSameShape(CanType lhs, CanType rhs) {
diff --git a/test/SILOptimizer/variadic_generics.sil b/test/SILOptimizer/variadic_generics.sil
