AST: Introduce PackType::unwrapSingletonPackExpansion()

Author: slavapestov

include/swift/AST/Types.h

@@ -6845,6 +6845,8 @@ class PackType final : public TypeBase, public llvm::FoldingSetNode,
 
   bool containsPackExpansionType() const;
 
+  PackExpansionType *unwrapSingletonPackExpansion() const;
+
   CanTypeWrapper<PackType> getReducedShape();
 
 public:
@@ -6883,6 +6885,8 @@ BEGIN_CAN_TYPE_WRAPPER(PackType, Type)
   CanTypeArrayRef getElementTypes() const {
     return CanTypeArrayRef(getPointer()->getElementTypes());
   }
+
+  CanTypeWrapper<PackExpansionType> unwrapSingletonPackExpansion() const;
 END_CAN_TYPE_WRAPPER(PackType, Type)
 
 inline CanPackType CanTupleType::getInducedPackType() const {
@@ -6971,6 +6975,13 @@ BEGIN_CAN_TYPE_WRAPPER(PackExpansionType, Type)
   }
 END_CAN_TYPE_WRAPPER(PackExpansionType, Type)
 
+
+inline CanTypeWrapper<PackExpansionType>
+CanPackType::unwrapSingletonPackExpansion() const {
+  return CanPackExpansionType(
+      getPointer()->unwrapSingletonPackExpansion());
+}
+
 /// getASTContext - Return the ASTContext that this type belongs to.
 inline ASTContext &TypeBase::getASTContext() const {
   // If this type is canonical, it has the ASTContext in it.

lib/AST/ParameterPack.cpp

@@ -300,6 +300,18 @@ CanPackType CanPackType::getSingletonPackExpansion(CanType param) {
   return CanPackType(PackType::getSingletonPackExpansion(param));
 }
 
+PackExpansionType *PackType::unwrapSingletonPackExpansion() const {
+  if (getNumElements() == 1) {
+    if (auto expansion = getElementTypes()[0]->getAs<PackExpansionType>()) {
+      auto pattern = expansion->getPatternType();
+      if (pattern->isParameterPack() || pattern->is<PackArchetypeType>())
+        return expansion;
+    }
+  }
+
+  return nullptr;
+}
+
 bool SILPackType::containsPackExpansionType() const {
   for (auto type : getElementTypes()) {
     if (isa<PackExpansionType>(type))

lib/IRGen/Fulfillment.cpp

@@ -182,14 +182,12 @@ bool FulfillmentMap::searchTypeMetadata(IRGenModule &IGM, CanType type,
 static CanType getSingletonPackExpansionParameter(CanPackType packType,
                     const FulfillmentMap::InterestingKeysCallback &keys,
                     Optional<unsigned> &packExpansionComponent) {
-  if (packType->getNumElements() != 1)
-    return CanType();
-  auto expansion = dyn_cast<PackExpansionType>(packType.getElementType(0));
-  if (!expansion || !keys.isInterestingPackExpansion(expansion))
-    return CanType();
-
-  packExpansionComponent = 0;
-  return expansion.getPatternType();
+  if (auto expansion = packType.unwrapSingletonPackExpansion()) {
+    if (keys.isInterestingPackExpansion(expansion)) {
+      packExpansionComponent = 0;
+      return expansion.getPatternType();
+    }
+  }
 }
 
 bool FulfillmentMap::searchTypeMetadataPack(IRGenModule &IGM,
@@ -405,10 +403,7 @@ bool FulfillmentMap::searchShapeRequirement(IRGenModule &IGM, CanType argType,
   // there aren't expansions over pack parameters with different shapes,
   // we should always be able to turn this into the equation
   // `ax + b = <fulfilled count>` and then solve that.
-  if (packType->getNumElements() != 1)
-    return false;
-  auto expansion =
-    dyn_cast<PackExpansionType>(packType.getElementType(0));
+  auto expansion = packType.unwrapSingletonPackExpansion();
   if (!expansion)
     return false;
 

lib/IRGen/GenPack.cpp

@@ -37,15 +37,9 @@ using namespace irgen;
 
 static CanPackArchetypeType
 getForwardedPackArchetypeType(CanPackType packType) {
-  if (packType->getNumElements() != 1)
-    return CanPackArchetypeType();
-  auto uncastElement = packType.getElementType(0);
-  auto element = dyn_cast<PackExpansionType>(uncastElement);
-  if (!element)
-    return CanPackArchetypeType();
-  auto patternType = element.getPatternType();
-  auto packArchetype = dyn_cast<PackArchetypeType>(patternType);
-  return packArchetype;
+  if (auto expansion = packType.unwrapSingletonPackExpansion())
+    return dyn_cast<PackArchetypeType>(expansion.getPatternType());
+  return CanPackArchetypeType();
 }
 
 static MetadataResponse

lib/IRGen/GenProto.cpp

@@ -3730,13 +3730,11 @@ void irgen::bindGenericRequirement(IRGenFunction &IGF,
   // FIXME: Remove this
   bool wasUnwrappedPack = false;
   if (auto packType = dyn_cast<PackType>(type)) {
-    if (packType->getNumElements() == 1) {
-      auto eltType = packType.getElementType(0);
-      if (auto expansionType = dyn_cast<PackExpansionType>(eltType)) {
-        if (auto archetypeType = dyn_cast<PackArchetypeType>(expansionType.getPatternType())) {
-          type = archetypeType;
-          wasUnwrappedPack = true;
-        }
+    if (auto expansionType = packType.unwrapSingletonPackExpansion()) {
+      if (auto archetypeType = dyn_cast_or_null<PackArchetypeType>(
+            expansionType.getPatternType())) {
+        type = archetypeType;
+        wasUnwrappedPack = true;
       }
     }
   }