SIL: Fix SubstFunctionTypePatternVisitor handling of pack types

The included test case also covers the previous commit.

Author: slavapestov

include/swift/SIL/AbstractionPattern.h

@@ -37,6 +37,7 @@ namespace swift {
 namespace Lowering {
 class FunctionParamGenerator;
 class TupleElementGenerator;
+class PackElementGenerator;
 
 /// A pattern for the abstraction of a value.
 ///
@@ -1450,6 +1451,15 @@ class AbstractionPattern {
     llvm_unreachable("bad kind");
   }
 
+  /// Perform a parallel visitation of the elements of a pack 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 pack pattern.
+  void forEachPackElement(CanPackType substPackType,
+         llvm::function_ref<void(PackElementGenerator &element)> fn) const;
+
   /// Given that the value being abstracted is a move only type, return the
   /// abstraction pattern with the move only bit removed.
   AbstractionPattern removingMoveOnlyWrapper() const;

include/swift/SIL/AbstractionPatternGenerators.h

@@ -335,6 +335,122 @@ class TupleElementGenerator {
   }
 };
 
+/// A generator for traversing the formal elements of a pack type
+/// while properly respecting variadic generics.
+class PackElementGenerator {
+  // The steady state of the generator.
+
+  /// The abstraction pattern of the entire pack type.  Set once
+  /// during construction.
+  AbstractionPattern origPackType;
+
+  /// The substituted pack type.  Set once during construction.
+  CanPackType substPackType;
+
+  /// The number of orig elements to traverse.  Set once during
+  /// construction.
+  unsigned numOrigElts;
+
+  /// The index of the current orig element.
+  /// Incremented during advance().
+  unsigned origEltIndex = 0;
+
+  /// The (start) index of the current subst elements.
+  /// Incremented during advance().
+  unsigned substEltIndex = 0;
+
+  /// The number of subst elements corresponding to the current
+  /// orig element.
+  unsigned numSubstEltsForOrigElt;
+
+  /// Whether the current orig element is a pack expansion.
+  bool origEltIsExpansion;
+
+  /// The abstraction pattern of the current orig element.
+  /// If it is a pack expansion, this is the expansion type, not the
+  /// pattern type.
+  AbstractionPattern origEltType = AbstractionPattern::getInvalid();
+
+  /// Load the informaton for the current orig element into the
+  /// fields above for it.
+  void loadElement() {
+    origEltType = origPackType.getPackElementType(origEltIndex);
+    origEltIsExpansion = origEltType.isPackExpansion();
+    numSubstEltsForOrigElt =
+      (origEltIsExpansion
+         ? origEltType.getNumPackExpandedComponents()
+         : 1);
+  }
+
+public:
+  PackElementGenerator(AbstractionPattern origTupleType,
+                       CanPackType substPackType);
+
+  /// Is the traversal finished?  If so, none of the getters below
+  /// are allowed to be called.
+  bool isFinished() const {
+    return origEltIndex == numOrigElts;
+  }
+
+  /// Advance to the next orig element.
+  void advance() {
+    assert(!isFinished());
+    origEltIndex++;
+    substEltIndex += numSubstEltsForOrigElt;
+    if (!isFinished()) loadElement();
+  }
+
+  /// Return the index of the current orig element.
+  unsigned getOrigIndex() const {
+    assert(!isFinished());
+    return origEltIndex;
+  }
+
+  /// Return the index of the (first) subst element corresponding
+  /// to the current orig element.
+  unsigned getSubstIndex() const {
+    assert(!isFinished());
+    return substEltIndex;
+  }
+
+  IntRange<unsigned> getSubstIndexRange() const {
+    assert(!isFinished());
+    return IntRange<unsigned>(substEltIndex,
+                              substEltIndex + numSubstEltsForOrigElt);
+  }
+
+  /// Return the type of the current orig element.
+  const AbstractionPattern &getOrigType() const {
+    assert(!isFinished());
+    return origEltType;
+  }
+
+  /// Return whether the current orig element type is a pack expansion.
+  bool isOrigPackExpansion() const {
+    assert(!isFinished());
+    return origEltIsExpansion;
+  }
+
+  /// Return the substituted elements corresponding to the current
+  /// orig element type.  If the current orig element is not a
+  /// pack expansion, this will have exactly one element.
+  CanTypeArrayRef getSubstTypes() const {
+    assert(!isFinished());
+    return substPackType
+             .getElementTypes().slice(substEltIndex,
+                                      numSubstEltsForOrigElt);
+  }
+
+  /// Call this to finalize the traversal and assert that it was done
+  /// properly.
+  void finish() {
+    assert(isFinished() && "didn't finish the traversal");
+    assert(substEltIndex == substPackType->getNumElements() &&
+           "didn't exhaust subst elements; possible missing subs on "
+           "orig pack type");
+  }
+};
+
 } // end namespace Lowering
 } // end namespace swift
 

lib/SIL/IR/AbstractionPattern.cpp

@@ -741,6 +741,72 @@ bool AbstractionPattern::matchesPack(CanPackType substType) {
   llvm_unreachable("bad kind");
 }
 
+void AbstractionPattern::forEachPackElement(CanPackType substType,
+      llvm::function_ref<void(PackElementGenerator &)> handleElement) const {
+  PackElementGenerator elt(*this, substType);
+  for (; !elt.isFinished(); elt.advance()) {
+    handleElement(elt);
+  }
+  elt.finish();
+}
+
+void AbstractionPattern::forEachExpandedPackElement(CanPackType substType,
+                      llvm::function_ref<void(AbstractionPattern origEltType,
+                                              CanType substEltType)>
+                        handleElement) const {
+  assert(matchesPack(substType));
+
+  auto substEltTypes = substPackType.getElementTypes();
+
+  // Handle opaque patterns by just iterating the substituted components.
+  if (!isPack()) {
+    for (auto i : indices(substEltTypes)) {
+      handleElement(getPackElementType(i), substEltTypes[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(getNumPackElements())) {
+    auto origEltType = getPackElementType(origEltIndex);
+    if (!origEltType.isPackExpansion()) {
+      handleElement(origEltType, substEltTypes[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);
+        substEltIndex++;
+      }
+    }
+  }
+  assert(substEltIndex == substEltTypes.size());
+}
+
+PackElementGenerator::PackElementGenerator(
+                            AbstractionPattern origPackType,
+                            CanPackType substPackType)
+    : origPackType(origPackType), substPackType(substPackType) {
+  assert(origPackType.isPack());
+  assert(origPackType.matchesPack(substPackType));
+  numOrigElts = origPackType.getNumPackElements();
+
+  if (!isFinished()) loadElement();
+}
+
 AbstractionPattern
 AbstractionPattern::getPackExpansionComponentType(CanType substType) const {
   return getPackExpansionComponentType(isa<PackExpansionType>(substType));
@@ -2214,12 +2280,25 @@ class SubstFunctionTypePatternVisitor
   }
 
   CanType visitPackType(CanPackType pack, AbstractionPattern pattern) {
+    assert(pattern.isPack());
+
     // Break down the pack.
     SmallVector<CanType, 4> packElts;
-    for (auto i : range(pack->getNumElements())) {
-      packElts.push_back(visit(pack.getElementType(i),
-                               pattern.getPackElementType(i)));
-    }
+
+    pattern.forEachPackElement(pack, [&](PackElementGenerator &elt) {
+      auto substEltTypes = elt.getSubstTypes();
+      CanType eltTy;
+      if (!elt.isOrigPackExpansion()) {
+        eltTy = visit(substEltTypes[0], elt.getOrigType());
+      } else {
+        CanType candidateSubstType;
+        if (!substEltTypes.empty())
+          candidateSubstType = substEltTypes[0];
+        eltTy = handlePackExpansion(elt.getOrigType(), candidateSubstType);
+      }
+
+      packElts.push_back(eltTy);
+    });
 
     return CanPackType::get(TC.Context, packElts);
   }

test/SILGen/variadic_generic_keypath_descriptors.swift

@@ -0,0 +1,16 @@
+// RUN: %target-swift-emit-silgen -enable-library-evolution -disable-availability-checking %s | %FileCheck %s
+
+// rdar://problem/112474421
+
+public struct G<each T> {
+   public var property1: (repeat each T) -> ()
+   public var property2: ((repeat each T)) -> ()
+   public var property3: (G<repeat each T>) -> ()
+}
+
+// We don't attempt to emit a keypath descriptor for property1 and property2,
+// and we shouldn't crash while emitting a keypath descriptor for property3.
+
+// CHECK-NOT: sil_property #G.property1<each τ_0_0>
+// CHECK-NOT: sil_property #G.property2<each τ_0_0>
+// CHECK-LABEL: sil_property #G.property3<each τ_0_0>