[ConstraintSystem] Record both the UUID and the shape class for opened pack

element environments.

This allows the constraint system to ensure that for a given pack expansion locator,
the given shape class is always the same when requesting the element environment.
If the shape class differs, it means there's a same-shape requirement failure, which
will be diagnosed via the ShapeOf constraint simplification.

Author: hborla

include/swift/Sema/ConstraintSystem.h

@@ -1505,7 +1505,8 @@ class Solution {
 
   /// The pack expansion environment that can open pack elements for
   /// a given locator.
-  llvm::DenseMap<ConstraintLocator *, UUID> PackExpansionEnvironments;
+  llvm::DenseMap<ConstraintLocator *, std::pair<UUID, Type>>
+      PackExpansionEnvironments;
 
   /// The locators of \c Defaultable constraints whose defaults were used.
   llvm::SmallPtrSet<ConstraintLocator *, 2> DefaultedConstraints;
@@ -3067,7 +3068,8 @@ class ConstraintSystem {
   llvm::SmallMapVector<ConstraintLocator *, OpenedArchetypeType *, 4>
       OpenedExistentialTypes;
 
-  llvm::SmallMapVector<ConstraintLocator *, UUID, 4> PackExpansionEnvironments;
+  llvm::SmallMapVector<ConstraintLocator *, std::pair<UUID, Type>, 4>
+      PackExpansionEnvironments;
 
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
@@ -4031,9 +4033,6 @@ class ConstraintSystem {
   std::pair<Type, OpenedArchetypeType *> openExistentialType(
       Type type, ConstraintLocator *locator);
 
-  /// Add the given pack expansion as an opened pack element environment.
-  void addPackElementEnvironment(PackExpansionExpr *expr);
-
   /// Get the opened element generic environment for the given locator.
   GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
                                                 CanType shapeClass);

lib/AST/ASTContext.cpp

@@ -5147,6 +5147,7 @@ GenericEnvironment::forOpenedElement(GenericSignature signature,
   if (found != openedElementEnvironments.end()) {
     auto *existingEnv = found->second;
     assert(existingEnv->getGenericSignature().getPointer() == signature.getPointer());
+    assert(existingEnv->getOpenedElementShapeClass()->isEqual(shapeClass));
     assert(existingEnv->getOpenedElementUUID() == uuid);
 
     return existingEnv;

lib/Sema/CSGen.cpp

@@ -1101,7 +1101,6 @@ namespace {
     ConstraintSystem &getConstraintSystem() const { return CS; }
 
     void addPackElementEnvironment(PackExpansionExpr *expr) {
-      CS.addPackElementEnvironment(expr);
       PackElementEnvironments.push_back(expr);
     }
 

lib/Sema/ConstraintSystem.cpp

@@ -649,29 +649,31 @@ std::pair<Type, OpenedArchetypeType *> ConstraintSystem::openExistentialType(
   return {result, opened};
 }
 
-void ConstraintSystem::addPackElementEnvironment(PackExpansionExpr *expr) {
-  auto *locator = getConstraintLocator(expr);
-  PackExpansionEnvironments[locator] = UUID::fromTime();
-}
-
 GenericEnvironment *
 ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
                                             CanType shapeClass) {
+  assert(locator->directlyAt<PackExpansionExpr>());
+
+  std::pair<UUID, Type> uuidAndShape;
   auto result = PackExpansionEnvironments.find(locator);
-  if (result == PackExpansionEnvironments.end())
-    return nullptr;
+  if (result == PackExpansionEnvironments.end()) {
+    uuidAndShape = std::make_pair(UUID::fromTime(), shapeClass);
+    PackExpansionEnvironments[locator] = uuidAndShape;
+  } else {
+    uuidAndShape = result->second;
+  }
 
-  if (!shapeClass->is<PackArchetypeType>())
+  if (!shapeClass->is<PackArchetypeType>() ||
+      !shapeClass->isEqual(uuidAndShape.second))
     return nullptr;
 
-  auto uuid = result->second;
   auto &ctx = getASTContext();
   auto elementSig = ctx.getOpenedElementSignature(
       DC->getGenericSignatureOfContext().getCanonicalSignature(), shapeClass);
   auto *contextEnv = DC->getGenericEnvironmentOfContext();
   auto contextSubs = contextEnv->getForwardingSubstitutionMap();
-  return GenericEnvironment::forOpenedElement(elementSig, uuid, shapeClass,
-                                              contextSubs);
+  return GenericEnvironment::forOpenedElement(elementSig, uuidAndShape.first,
+                                              shapeClass, contextSubs);
 }
 
 /// Extend the given depth map by adding depths for all of the subexpressions