[NFC] Add better APIs for parallel destructuring of orig+subst types

As I've been iterating on this work, I've been gradually mulling these
over, and I think this is the way to go for now.  These should make it
a lot less cumbersome to write these kinds of traversals correctly.

The intent is to the sunset the existing expanded-components stuff
after I do a similar pass for function parameters.

Author: rjmccall

include/swift/SIL/AbstractionPattern.h

@@ -913,6 +913,9 @@ class AbstractionPattern {
   bool hasCachingKey() const {
     // Only the simplest Kind::Type pattern has a caching key; we
     // don't want to try to unique by Clang node.
+    //
+    // Even if we support Clang nodes someday, we *cannot* cache
+    // by the open-coded patterns like Tuple and PackExpansion.
     return getKind() == Kind::Type || getKind() == Kind::Opaque
         || getKind() == Kind::Discard;
   }
@@ -1216,9 +1219,7 @@ class AbstractionPattern {
     case Kind::Invalid:
       llvm_unreachable("querying invalid abstraction pattern!");
     case Kind::Opaque:
-      return typename CanTypeWrapperTraits<TYPE>::type();
     case Kind::Tuple:
-      return typename CanTypeWrapperTraits<TYPE>::type();
     case Kind::OpaqueFunction:
     case Kind::OpaqueDerivativeFunction:
       return typename CanTypeWrapperTraits<TYPE>::type();
@@ -1275,7 +1276,7 @@ class AbstractionPattern {
 
   /// Is the given tuple type a valid substitution of this abstraction
   /// pattern?
-  bool matchesTuple(CanTupleType substType);
+  bool matchesTuple(CanTupleType substType) const;
 
   bool isTuple() const {
     switch (getKind()) {
@@ -1346,6 +1347,40 @@ class AbstractionPattern {
     return { { this, 0 }, { this, getNumTupleElements() } };
   }
 
+  /// Perform a parallel visitation of the elements of a tuple type,
+  /// preserving structure about where pack expansions appear in the
+  /// original type and how many elements of the substituted type they
+  /// expand to.
+  ///
+  /// This pattern must be a tuple pattern.
+  ///
+  /// Calls handleScalar or handleExpansion as appropriate for each
+  /// element of the original tuple, in order.
+  void forEachTupleElement(CanTupleType substType,
+           llvm::function_ref<void(unsigned origEltIndex,
+                                   unsigned substEltIndex,
+                                   AbstractionPattern origEltType,
+                                   CanType substEltType)>
+               handleScalar,
+           llvm::function_ref<void(unsigned origEltIndex,
+                                   unsigned substEltIndex,
+                                   AbstractionPattern origExpansionType,
+                                   CanTupleEltTypeArrayRef substEltTypes)>
+               handleExpansion) const;
+
+  /// Perform a parallel visitation of the elements of a tuple type,
+  /// expanding the elements of the type.  This preserves the structure
+  /// of the *substituted* tuple type: it will be called once per element
+  /// of the substituted type, in order.  The original element trappings
+  /// are also provided for convenience.
+  ///
+  /// This pattern must match the substituted type, but it may be an
+  /// opaque pattern.
+  void forEachExpandedTupleElement(CanTupleType substType,
+      llvm::function_ref<void(AbstractionPattern origEltType,
+                              CanType substEltType,
+                              const TupleTypeElt &elt)> handleElement) const;
+
   /// Is the given pack type a valid substitution of this abstraction
   /// pattern?
   bool matchesPack(CanPackType substType);
@@ -1420,13 +1455,20 @@ class AbstractionPattern {
   /// the abstraction pattern for an element type.
   AbstractionPattern getPackElementType(unsigned index) const;
 
-  /// Give that the value being abstracted is a pack expansion type, return the
-  /// underlying pattern type.
+  /// Given that the value being abstracted is a pack expansion type,
+  /// return the underlying pattern type.
+  ///
+  /// If you're looking for getPackExpansionCountType(), it deliberately
+  /// does not exist.  Count types are not lowered types, and the original
+  /// count types are not relevant to lowering.  Only the substituted
+  /// components and expansion counts are significant.
   AbstractionPattern getPackExpansionPatternType() const;
 
-  /// Give that the value being abstracted is a pack expansion type, return the
-  /// underlying count type.
-  AbstractionPattern getPackExpansionCountType() const;
+  /// Given that the value being abstracted is a pack expansion type,
+  /// return the appropriate pattern type for the given expansion
+  /// component.
+  AbstractionPattern getPackExpansionComponentType(CanType substType) const;
+  AbstractionPattern getPackExpansionComponentType(bool isExpansion) const;
 
   /// Given that the value being abstracted is a function, return the
   /// abstraction pattern for its result type.
@@ -1486,6 +1528,8 @@ class AbstractionPattern {
   void forEachPackExpandedComponent(
       llvm::function_ref<void(AbstractionPattern pattern)> fn) const;
 
+  size_t getNumPackExpandedComponents() const;
+
   SmallVector<AbstractionPattern, 4> getPackExpandedComponents() const;
 
   /// If this pattern refers to a foreign ObjC method that was imported as 

lib/SIL/IR/AbstractionPattern.cpp

@@ -282,7 +282,7 @@ LayoutConstraint AbstractionPattern::getLayoutConstraint() const {
   }
 }
 
-bool AbstractionPattern::matchesTuple(CanTupleType substType) {
+bool AbstractionPattern::matchesTuple(CanTupleType substType) const {
   switch (getKind()) {
   case Kind::Invalid:
     llvm_unreachable("querying invalid abstraction pattern!");
@@ -311,26 +311,25 @@ bool AbstractionPattern::matchesTuple(CanTupleType substType) {
     LLVM_FALLTHROUGH;
   case Kind::Tuple: {
     size_t nextSubstIndex = 0;
-    auto nextComponentIsAcceptable =
-        [&](AbstractionPattern origComponentType) -> bool {
+    auto nextComponentIsAcceptable = [&](bool isPackExpansion) -> bool {
       if (nextSubstIndex == substType->getNumElements())
         return false;
       auto substComponentType = substType.getElementType(nextSubstIndex++);
-      return (origComponentType.isPackExpansion() ==
-              isa<PackExpansionType>(substComponentType));
+      return (isPackExpansion == isa<PackExpansionType>(substComponentType));
     };
-    for (size_t i = 0, n = getNumTupleElements(); i != n; ++i) {
-      auto elt = getTupleElementType(i);
-      if (elt.isPackExpansion()) {
-        bool fail = false;
-        elt.forEachPackExpandedComponent([&](AbstractionPattern component) {
-          if (!nextComponentIsAcceptable(component))
-            fail = true;
-        });
-        if (fail) return false;
-      } else {
-        if (!nextComponentIsAcceptable(elt))
-          return false;
+    for (auto elt : getTupleElementTypes()) {
+      bool isPackExpansion = elt.isPackExpansion();
+      if (isPackExpansion && elt.GenericSubs) {
+        auto origExpansion = cast<PackExpansionType>(elt.getType());
+        auto substShape = cast<PackType>(
+          origExpansion.getCountType().subst(elt.GenericSubs)
+            ->getCanonicalType());
+        for (auto shapeElt : substShape.getElementTypes()) {
+          if (!nextComponentIsAcceptable(isa<PackExpansionType>(shapeElt)))
+            return false;
+        }
+      } else if (!nextComponentIsAcceptable(isPackExpansion)) {
+        return false;
       }
     }
     return nextSubstIndex == substType->getNumElements();
@@ -469,6 +468,87 @@ bool AbstractionPattern::doesTupleContainPackExpansionType() const {
   llvm_unreachable("bad kind");
 }
 
+void AbstractionPattern::forEachTupleElement(CanTupleType substType,
+                      llvm::function_ref<void(unsigned origEltIndex,
+                                              unsigned substEltIndex,
+                                              AbstractionPattern origEltType,
+                                              CanType substEltType)>
+                        handleScalar,
+                      llvm::function_ref<void(unsigned origEltIndex,
+                                              unsigned substEltIndex,
+                                              AbstractionPattern origExpansionType,
+                                              CanTupleEltTypeArrayRef substEltTypes)>
+                        handleExpansion) const {
+  assert(isTuple() && "can only call on a tuple expansion");
+  assert(matchesTuple(substType));
+
+  size_t substEltIndex = 0;
+  auto substEltTypes = substType.getElementTypes();
+  for (size_t origEltIndex : range(getNumTupleElements())) {
+    auto origEltType = getTupleElementType(origEltIndex);
+    if (!origEltType.isPackExpansion()) {
+      handleScalar(origEltIndex, substEltIndex,
+                   origEltType, substEltTypes[substEltIndex]);
+      substEltIndex++;
+    } else {
+      auto numComponents = origEltType.getNumPackExpandedComponents();
+      handleExpansion(origEltIndex, substEltIndex, origEltType,
+                      substEltTypes.slice(substEltIndex, numComponents));
+      substEltIndex += numComponents;
+    }
+  }
+  assert(substEltIndex == substEltTypes.size());
+}
+
+void AbstractionPattern::forEachExpandedTupleElement(CanTupleType substType,
+                      llvm::function_ref<void(AbstractionPattern origEltType,
+                                              CanType substEltType,
+                                              const TupleTypeElt &elt)>
+                        handleElement) const {
+  assert(matchesTuple(substType));
+
+  auto substEltTypes = substType.getElementTypes();
+
+  // Handle opaque patterns by just iterating the substituted components.
+  if (!isTuple()) {
+    for (auto i : indices(substEltTypes)) {
+      handleElement(getTupleElementType(i), substEltTypes[i],
+                    substType->getElement(i));
+    }
+    return;
+  }
+
+  // For non-opaque patterns, we have to iterate the original components
+  // in order to match things up properly, but we'll still end up calling
+  // once per substituted element.
+  size_t substEltIndex = 0;
+  for (size_t origEltIndex : range(getNumTupleElements())) {
+    auto origEltType = getTupleElementType(origEltIndex);
+    if (!origEltType.isPackExpansion()) {
+      handleElement(origEltType, substEltTypes[substEltIndex],
+                    substType->getElement(substEltIndex));
+      substEltIndex++;
+    } else {
+      auto origPatternType = origEltType.getPackExpansionPatternType();
+      for (auto i : range(origEltType.getNumPackExpandedComponents())) {
+        (void) i;
+        auto substEltType = substEltTypes[substEltIndex];
+        // When the substituted type is a pack expansion, pass down
+        // the original element type so that it's *also* a pack expansion.
+        // Clients expect to look through this structure in parallel on
+        // both types.  The count is misleading, but normal usage won't
+        // access it, and there's nothing we could provide that *wouldn't*
+        // be misleading in one way or another.
+        handleElement(isa<PackExpansionType>(substEltType)
+                        ? origEltType : origPatternType,
+                      substEltType, substType->getElement(substEltIndex));
+        substEltIndex++;
+      }
+    }
+  }
+  assert(substEltIndex == substEltTypes.size());
+}
+
 static CanType getCanPackElementType(CanType type, unsigned index) {
   return cast<PackType>(type).getElementType(index);
 }
@@ -541,6 +621,17 @@ bool AbstractionPattern::matchesPack(CanPackType substType) {
   llvm_unreachable("bad kind");
 }
 
+AbstractionPattern
+AbstractionPattern::getPackExpansionComponentType(CanType substType) const {
+  return getPackExpansionComponentType(isa<PackExpansionType>(substType));
+}
+
+AbstractionPattern
+AbstractionPattern::getPackExpansionComponentType(bool isExpansion) const {
+  assert(isPackExpansion());
+  return isExpansion ? *this : getPackExpansionPatternType();
+}
+
 static CanType getPackExpansionPatternType(CanType type) {
   return cast<PackExpansionType>(type).getPatternType();
 }
@@ -584,49 +675,6 @@ AbstractionPattern AbstractionPattern::getPackExpansionPatternType() const {
   llvm_unreachable("bad kind");
 }
 
-static CanType getPackExpansionCountType(CanType type) {
-  return cast<PackExpansionType>(type).getCountType();
-}
-
-AbstractionPattern AbstractionPattern::getPackExpansionCountType() const {
-  switch (getKind()) {
-  case Kind::Invalid:
-    llvm_unreachable("querying invalid abstraction pattern!");
-  case Kind::ObjCMethodType:
-  case Kind::CurriedObjCMethodType:
-  case Kind::PartialCurriedObjCMethodType:
-  case Kind::CFunctionAsMethodType:
-  case Kind::CurriedCFunctionAsMethodType:
-  case Kind::PartialCurriedCFunctionAsMethodType:
-  case Kind::CXXMethodType:
-  case Kind::CurriedCXXMethodType:
-  case Kind::PartialCurriedCXXMethodType:
-  case Kind::Tuple:
-  case Kind::OpaqueFunction:
-  case Kind::OpaqueDerivativeFunction:
-  case Kind::ObjCCompletionHandlerArgumentsType:
-  case Kind::ClangType:
-    llvm_unreachable("pattern for function or tuple cannot be for "
-                     "pack expansion type");
-
-  case Kind::Opaque:
-    return *this;
-
-  case Kind::Type:
-    if (isTypeParameterOrOpaqueArchetype())
-      return AbstractionPattern::getOpaque();
-    return AbstractionPattern(getGenericSubstitutions(),
-                              getGenericSignature(),
-                              ::getPackExpansionCountType(getType()));
-
-  case Kind::Discard:
-    return AbstractionPattern::getDiscard(
-        getGenericSubstitutions(), getGenericSignature(),
-        ::getPackExpansionCountType(getType()));
-  }
-  llvm_unreachable("bad kind");
-}
-
 SmallVector<AbstractionPattern, 4>
 AbstractionPattern::getPackExpandedComponents() const {
   SmallVector<AbstractionPattern, 4> result;
@@ -636,6 +684,21 @@ AbstractionPattern::getPackExpandedComponents() const {
   return result;
 }
 
+size_t AbstractionPattern::getNumPackExpandedComponents() const {
+  assert(isPackExpansion());
+  assert(getKind() == Kind::Type || getKind() == Kind::Discard);
+
+  // If we don't have substitutions, we should be walking parallel
+  // structure; take a single element.
+  if (!GenericSubs) return 1;
+
+  // Otherwise, substitute the expansion shape.
+  auto origExpansion = cast<PackExpansionType>(getType());
+  auto substShape = cast<PackType>(
+    origExpansion.getCountType().subst(GenericSubs)->getCanonicalType());
+  return substShape->getNumElements();
+}
+
 void AbstractionPattern::forEachPackExpandedComponent(
           llvm::function_ref<void (AbstractionPattern)> fn) const {
   assert(isPackExpansion());
@@ -665,7 +728,7 @@ void AbstractionPattern::forEachPackExpandedComponent(
   }
 
   default:
-    return fn(*this);
+    llvm_unreachable("not a pack expansion");
   }
   llvm_unreachable("bad kind");
 }
@@ -692,8 +755,9 @@ AbstractionPattern AbstractionPattern::removingMoveOnlyWrapper() const {
     llvm_unreachable("not handled yet");
   case Kind::Discard:
     llvm_unreachable("operation not needed on discarded abstractions yet");
-  case Kind::Opaque:
   case Kind::Tuple:
+    llvm_unreachable("cannot apply move-only wrappers to open-coded patterns");
+  case Kind::Opaque:
   case Kind::Type:
     if (auto mvi = dyn_cast<SILMoveOnlyWrappedType>(getType())) {
       return AbstractionPattern(getGenericSubstitutions(),
@@ -727,8 +791,9 @@ AbstractionPattern AbstractionPattern::addingMoveOnlyWrapper() const {
     llvm_unreachable("not handled yet");
   case Kind::Discard:
     llvm_unreachable("operation not needed on discarded abstractions yet");
-  case Kind::Opaque:
   case Kind::Tuple:
+    llvm_unreachable("cannot add move only wrapper to open-coded pattern");
+  case Kind::Opaque:
   case Kind::Type:
     if (isa<SILMoveOnlyWrappedType>(getType()))
       return *this;
@@ -1686,7 +1751,7 @@ AbstractionPattern::operator==(const AbstractionPattern &other) const {
       }
     }
     return true;
-  
+
   case Kind::Type:
   case Kind::Discard:
     return OrigType == other.OrigType
@@ -1996,7 +2061,7 @@ class SubstFunctionTypePatternVisitor
     auto substPatternType = visit(pack->getPatternType(),
                                   pattern.getPackExpansionPatternType());
     auto substCountType = visit(pack->getCountType(),
-                                pattern.getPackExpansionCountType());
+                                AbstractionPattern::getOpaque());
 
     SmallVector<Type> rootParameterPacks;
     substPatternType->getTypeParameterPacks(rootParameterPacks);