[Constraint System] Fix the shape class and context substitiutions for

opened element generic environments containing same-element requirements.

Author: hborla

lib/AST/GenericSignature.cpp

@@ -105,8 +105,16 @@ void GenericSignatureImpl::forEachParam(
 
   for (auto req : getRequirements()) {
     GenericTypeParamType *gp;
+    bool isCanonical = false;
     switch (req.getKind()) {
     case RequirementKind::SameType: {
+      if (req.getSecondType()->isParameterPack() != 
+          req.getFirstType()->isParameterPack()) {
+        // This is a same-element requirement, which does not make
+        // type parameters non-canonical.
+        isCanonical = true;
+      }
+
       if (auto secondGP = req.getSecondType()->getAs<GenericTypeParamType>()) {
         // If two generic parameters are same-typed, then the right-hand one
         // is non-canonical.
@@ -136,7 +144,7 @@ void GenericSignatureImpl::forEachParam(
     }
 
     unsigned index = GenericParamKey(gp).findIndexIn(genericParams);
-    genericParamsAreCanonical[index] = false;
+    genericParamsAreCanonical[index] = isCanonical;
   }
 
   // Call the callback with each parameter and the result of the above analysis.

lib/Sema/CSApply.cpp

@@ -3843,6 +3843,11 @@ namespace {
       auto *locator = cs.getConstraintLocator(expr);
       auto *environment = cs.getPackElementEnvironment(locator,
           expansionTy->getCountType()->getCanonicalType());
+
+      // Assert that we have an opened element environment, otherwise we'll get
+      // an ASTVerifier crash when pack archetypes or element archetypes appear
+      // inside the pack expansion expression.
+      assert(environment);
       expr->setGenericEnvironment(environment);
 
       return expr;

lib/Sema/CSGen.cpp

@@ -3143,18 +3143,9 @@ namespace {
         auto expansionType =
             CS.getType(packEnvironment)->castTo<PackExpansionType>();
         CS.addConstraint(ConstraintKind::ShapeOf, expansionType->getCountType(),
-                         elementType,
+                         packType,
                          CS.getConstraintLocator(packEnvironment,
                                                  ConstraintLocator::PackShape));
-        auto *elementShape = CS.createTypeVariable(
-            CS.getConstraintLocator(expr, ConstraintLocator::PackShape),
-            TVO_CanBindToPack);
-        CS.addConstraint(
-            ConstraintKind::ShapeOf, elementShape, elementType,
-            CS.getConstraintLocator(expr, ConstraintLocator::PackShape));
-        CS.addConstraint(
-            ConstraintKind::Equal, elementShape, expansionType->getCountType(),
-            CS.getConstraintLocator(expr, ConstraintLocator::PackShape));
       } else {
         CS.recordFix(AllowInvalidPackReference::create(
             CS, packType, CS.getConstraintLocator(expr->getPackRefExpr())));

lib/Sema/CSSimplify.cpp

@@ -13585,6 +13585,16 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
     return SolutionKind::Solved;
   }
 
+  if (isSingleUnlabeledPackExpansionTuple(packTy)) {
+    auto *packVar =
+        createTypeVariable(getConstraintLocator(locator), TVO_CanBindToPack);
+    addConstraint(ConstraintKind::MaterializePackExpansion, packTy,
+                  packVar,
+                  getConstraintLocator(locator, {ConstraintLocator::Member}));
+    addConstraint(ConstraintKind::ShapeOf, shapeTy, packVar, locator);
+    return SolutionKind::Solved;
+  }
+
   // Map element archetypes to the pack context to check for equality.
   if (packTy->hasElementArchetype()) {
     auto *packEnv = DC->getGenericEnvironmentOfContext();

test/Constraints/pack-expansion-expressions.swift

@@ -63,8 +63,7 @@ func outerArchetype<each T, U>(t: repeat each T, u: U) where repeat each T: P {
 }
 
 func sameElement<each T, U>(t: repeat each T, u: U) where repeat each T: P, repeat each T == U {
-  // FIXME
-//  let _: (repeat each T) = (repeat (each t).f(u))
+  let _ = (repeat (each t).f(u))
 }
 
 func forEachEach<each C, U>(c: repeat each C, function: (U) -> Void)