Perform component-wise substitution of pack expansions immediately.

Substitution of a pack expansion type may now produce a pack type.
We immediately expand that pack when transforming a tuple, a function
parameter, or a pack.

I had to duplicate the component-wise transformation logic in the
simplifyType transform, which I'm not pleased about, but a little
code duplication seemed a lot better than trying to unify the code
in two very different places.

I think we're very close to being able to assert that pack expansion
shapes are either pack archetypes or pack parameters; unfortunately,
the pack matchers intentionally produce expansions of packs, and I
didn't want to add that to an already-large patch.

Author: rjmccall

include/swift/AST/InFlightSubstitution.h

@@ -32,6 +32,12 @@ class InFlightSubstitution {
   TypeSubstitutionFn BaselineSubstType;
   LookupConformanceFn BaselineLookupConformance;
 
+  struct ActivePackExpansion {
+    bool isSubstExpansion = false;
+    unsigned expansionIndex = 0;
+  };
+  SmallVector<ActivePackExpansion, 4> ActivePackExpansions;
+
 public:
   InFlightSubstitution(TypeSubstitutionFn substType,
                        LookupConformanceFn lookupConformance,
@@ -43,16 +49,72 @@ class InFlightSubstitution {
   InFlightSubstitution(const InFlightSubstitution &) = delete;
   InFlightSubstitution &operator=(const InFlightSubstitution &) = delete;
 
-  Type substType(SubstitutableType *ty) {
-    return BaselineSubstType(ty);
-  }
-
+  // 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);
+
+  /// Perform primitive conformance lookup on the given type.
   ProtocolConformanceRef lookupConformance(CanType dependentType,
                                            Type conformingReplacementType,
-                                           ProtocolDecl *conformedProtocol) {
-    return BaselineLookupConformance(dependentType,
-                                     conformingReplacementType,
-                                     conformedProtocol);
+                                           ProtocolDecl *conformedProtocol);
+
+  /// Given the shape type of a pack expansion, invoke the given callback
+  /// for each expanded component of it.  If the substituted component
+  /// is an expansion component, the desired shape of that expansion
+  /// is passed as the argument; otherwise, the argument is Type().
+  /// In either case, an active expansion is entered on this IFS for
+  /// the duration of the call to handleComponent, and subsequent
+  /// pack-element type references will substitute to the corresponding
+  /// element of the substitution of the pack.
+  void expandPackExpansionShape(Type origShape,
+      llvm::function_ref<void(Type substComponentShape)> handleComponent);
+
+  /// Call the given function for each expanded component type of the
+  /// given pack expansion type.  The function will be invoked with the
+  /// active expansion still active.
+  void expandPackExpansionType(PackExpansionType *origExpansionType,
+      llvm::function_ref<void(Type substType)> handleComponentType) {
+    expandPackExpansionShape(origExpansionType->getCountType(),
+                             [&](Type substComponentShape) {
+      auto origPatternType = origExpansionType->getPatternType();
+      auto substEltType = origPatternType.subst(*this);
+
+      auto substComponentType =
+        (substComponentShape
+           ? PackExpansionType::get(substEltType, substComponentShape)
+           : substEltType);
+      handleComponentType(substComponentType);
+    });
+  }
+
+  /// Return a list of component types that the pack expansion expands to.
+  SmallVector<Type, 8>
+  expandPackExpansionType(PackExpansionType *origExpansionType) {
+    SmallVector<Type, 8> substComponentTypes;
+    expandPackExpansionType(origExpansionType, substComponentTypes);
+    return substComponentTypes;
+  }
+
+  /// Expand the list of component types that the pack expansion expands
+  /// to into the given array.
+  void expandPackExpansionType(PackExpansionType *origExpansionType,
+                               SmallVectorImpl<Type> &substComponentTypes) {
+    expandPackExpansionType(origExpansionType, [&](Type substComponentType) {
+      substComponentTypes.push_back(substComponentType);
+    });
   }
 
   class OptionsAdjustmentScope {

lib/AST/ASTContext.cpp

@@ -3257,6 +3257,12 @@ CanPackExpansionType::get(CanType patternType, CanType countType) {
 }
 
 PackExpansionType *PackExpansionType::get(Type patternType, Type countType) {
+  assert(!patternType->is<PackExpansionType>());
+  assert(!countType->is<PackExpansionType>());
+  // FIXME: stop doing this deliberately in PackExpansionMatcher
+  //assert(!patternType->is<PackType>());
+  //assert(!countType->is<PackType>());
+
   auto properties = patternType->getRecursiveProperties();
   properties |= countType->getRecursiveProperties();
 

lib/AST/PackConformance.cpp

@@ -167,207 +167,65 @@ ProtocolConformanceRef PackConformance::subst(SubstitutionMap subMap,
   return subst(IFS);
 }
 
-// TODO: Move this elsewhere since it's generally useful
-static bool arePackShapesEqual(PackType *lhs, PackType *rhs) {
-  if (lhs->getNumElements() != rhs->getNumElements())
-    return false;
-
-  for (unsigned i = 0, e = lhs->getNumElements(); i < e; ++i) {
-    auto lhsElt = lhs->getElementType(i);
-    auto rhsElt = rhs->getElementType(i);
-
-    if (lhsElt->is<PackExpansionType>() != rhsElt->is<PackExpansionType>())
-      return false;
-  }
-
-  return true;
-}
-
-static bool isRootParameterPack(Type t) {
-  if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
-    return paramTy->isParameterPack();
-  } else if (auto *archetypeTy = t->getAs<PackArchetypeType>()) {
-    return archetypeTy->isRoot();
-  }
-
-  return false;
-}
-
-static bool isRootedInParameterPack(Type t) {
-  if (auto *archetypeTy = t->getAs<PackArchetypeType>()) {
-    return true;
-  }
-
-  return t->getRootGenericParam()->isParameterPack();
-}
-
 namespace {
 
-template<typename ImplClass>
-class PackExpander {
-protected:
+struct PackConformanceExpander {
   InFlightSubstitution &IFS;
+  ArrayRef<ProtocolConformanceRef> origConformances;
 
-  PackExpander(InFlightSubstitution &IFS) : IFS(IFS) {}
-
-  ImplClass *asImpl() {
-    return static_cast<ImplClass *>(this);
-  }
-
-  /// We're replacing a pack expansion type with a pack -- flatten the pack
-  /// using the pack expansion's pattern.
-  void addExpandedExpansion(Type origPatternType, PackType *expandedCountType,
-                            unsigned i) {
-
-    // Get all pack parameters referenced from the pattern.
-    SmallVector<Type, 2> rootParameterPacks;
-    origPatternType->getTypeParameterPacks(rootParameterPacks);
-
-    // Each pack parameter referenced from the pattern must be replaced
-    // with a pack type, and all pack types must have the same shape as
-    // the expanded count pack type.
-    llvm::SmallDenseMap<Type, PackType *, 2> expandedPacks;
-    for (auto origParamType : rootParameterPacks) {
-      auto substParamType = origParamType.subst(IFS);
-
-      if (auto expandedParamType = substParamType->template getAs<PackType>()) {
-        assert(arePackShapesEqual(expandedParamType, expandedCountType) &&
-               "TODO: Return an invalid conformance if this fails");
-
-        auto inserted = expandedPacks.insert(
-            std::make_pair(origParamType->getCanonicalType(),
-                           expandedParamType)).second;
-        assert(inserted &&
-               "getTypeParameterPacks() should not return duplicates");
-      } else {
-        assert(false &&
-               "TODO: Return an invalid conformance if this fails");
-      }
-    }
-
-    // For each element of the expanded count, compute the substituted
-    // pattern type.
-    for (unsigned j = 0, ee = expandedCountType->getNumElements(); j < ee; ++j) {
-      auto projectedSubs = [&](SubstitutableType *type) -> Type {
-        // Nested sequence archetypes get passed in here, but we must
-        // handle them via the standard nested type path.
-        if (auto *archetypeType = dyn_cast<ArchetypeType>(type)) {
-          if (!archetypeType->isRoot())
-            return Type();
-        }
-
-        // Compute the substituted type using our parent substitutions.
-        auto substType = Type(type).subst(IFS);
-
-        // If the substituted type is a pack, project the jth element.
-        if (isRootParameterPack(type)) {
-          // FIXME: What if you have something like G<T...>... where G<> is
-          // variadic?
-          assert(substType->template is<PackType>() &&
-                 "TODO: Return an invalid conformance if this fails");
-          auto *packType = substType->template castTo<PackType>();
-          assert(arePackShapesEqual(packType, expandedCountType) &&
-                 "TODO: Return an invalid conformance if this fails");
-
-          return packType->getElementType(j);
-        }
-
-        return IFS.substType(type);
-      };
-
-      auto projectedConformances = [&](CanType origType, Type substType,
-                                       ProtocolDecl *proto) -> ProtocolConformanceRef {
-        auto substConformance =
-          IFS.lookupConformance(origType, substType, proto);
-
-        // If the substituted conformance is a pack, project the jth element.
-        if (isRootedInParameterPack(origType)) {
-          return substConformance.getPack()->getPatternConformances()[j];
-        }
-
-        return substConformance;
-      };
-
-      auto origCountElement = expandedCountType->getElementType(j);
-      auto substCountElement = origCountElement.subst(
-          projectedSubs, projectedConformances, IFS.getOptions());
-
-      asImpl()->add(origCountElement, substCountElement, i);
-    }
-  }
-
-  /// A pack expansion remains unexpanded, so we substitute the pattern and
-  /// form a new pack expansion.
-  void addUnexpandedExpansion(Type origPatternType, Type substCountType,
-                              unsigned i) {
-    auto substPatternType = origPatternType.subst(IFS);
-    auto substExpansion = PackExpansionType::get(substPatternType, substCountType);
+public:
+  // Results built up by the expansion.
+  SmallVector<Type, 4> substElementTypes;
+  SmallVector<ProtocolConformanceRef, 4> substConformances;
 
-    asImpl()->add(origPatternType, substExpansion, i);
-  }
+  PackConformanceExpander(InFlightSubstitution &IFS,
+                          ArrayRef<ProtocolConformanceRef> origConformances)
+    : IFS(IFS), origConformances(origConformances) {}
 
-  /// Scalar elements of the original pack are substituted and added to the
-  /// flattened pack.
-  void addScalar(Type origElement, unsigned i) {
-    auto substElement = origElement.subst(IFS);
+private:
+  /// Substitute a scalar element of the original pack.
+  void substScalar(Type origElementType,
+                   ProtocolConformanceRef origConformance) {
+    auto substElementType = origElementType.subst(IFS);
+    auto substConformance = origConformance.subst(origElementType, IFS);
 
-    asImpl()->add(origElement, substElement, i);
+    substElementTypes.push_back(substElementType);
+    substConformances.push_back(substConformance);
   }
 
-  /// Potentially expand an element of the original pack.
-  void maybeExpandExpansion(PackExpansionType *origExpansion, unsigned i) {
-    auto origPatternType = origExpansion->getPatternType();
-    auto origCountType = origExpansion->getCountType();
-
-    auto substCountType = origCountType.subst(IFS);
-
-    // If the substituted count type is a pack, we're expanding the
-    // original element.
-    if (auto *expandedCountType = substCountType->template getAs<PackType>()) {
-      addExpandedExpansion(origPatternType, expandedCountType, i);
-      return;
-    }
-
-    addUnexpandedExpansion(origPatternType, substCountType, i);
+  /// Substitute and expand an expansion element of the original pack.
+  void substExpansion(PackExpansionType *origExpansionType,
+                       ProtocolConformanceRef origConformance) {
+    IFS.expandPackExpansionType(origExpansionType,
+                                [&](Type substComponentType) {
+      auto origPatternType = origExpansionType->getPatternType();
+
+      // Just substitute the conformance.  We don't directly represent
+      // pack expansion conformances here; it's sort of implicit in the
+      // corresponding pack element type.
+      auto substConformance = origConformance.subst(origPatternType, IFS);
+
+      substElementTypes.push_back(substComponentType);
+      substConformances.push_back(substConformance);
+    });
   }
 
 public:
   void expand(PackType *origPackType) {
-    for (unsigned i = 0, e = origPackType->getNumElements(); i < e; ++i) {
-      auto origElement = origPackType->getElementType(i);
+    assert(origPackType->getNumElements() == origConformances.size());
 
-      // Check if the original element is potentially being expanded.
-      if (auto *origExpansion = origElement->getAs<PackExpansionType>()) {
-        maybeExpandExpansion(origExpansion, i);
-        continue;
+    for (auto i : range(origPackType->getNumElements())) {
+      auto origElementType = origPackType->getElementType(i);
+      if (auto *origExpansion = origElementType->getAs<PackExpansionType>()) {
+        substExpansion(origExpansion, origConformances[i]);
+      } else {
+        substScalar(origElementType, origConformances[i]);
       }
-
-      addScalar(origElement, i);
     }
   }
 };
 
-class PackConformanceExpander : public PackExpander<PackConformanceExpander> {
-public:
-  SmallVector<Type, 4> substElements;
-  SmallVector<ProtocolConformanceRef, 4> substConformances;
-
-  ArrayRef<ProtocolConformanceRef> origConformances;
-
-  PackConformanceExpander(InFlightSubstitution &IFS,
-                          ArrayRef<ProtocolConformanceRef> origConformances)
-    : PackExpander(IFS), origConformances(origConformances) {}
-
-  void add(Type origType, Type substType, unsigned i) {
-    substElements.push_back(substType);
-
-    // FIXME: Pass down projection callbacks
-    substConformances.push_back(origConformances[i].subst(
-        origType, IFS));
-  }
-};
-
-}
+} // end anonymous namespace
 
 ProtocolConformanceRef PackConformance::subst(TypeSubstitutionFn subs,
                                               LookupConformanceFn conformances,
@@ -382,7 +240,7 @@ PackConformance::subst(InFlightSubstitution &IFS) const {
   expander.expand(ConformingType);
 
   auto &ctx = Protocol->getASTContext();
-  auto *substConformingType = PackType::get(ctx, expander.substElements);
+  auto *substConformingType = PackType::get(ctx, expander.substElementTypes);
 
   auto substConformance = PackConformance::get(substConformingType, Protocol,
                                                expander.substConformances);

lib/AST/ParameterPack.cpp

@@ -449,15 +449,15 @@ PackType *PackType::get(const ASTContext &C,
     auto arg = args[i];
 
     if (params[i]->isParameterPack()) {
-      wrappedArgs.push_back(PackExpansionType::get(
-          arg, arg->getReducedShape()));
+      auto argPackElements = arg->castTo<PackType>()->getElementTypes();
+      wrappedArgs.append(argPackElements.begin(), argPackElements.end());
       continue;
     }
 
     wrappedArgs.push_back(arg);
   }
 
-  return get(C, wrappedArgs)->flattenPackTypes();
+  return get(C, wrappedArgs);
 }
 
 PackType *PackType::getSingletonPackExpansion(Type param) {