AST: Introduce PackElementType

Author: slavapestov

include/swift/AST/TypeDifferenceVisitor.h

@@ -149,6 +149,11 @@ class CanTypeDifferenceVisitor : public CanTypePairVisitor<Impl, bool> {
     return asImpl().visit(type1.getPatternType(), type2.getPatternType());
   }
 
+  bool visitPackElementType(CanPackElementType type1,
+                            CanPackElementType type2) {
+    return asImpl().visit(type1.getPackType(), type2.getPackType());
+  }
+
   bool visitTupleType(CanTupleType type1, CanTupleType type2) {
     return visitComponentArray(type1, type2,
                                type1->getElements(), type2->getElements());

include/swift/AST/TypeMatcher.h

@@ -202,6 +202,18 @@ class TypeMatcher {
       return mismatch(firstPE.getPointer(), secondType, sugaredFirstType);
     }
 
+    bool visitPackElementType(CanPackElementType firstElement, Type secondType,
+                              Type sugaredFirstType) {
+      if (auto secondElement = secondType->getAs<PackElementType>()) {
+        return this->visit(firstElement.getPackType(),
+                           secondElement->getPackType(),
+                           sugaredFirstType->castTo<PackElementType>()
+                             ->getPackType());
+      }
+
+      return mismatch(firstElement.getPointer(), secondType, sugaredFirstType);
+    }
+
     bool visitReferenceStorageType(CanReferenceStorageType firstStorage,
                                    Type secondType, Type sugaredFirstType) {
       if (firstStorage->getKind() == secondType->getDesugaredType()->getKind()) {

include/swift/AST/TypeNodes.def

@@ -186,6 +186,7 @@ TYPE(LValue, Type)
 TYPE(InOut, Type)
 TYPE(Pack, Type)
 TYPE(PackExpansion, Type)
+TYPE(PackElement, Type)
 UNCHECKED_TYPE(TypeVariable, Type)
 ABSTRACT_SUGARED_TYPE(Sugar, Type)
   SUGARED_TYPE(Paren, SugarType)

include/swift/AST/Types.h

@@ -6994,13 +6994,71 @@ BEGIN_CAN_TYPE_WRAPPER(PackExpansionType, Type)
   }
 END_CAN_TYPE_WRAPPER(PackExpansionType, Type)
 
-
 inline CanTypeWrapper<PackExpansionType>
 CanPackType::unwrapSingletonPackExpansion() const {
   return CanPackExpansionType(
       getPointer()->unwrapSingletonPackExpansion());
 }
 
+/// Represents a reference to a pack from an outer expansion. This comes up
+/// after substitution. For example, given these declarations:
+///
+/// typealias A<each T, U> = (repeat (each T, U))
+/// typealias B<each X, repeat each Y> = (repeat A<repeat each X, each Y>)
+///
+/// Naively substituting replacing {T := repeat each X, U := each Y} in the
+/// underlying type of A would give us:
+///
+///   '(repeat (repeat (each X, each Y)))'
+///
+/// However, this is wrong; we're not expanding X and Y in parallel (they
+/// might not even have the same shape). Instead, we're expanding X, and
+/// then on each iteration, expanding Y.
+///
+/// If we annotate each 'repeat' and its corresponding 'each', we instead see
+/// that the above should give us:
+///
+///   '(repeat[1] (repeat[0] (each[0], each[1] U)))'
+///
+/// We number PackExpansionTypes from the innermost one outwards, assigning
+/// a level of 0 to the innermost one. Then, a PackElementType represents a
+/// reference to a parameter pack from an expansion with level > 0.
+class PackElementType : public TypeBase, public llvm::FoldingSetNode {
+  friend class ASTContext;
+
+  Type packType;
+  unsigned level;
+
+  PackElementType(Type packType, unsigned level,
+                  RecursiveTypeProperties properties,
+                  const ASTContext *ctx);
+
+public:
+  static PackElementType *get(Type packType, unsigned level);
+
+  Type getPackType() const { return packType; }
+
+  unsigned getLevel() const { return level; }
+
+  void Profile(llvm::FoldingSetNodeID &ID) {
+    Profile(ID, getPackType(), getLevel());
+  }
+
+  static void Profile(llvm::FoldingSetNodeID &ID, Type packType, unsigned level);
+
+  // Implement isa/cast/dyncast/etc.
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::PackElement;
+  }
+};
+BEGIN_CAN_TYPE_WRAPPER(PackElementType, Type)
+  static CanPackElementType get(CanType pack);
+
+  CanType getPackType() const {
+    return CanType(getPointer()->getPackType());
+  }
+END_CAN_TYPE_WRAPPER(PackElementType, Type)
+
 /// 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.

include/swift/Demangling/DemangleNodes.def

@@ -241,6 +241,7 @@ NODE(Pack)
 NODE(SILPackDirect)
 NODE(SILPackIndirect)
 NODE(PackExpansion)
+NODE(PackElement)
 NODE(Type)
 CONTEXT_NODE(TypeSymbolicReference)
 CONTEXT_NODE(TypeAlias)

lib/AST/ASTContext.cpp

@@ -417,6 +417,7 @@ struct ASTContext::Implementation {
     llvm::FoldingSet<TupleType> TupleTypes;
     llvm::FoldingSet<PackType> PackTypes;
     llvm::FoldingSet<PackExpansionType> PackExpansionTypes;
+    llvm::FoldingSet<PackElementType> PackElementTypes;
     llvm::DenseMap<llvm::PointerIntPair<TypeBase*, 3, unsigned>,
                    MetatypeType*> MetatypeTypes;
     llvm::DenseMap<llvm::PointerIntPair<TypeBase*, 3, unsigned>,
@@ -3290,6 +3291,47 @@ PackType *PackType::getEmpty(const ASTContext &C) {
   return cast<PackType>(CanType(C.TheEmptyPackType));
 }
 
+PackElementType::PackElementType(Type packType, unsigned level,
+                                 RecursiveTypeProperties properties,
+                                 const ASTContext *canCtx)
+  : TypeBase(TypeKind::PackElement, canCtx, properties),
+    packType(packType), level(level) {
+  assert(packType->isParameterPack() ||
+         packType->is<PackArchetypeType>() ||
+         packType->is<TypeVariableType>());
+  assert(level > 0);
+}
+
+PackElementType *PackElementType::get(Type packType, unsigned level) {
+  auto properties = packType->getRecursiveProperties();
+  auto arena = getArena(properties);
+
+  auto &context = packType->getASTContext();
+  llvm::FoldingSetNodeID id;
+  PackElementType::Profile(id, packType, level);
+
+  void *insertPos;
+  if (PackElementType *elementType =
+        context.getImpl().getArena(arena)
+          .PackElementTypes.FindNodeOrInsertPos(id, insertPos))
+    return elementType;
+
+  const ASTContext *canCtx = packType->isCanonical()
+      ? &context : nullptr;
+  PackElementType *elementType =
+      new (context, arena) PackElementType(packType, level, properties,
+                                           canCtx);
+  context.getImpl().getArena(arena).PackElementTypes.InsertNode(elementType,
+                                                                insertPos);
+  return elementType;
+}
+
+void PackElementType::Profile(llvm::FoldingSetNodeID &ID,
+                              Type packType, unsigned level) {
+  ID.AddPointer(packType.getPointer());
+  ID.AddInteger(level);
+}
+
 CanPackType CanPackType::get(const ASTContext &C, ArrayRef<CanType> elements) {
   SmallVector<Type, 8> ncElements(elements.begin(), elements.end());
   return CanPackType(PackType::get(C, ncElements));

lib/AST/ASTDumper.cpp

@@ -3873,6 +3873,13 @@ namespace {
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
 
+    void visitPackElementType(PackElementType *T, StringRef label) {
+      printCommon(label, "element_type");
+      printField("level", T->getLevel());
+      printRec("pack", T->getPackType());
+      PrintWithColorRAII(OS, ParenthesisColor) << ')';
+    }
+
     void visitParenType(ParenType *T, StringRef label) {
       printCommon(label, "paren_type");
       printRec(T->getUnderlyingType());

lib/AST/ASTMangler.cpp

@@ -1285,6 +1285,15 @@ void ASTMangler::appendType(Type type, GenericSignature sig,
       return;
     }
 
+    case TypeKind::PackElement: {
+      auto elementType = cast<PackElementType>(tybase);
+      appendType(elementType->getPackType(), sig, forDecl);
+
+      // FIXME: append expansion depth
+
+      return;
+    }
+
     case TypeKind::Pack: {
       auto packTy = cast<PackType>(tybase);
 

lib/AST/ASTPrinter.cpp

@@ -6049,6 +6049,11 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     visit(T->getPatternType());
   }
 
+  void visitPackElementType(PackElementType *T) {
+    Printer << "@level(" << T->getLevel() << ") ";
+    visit(T->getPackType());
+  }
+
   void visitTupleType(TupleType *T) {
     Printer.callPrintStructurePre(PrintStructureKind::TupleType);
     SWIFT_DEFER { Printer.printStructurePost(PrintStructureKind::TupleType); };

lib/AST/ExistentialGeneralization.cpp

@@ -158,6 +158,7 @@ class Generalizer : public CanTypeVisitor<Generalizer, Type> {
   NO_PRESERVABLE_STRUCTURE(Module)
   NO_PRESERVABLE_STRUCTURE(Pack)
   NO_PRESERVABLE_STRUCTURE(PackExpansion)
+  NO_PRESERVABLE_STRUCTURE(PackElement)
 #undef NO_PRESERVABLE_STRUCTURE
 
   // These types simply shouldn't appear in types that we generalize at all.