[CSGen] Use PackExpansionExpr::getExpandedPacks to generate ShapeOf constraints

rather than relying on PackElementExprs collected by preCheck.

This handles pack element expressions and pack element type reprs, and enforces that
all packs expanded by a given pack expansion expression all have the same shape.

Once the parser creates PackExpansionExpr directly (based on a dedicated syntax instead
of postfix '...'), the code in preCheck for identifying and creating pack expansions
can simply be deleted.

Author: hborla

include/swift/AST/Types.h

@@ -72,6 +72,7 @@ class InOutType;
 class OpaqueTypeDecl;
 class OpenedArchetypeType;
 class PackType;
+class PackReferenceTypeRepr;
 class PlaceholderTypeRepr;
 enum class ReferenceCounting : uint8_t;
 enum class ResilienceExpansion : unsigned;

include/swift/Sema/ConstraintSystem.h

@@ -6130,7 +6130,7 @@ class OpenPackElementType {
     this->locator = cs.getConstraintLocator(locator);
   }
 
-  Type operator()(Type packType) const {
+  Type operator()(Type packType, PackReferenceTypeRepr *packRepr) const {
     // Only assert we have an element environment when invoking the function
     // object. In cases where pack elements are referenced outside of a
     // pack expansion, type resolution will error before opening the pack
@@ -6141,6 +6141,14 @@ class OpenPackElementType {
     auto elementLoc = cs.getConstraintLocator(locator,
         LocatorPathElt::OpenedPackElement(elementEnv));
 
+    // If we're opening a pack element from an explicit type repr,
+    // set the type repr types in the constraint system for generating
+    // ShapeOf constraints when visiting the PackExpansionExpr.
+    if (packRepr) {
+      cs.setType(packRepr->getPackType(), packType);
+      cs.setType(packRepr, elementType);
+    }
+
     cs.addConstraint(ConstraintKind::PackElementOf, elementType,
                      packType, elementLoc);
     return elementType;

lib/Sema/CSGen.cpp

@@ -3009,18 +3009,24 @@ namespace {
       auto *shapeTypeVar = CS.createTypeVariable(shapeLoc,
                                                  TVO_CanBindToPack |
                                                  TVO_CanBindToHole);
-      Type packType;
-      if (!expr->getPackElements().empty()) {
-        auto packReference = expr->getPackElements().front()->getPackRefExpr();
-        packType = CS.simplifyType(CS.getType(packReference));
-      } else {
-        // FIXME: The generic environment needs to be per-shape-class.
-        llvm::SmallVector<GenericEnvironment::PackElementBinding, 2> bindings;
-        elementEnv->getPackElementBindings(bindings);
-        packType = bindings.front().second;
+
+      // Generate ShapeOf constraints between all packs expanded by this
+      // pack expansion expression through the shape type variable.
+      SmallVector<ASTNode, 2> expandedPacks;
+      expr->getExpandedPacks(expandedPacks);
+      for (auto pack : expandedPacks) {
+        Type packType;
+        if (auto *elementExpr = getAsExpr<PackElementExpr>(pack)) {
+          packType = CS.getType(elementExpr->getPackRefExpr());
+        } else if (auto *elementType = getAsTypeRepr<PackReferenceTypeRepr>(pack)) {
+          packType = CS.getType(elementType->getPackType());
+        } else {
+          llvm_unreachable("unsupported pack reference ASTNode");
+        }
+
+        CS.addConstraint(ConstraintKind::ShapeOf, packType, shapeTypeVar,
+                         CS.getConstraintLocator(expr));
       }
-      CS.addConstraint(ConstraintKind::ShapeOf, packType, shapeTypeVar,
-                       CS.getConstraintLocator(expr));
 
       return PackExpansionType::get(patternTy, shapeTypeVar);
     }
@@ -3032,7 +3038,7 @@ namespace {
       // of the pack reference.
       OpenPackElementType openPackElement(CS, CS.getConstraintLocator(expr),
                                           PackElementEnvironments.back());
-      return openPackElement(packType);
+      return openPackElement(packType, /*packRepr*/ nullptr);
     }
 
     Type visitDynamicTypeExpr(DynamicTypeExpr *expr) {

lib/Sema/ConstraintSystem.cpp

@@ -781,7 +781,7 @@ Type ConstraintSystem::openUnboundGenericType(GenericTypeDecl *decl,
           DC, None,
           [](auto) -> Type { llvm_unreachable("should not be used"); },
           [](auto &, auto) -> Type { llvm_unreachable("should not be used"); },
-          [](auto) -> Type { llvm_unreachable("should not be used"); })
+          [](auto, auto) -> Type { llvm_unreachable("should not be used"); })
           .applyUnboundGenericArguments(decl, parentTy, SourceLoc(), arguments);
   if (!parentTy && !isTypeResolution) {
     result = DC->mapTypeIntoContext(result);

lib/Sema/PreCheckExpr.cpp

@@ -400,6 +400,10 @@ static BinaryExpr *getCompositionExpr(Expr *expr) {
 }
 
 static Expr *getPackExpansion(DeclContext *dc, Expr *expr, SourceLoc opLoc) {
+  // FIXME: The parser should create PackExpansionExprs directly, pack
+  // elements should be discovered via PackExpansionExpr::getExpandedPacks,
+  // and the generic environment should be created lazily when solving
+  // PackElementOf constraints.
   struct PackReferenceFinder : public ASTWalker {
     DeclContext *dc;
     llvm::SmallVector<PackElementExpr *, 2> packElements;

lib/Sema/TypeCheckType.cpp

@@ -4288,7 +4288,7 @@ NeverNullType TypeResolver::resolvePackReference(PackReferenceTypeRepr *repr,
   // Open the pack reference to an element archetype if requested.
   if (auto openPackElement = resolution.getPackElementOpener()) {
     auto *env = resolution.getGenericSignature().getGenericEnvironment();
-    return openPackElement(env->mapTypeIntoContext(packReference));
+    return openPackElement(env->mapTypeIntoContext(packReference), repr);
   }
 
   return packReference;

lib/Sema/TypeCheckType.h

@@ -482,7 +482,8 @@ using HandlePlaceholderTypeReprFn =
 
 /// A function reference used to replace pack references with opened
 /// element archetypes when resolving a \c PackReferenceTypeRepr.
-using OpenPackElementFn = llvm::function_ref<Type(Type)>;
+using OpenPackElementFn =
+    llvm::function_ref<Type(Type, PackReferenceTypeRepr *)>;
 
 /// Handles the resolution of types within a given declaration context,
 /// which might involve resolving generic parameters to a particular