AST: Introduce {TypeAliasType,BoundGenericType}::getExpandedGenericArgsPack()

Author: slavapestov

include/swift/AST/Types.h

@@ -71,6 +71,7 @@ class Identifier;
 class InOutType;
 class OpaqueTypeDecl;
 class OpenedArchetypeType;
+class PackType;
 class PlaceholderTypeRepr;
 enum class ReferenceCounting : uint8_t;
 enum class ResilienceExpansion : unsigned;
@@ -2002,6 +2003,8 @@ class TypeAliasType final
   /// this type references.
   ArrayRef<Type> getDirectGenericArgs() const;
 
+  PackType *getExpandedGenericArgsPack();
+
   // Support for FoldingSet.
   void Profile(llvm::FoldingSetNodeID &id) const;
 
@@ -2420,6 +2423,8 @@ class BoundGenericType : public NominalOrBoundGenericNominalType,
     return {getTrailingObjectsPointer(), Bits.BoundGenericType.GenericArgCount};
   }
 
+  PackType *getExpandedGenericArgsPack();
+
   void Profile(llvm::FoldingSetNodeID &ID) {
     Profile(ID, getDecl(), getParent(), getGenericArgs());
   }
@@ -6478,6 +6483,10 @@ class PackType final : public TypeBase, public llvm::FoldingSetNode,
   /// Creates a pack from the types in \p elements.
   static PackType *get(const ASTContext &C, ArrayRef<Type> elements);
 
+  static PackType *get(const ASTContext &C,
+                       TypeArrayView<GenericTypeParamType> params,
+                       ArrayRef<Type> args);
+
 public:
   /// Retrieves the number of elements in this pack.
   unsigned getNumElements() const { return Bits.PackType.Count; }

lib/AST/ParameterPack.cpp

@@ -17,6 +17,7 @@
 
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/GenericParamList.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Type.h"
 #include "swift/AST/Types.h"
@@ -35,6 +36,24 @@ struct PackTypeParameterCollector: TypeWalker {
     if (t->is<PackExpansionType>())
       return Action::SkipChildren;
 
+    if (auto *boundGenericType = dyn_cast<BoundGenericType>(t.getPointer())) {
+      if (auto parentType = boundGenericType->getParent())
+        parentType.walk(*this);
+
+      Type(boundGenericType->getExpandedGenericArgsPack()).walk(*this);
+      return Action::SkipChildren;
+    }
+
+    if (auto *typeAliasType = dyn_cast<TypeAliasType>(t.getPointer())) {
+      if (typeAliasType->getDecl()->isGeneric()) {
+        if (auto parentType = typeAliasType->getParent())
+          parentType.walk(*this);
+
+        Type(typeAliasType->getExpandedGenericArgsPack()).walk(*this);
+        return Action::SkipChildren;
+      }
+    }
+
     if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
       if (paramTy->isParameterPack())
         typeParams.insert(paramTy);
@@ -348,3 +367,51 @@ unsigned ParameterList::getOrigParamIndex(SubstitutionMap subMap,
   dump(llvm::errs());
   abort();
 }
+
+/// <T...> Foo<T, Pack{Int, String}> => Pack{T..., Int, String}
+PackType *BoundGenericType::getExpandedGenericArgsPack() {
+  // It would be nicer to use genericSig.getInnermostGenericParams() here,
+  // but that triggers a request cycle if we're in the middle of computing
+  // the generic signature already.
+  SmallVector<Type, 2> params;
+  for (auto *paramDecl : getDecl()->getGenericParams()->getParams()) {
+    params.push_back(paramDecl->getDeclaredInterfaceType());
+  }
+
+  return PackType::get(getASTContext(),
+                       TypeArrayView<GenericTypeParamType>(params),
+                       getGenericArgs());
+}
+
+/// <T...> Foo<T, Pack{Int, String}> => Pack{T..., Int, String}
+PackType *TypeAliasType::getExpandedGenericArgsPack() {
+  if (!getDecl()->isGeneric())
+    return nullptr;
+
+  auto genericSig = getGenericSignature();
+  return PackType::get(getASTContext(),
+                       genericSig.getInnermostGenericParams(),
+                       getDirectGenericArgs());
+}
+
+/// <T...> Pack{T, Pack{Int, String}} => Pack{T..., Int, String}
+PackType *PackType::get(const ASTContext &C,
+                        TypeArrayView<GenericTypeParamType> params,
+                        ArrayRef<Type> args) {
+  SmallVector<Type, 2> wrappedArgs;
+
+  assert(params.size() == args.size());
+  for (unsigned i = 0, e = params.size(); i < e; ++i) {
+    auto arg = args[i];
+
+    if (params[i]->isParameterPack()) {
+      wrappedArgs.push_back(PackExpansionType::get(
+          arg, arg->getReducedShape()));
+      continue;
+    }
+
+    wrappedArgs.push_back(arg);
+  }
+
+  return get(C, wrappedArgs)->flattenPackTypes();
+}