Merge pull request #66382 from slavapestov/pack-element-type-5.9

Preliminary support for pack element references to outer pack expansions [5.9]

Author: slavapestov

docs/ABI/Mangling.rst

@@ -675,6 +675,8 @@ Types
   type ::= assoc-type-list 'QZ'                      // shortcut for 'QYz'
   type ::= opaque-type-decl-name bound-generic-args 'Qo' INDEX // opaque type
   
+  type ::= pack-type 'Qe' INDEX              // pack element type
+  
   type ::= pattern-type count-type 'Qp'      // pack expansion type
   type ::= pack-element-list 'QP'            // pack type
   type ::= pack-element-list 'QS' DIRECTNESS // SIL pack type

include/swift/AST/InFlightSubstitution.h

@@ -49,27 +49,15 @@ class InFlightSubstitution {
   InFlightSubstitution(const InFlightSubstitution &) = delete;
   InFlightSubstitution &operator=(const InFlightSubstitution &) = delete;
 
-  // TODO: when we add PackElementType, we should recognize it during
-  // substitution and either call different methods on this class or
-  // pass an extra argument for the pack-expansion depth D.  We should
-  // be able to rely on that to mark a pack-element reference instead
-  // of checking whether the original type was a pack.  Substitution
-  // should use the D'th entry from the end of ActivePackExpansions to
-  // guide the element substitution:
-  //   - project the given index of the pack substitution
-  //   - wrap it in a PackElementType if it's a subst expansion
-  //   - the depth of that PackElementType is the number of subst
-  //     expansions between the depth entry and the end of
-  //     ActivePackExpansions
-
   /// Perform primitive substitution on the given type.  Returns Type()
   /// if the type should not be substituted as a whole.
-  Type substType(SubstitutableType *origType);
+  Type substType(SubstitutableType *origType, unsigned level);
 
   /// Perform primitive conformance lookup on the given type.
   ProtocolConformanceRef lookupConformance(CanType dependentType,
                                            Type conformingReplacementType,
-                                           ProtocolDecl *conformedProtocol);
+                                           ProtocolDecl *conformedProtocol,
+                                           unsigned level);
 
   /// Given the shape type of a pack expansion, invoke the given callback
   /// for each expanded component of it.  If the substituted component

include/swift/AST/Type.h

@@ -151,7 +151,11 @@ enum class SubstFlags {
   /// Map member types to their desugared witness type.
   DesugarMemberTypes = 0x02,
   /// Substitute types involving opaque type archetypes.
-  SubstituteOpaqueArchetypes = 0x04
+  SubstituteOpaqueArchetypes = 0x04,
+  /// Don't increase pack expansion level for free pack references.
+  /// Do not introduce new usages of this flag.
+  /// FIXME: Remove this.
+  PreservePackExpansionLevel = 0x08,
 };
 
 /// Options for performing substitutions into a type.

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

@@ -710,8 +710,10 @@ class alignas(1 << TypeAlignInBits) TypeBase
       SmallVectorImpl<OpenedArchetypeType *> &rootOpenedArchetypes) const;
 
   /// Retrieve the set of type parameter packs that occur within this type.
-  void getTypeParameterPacks(
-      SmallVectorImpl<Type> &rootParameterPacks);
+  void getTypeParameterPacks(SmallVectorImpl<Type> &rootParameterPacks);
+
+  /// Retrieve the set of type parameter packs that occur within this type.
+  void walkPackReferences(llvm::function_ref<bool (Type)> fn);
 
   /// Replace opened archetypes with the given root with their most
   /// specific non-dependent upper bounds throughout this type.
@@ -778,6 +780,9 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// this function will wrap into a pack containing a singleton expansion.
   PackType *getPackSubstitutionAsPackType();
 
+  /// Increase the expansion level of each parameter pack appearing in this type.
+  Type increasePackElementLevel(unsigned level);
+
   /// Determines whether this type is an lvalue. This includes both straight
   /// lvalue types as well as tuples or optionals of lvalues.
   bool hasLValueType() {
@@ -6412,7 +6417,7 @@ class PackArchetypeType final
       LayoutConstraint Layout);
 
   // Returns the reduced shape type for this pack archetype.
-  CanType getReducedShape() const;
+  CanType getReducedShape();
 
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::PackArchetype;
@@ -6994,13 +6999,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.
@@ -7035,16 +7098,6 @@ inline bool TypeBase::isTypeParameter() {
   return t->is<GenericTypeParamType>();
 }
 
-inline bool TypeBase::isParameterPack() {
-  Type t(this);
-
-  while (auto *memberTy = t->getAs<DependentMemberType>())
-    t = memberTy->getBase();
-
-  return t->is<GenericTypeParamType>() &&
-         t->castTo<GenericTypeParamType>()->isParameterPack();
-}
-
 // TODO: This will become redundant once InOutType is removed.
 inline bool TypeBase::isMaterializable() {
   return !(hasLValueType() || is<InOutType>());

include/swift/Demangling/DemangleNodes.def

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

include/swift/SIL/AbstractionPattern.h

@@ -938,10 +938,9 @@ class AbstractionPattern {
     case Kind::Discard: {
       auto type = getType();
       if (isa<DependentMemberType>(type) ||
-          isa<GenericTypeParamType>(type)) {
-        return true;
-      }
-      if (isa<ArchetypeType>(type)) {
+          isa<GenericTypeParamType>(type) ||
+          isa<PackElementType>(type) ||
+          isa<ArchetypeType>(type)) {
         return true;
       }
       return false;
@@ -960,7 +959,8 @@ class AbstractionPattern {
     case Kind::Discard: {
       auto type = getType();
       if (isa<DependentMemberType>(type) ||
-          isa<GenericTypeParamType>(type)) {
+          isa<GenericTypeParamType>(type) ||
+          isa<PackElementType>(type)) {
         return true;
       }
       if (auto archetype = dyn_cast<ArchetypeType>(type)) {
@@ -1448,6 +1448,10 @@ class AbstractionPattern {
   /// the abstraction pattern for an element type.
   AbstractionPattern getTupleElementType(unsigned index) const;
 
+  /// Given that the value being abstracted is a pack element type, return
+  /// the abstraction pattern for its pack type.
+  AbstractionPattern getPackElementPackType() const;
+
   /// Given that the value being abstracted is a pack type, return
   /// the abstraction pattern for an element type.
   AbstractionPattern getPackElementType(unsigned index) const;

include/swift/SIL/SILType.h

@@ -660,12 +660,12 @@ class SILType {
   SILType subst(Lowering::TypeConverter &tc, TypeSubstitutionFn subs,
                 LookupConformanceFn conformances,
                 CanGenericSignature genericSig = CanGenericSignature(),
-                bool shouldSubstituteOpaqueArchetypes = false) const;
+                SubstOptions options = None) const;
 
   SILType subst(SILModule &M, TypeSubstitutionFn subs,
                 LookupConformanceFn conformances,
                 CanGenericSignature genericSig = CanGenericSignature(),
-                bool shouldSubstituteOpaqueArchetypes = false) const;
+                SubstOptions options = None) const;
 
   SILType subst(Lowering::TypeConverter &tc,
                 InFlightSubstitution &IFS,