[CSGen] Use a ShapeOf constraint when generating constraints for a pack

expansion expression.

Author: hborla

lib/AST/ParameterPack.cpp

@@ -314,15 +314,36 @@ CanPackType PackType::getReducedShape() {
 }
 
 CanType TypeBase::getReducedShape() {
-  if (auto *packArchetype = getAs<PackArchetypeType>())
-    return packArchetype->getReducedShape();
-
   if (auto *packType = getAs<PackType>())
     return packType->getReducedShape();
 
   if (auto *expansionType = getAs<PackExpansionType>())
     return expansionType->getReducedShape();
 
+  struct PatternTypeWalker : public TypeWalker {
+    CanType shapeType;
+
+    Action walkToTypePre(Type type) override {
+      // Don't consider pack references inside nested pack
+      // expansion types.
+      if (type->is<PackExpansionType>()) {
+        return Action::Stop;
+      }
+
+      if (auto *archetype = type->getAs<PackArchetypeType>()) {
+        shapeType = archetype->getReducedShape();
+        return Action::Stop;
+      }
+
+      return Action::Continue;
+    }
+  } patternTypeWalker;
+
+  Type(this).walk(patternTypeWalker);
+
+  if (patternTypeWalker.shapeType)
+    return patternTypeWalker.shapeType;
+
   assert(!isTypeVariableOrMember());
   assert(!hasTypeParameter());
 

lib/Sema/CSGen.cpp

@@ -2925,20 +2925,15 @@ namespace {
       CS.addConstraint(ConstraintKind::PackElementOf, elementResultType,
                        patternTy, CS.getConstraintLocator(expr));
 
-      // FIXME: Use a ShapeOf constraint here.
-      Type shapeType;
-      auto *binding = expr->getBindings().front();
-      auto type = CS.simplifyType(CS.getType(binding));
-      type.visit([&](Type type) {
-        if (shapeType)
-          return;
-
-        if (auto archetype = type->getAs<PackArchetypeType>()) {
-          shapeType = archetype->getReducedShape();
-        }
-      });
+      auto *shapeLoc =
+          CS.getConstraintLocator(expr, ConstraintLocator::PackShape);
+      auto *shapeTypeVar = CS.createTypeVariable(shapeLoc,
+                                                 TVO_CanBindToPack |
+                                                 TVO_CanBindToHole);
+      CS.addConstraint(ConstraintKind::ShapeOf, patternTy, shapeTypeVar,
+                       CS.getConstraintLocator(expr));
 
-      return PackExpansionType::get(patternTy, shapeType);
+      return PackExpansionType::get(patternTy, shapeTypeVar);
     }
 
     Type visitDynamicTypeExpr(DynamicTypeExpr *expr) {