[CS] Clean up pack expansion environment handling a little

- Track environments for `PackExpansionExpr` directly
instead of using a locator.
- Split up the querying and creation of the environment
such that the mismatch logic can be done directly in
CSSimplify instead of duplicating it.
- Just store the environment directly instead of
the shape and UUID.

Author: hamishknight

include/swift/Sema/CSTrail.def

@@ -55,7 +55,6 @@ LOCATOR_CHANGE(RecordedAppliedDisjunction, AppliedDisjunctions)
 LOCATOR_CHANGE(RecordedMatchCallArgumentResult, argumentMatchingChoices)
 LOCATOR_CHANGE(RecordedOpenedTypes, OpenedTypes)
 LOCATOR_CHANGE(RecordedOpenedExistentialType, OpenedExistentialTypes)
-LOCATOR_CHANGE(RecordedPackExpansionEnvironment, PackExpansionEnvironments)
 LOCATOR_CHANGE(RecordedDefaultedConstraint, DefaultedConstraints)
 LOCATOR_CHANGE(ResolvedOverload, ResolvedOverloads)
 LOCATOR_CHANGE(RecordedArgumentList, ArgumentLists)
@@ -96,6 +95,7 @@ CHANGE(AddedFix)
 CHANGE(AddedFixedRequirement)
 CHANGE(RecordedOpenedPackExpansionType)
 CHANGE(RecordedPackElementExpansion)
+CHANGE(RecordedPackExpansionEnvironment)
 CHANGE(RecordedNodeType)
 CHANGE(RecordedKeyPathComponentType)
 CHANGE(RecordedResultBuilderTransform)

include/swift/Sema/CSTrail.h

@@ -150,6 +150,7 @@ class SolverTrail {
       ConstraintFix *TheFix;
       ConstraintLocator *TheLocator;
       PackExpansionType *TheExpansion;
+      PackExpansionExpr *TheExpansionExpr;
       PackElementExpr *TheElement;
       Expr *TheExpr;
       Stmt *TheStmt;
@@ -215,6 +216,10 @@ class SolverTrail {
     /// to its parent expansion expression.
     static Change RecordedPackElementExpansion(PackElementExpr *packElement);
 
+    /// Create a change that records the GenericEnvironment for a given
+    /// PackExpansionExpr.
+    static Change RecordedPackExpansionEnvironment(PackExpansionExpr *expr);
+
     /// Create a change that recorded an assignment of a type to an AST node.
     static Change RecordedNodeType(ASTNode node, Type oldType);
 

include/swift/Sema/ConstraintSystem.h

@@ -1560,9 +1560,9 @@ class Solution {
   llvm::DenseMap<PackExpansionType *, TypeVariableType *>
       OpenedPackExpansionTypes;
 
-  /// The pack expansion environment that can open pack elements for
-  /// a given locator.
-  llvm::DenseMap<ConstraintLocator *, std::pair<UUID, Type>>
+  /// The generic environment that can open pack elements for a given
+  /// pack expansion.
+  llvm::DenseMap<PackExpansionExpr *, GenericEnvironment *>
       PackExpansionEnvironments;
 
   /// The pack expansion expression for a given pack element.
@@ -1810,6 +1810,11 @@ class Solution {
     return Type();
   }
 
+  /// Retrieve the generic environment for the opened element of a given pack
+  /// expansion, or \c nullptr if no environment was recorded.
+  GenericEnvironment *
+  getPackExpansionEnvironment(PackExpansionExpr *expr) const;
+
   /// For a given locator describing a function argument conversion, or a
   /// constraint within an argument conversion, returns information about the
   /// application of the argument to its parameter. If the locator is not
@@ -2407,7 +2412,7 @@ class ConstraintSystem {
   llvm::SmallDenseMap<PackExpansionType *, TypeVariableType *, 4>
       OpenedPackExpansionTypes;
 
-  llvm::SmallDenseMap<ConstraintLocator *, std::pair<UUID, Type>, 4>
+  llvm::SmallDenseMap<PackExpansionExpr *, GenericEnvironment *, 4>
       PackExpansionEnvironments;
 
   llvm::SmallDenseMap<PackElementExpr *, PackExpansionExpr *, 2>
@@ -3370,13 +3375,26 @@ class ConstraintSystem {
   void recordOpenedExistentialType(ConstraintLocator *locator,
                                    OpenedArchetypeType *opened);
 
-  /// Get the opened element generic environment for the given locator.
-  GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
-                                                CanType shapeClass);
+  /// Retrieve the generic environment for the opened element of a given pack
+  /// expansion, or \c nullptr if no environment was recorded yet.
+  GenericEnvironment *
+  getPackExpansionEnvironment(PackExpansionExpr *expr) const;
+
+  /// Create a new opened element generic environment for the given pack
+  /// expansion.
+  GenericEnvironment *
+  createPackExpansionEnvironment(PackExpansionExpr *expr,
+                                 CanGenericTypeParamType shapeParam);
 
   /// Update PackExpansionEnvironments and record a change in the trail.
-  void recordPackExpansionEnvironment(ConstraintLocator *locator,
-                                      std::pair<UUID, Type> uuidAndShape);
+  void recordPackExpansionEnvironment(PackExpansionExpr *expr,
+                                      GenericEnvironment *env);
+
+  /// Undo the above change.
+  void removePackExpansionEnvironment(PackExpansionExpr *expr) {
+    bool erased = PackExpansionEnvironments.erase(expr);
+    ASSERT(erased);
+  }
 
   /// Get the pack expansion expr for the given pack element.
   PackExpansionExpr *

lib/Sema/CSApply.cpp

@@ -3942,21 +3942,15 @@ namespace {
     }
 
     Expr *visitPackExpansionExpr(PackExpansionExpr *expr) {
-      simplifyExprType(expr);
-
       // Set the opened pack element environment for this pack expansion.
-      auto expansionTy = cs.getType(expr)->castTo<PackExpansionType>();
-      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.
+      auto *environment = solution.getPackExpansionEnvironment(expr);
       assert(environment);
       expr->setGenericEnvironment(environment);
 
-      return expr;
+      return simplifyExprType(expr);
     }
 
     Expr *visitPackElementExpr(PackElementExpr *expr) {

lib/Sema/CSSimplify.cpp

@@ -9629,43 +9629,57 @@ ConstraintSystem::simplifyBindTupleOfFunctionParamsConstraint(
 ConstraintSystem::SolutionKind
 ConstraintSystem::matchPackElementType(Type elementType, Type patternType,
                                        ConstraintLocatorBuilder locator) {
-  auto *loc = getConstraintLocator(locator);
-  auto shapeClass = patternType->getReducedShape();
-  auto *elementEnv = getPackElementEnvironment(loc, shapeClass);
-
-  // Without an opened element environment, we cannot derive the
-  // element binding.
-  if (!elementEnv) {
+  auto tryFix = [&](llvm::function_ref<ConstraintFix *(void)> fix) {
     if (!shouldAttemptFixes())
       return SolutionKind::Error;
 
-    // `each` was applied to a concrete type.
-    if (!shapeClass->is<PackArchetypeType>()) {
-      if (recordFix(AllowInvalidPackElement::create(*this, patternType, loc)))
-        return SolutionKind::Error;
-    } else {
-      auto envShape = PackExpansionEnvironments.find(loc);
-      if (envShape == PackExpansionEnvironments.end()) {
-        return SolutionKind::Error;
-      }
-      auto *fix = SkipSameShapeRequirement::create(
-          *this, envShape->second.second, shapeClass,
-          getConstraintLocator(loc, ConstraintLocator::PackShape));
-      if (recordFix(fix)) {
-        return SolutionKind::Error;
-      }
-    }
+    if (recordFix(fix()))
+      return SolutionKind::Error;
 
     recordAnyTypeVarAsPotentialHole(elementType);
     return SolutionKind::Solved;
+  };
+
+  auto *loc = getConstraintLocator(locator);
+  ASSERT(loc->directlyAt<PackExpansionExpr>());
+  auto *packExpansion = castToExpr<PackExpansionExpr>(loc->getAnchor());
+
+  ASSERT(!patternType->hasTypeVariable());
+  auto shapeClass = patternType->getReducedShape();
+
+  // `each` was applied to a concrete type.
+  if (!shapeClass->is<PackArchetypeType>()) {
+    return tryFix([&]() {
+      return AllowInvalidPackElement::create(*this, patternType, loc);
+    });
+  }
+
+  auto shapeParam = CanGenericTypeParamType(cast<GenericTypeParamType>(
+      shapeClass->mapTypeOutOfContext()->getCanonicalType()));
+
+  auto *genericEnv = getPackExpansionEnvironment(packExpansion);
+  if (genericEnv) {
+    if (shapeParam != genericEnv->getOpenedElementShapeClass()) {
+      return tryFix([&]() {
+        auto envShape = genericEnv->mapTypeIntoContext(
+            genericEnv->getOpenedElementShapeClass());
+        if (auto *pack = dyn_cast<PackType>(envShape))
+          envShape = pack->unwrapSingletonPackExpansion()->getPatternType();
+
+        return SkipSameShapeRequirement::create(
+            *this, envShape, shapeClass,
+            getConstraintLocator(loc, ConstraintLocator::PackShape));
+      });
+    }
+  } else {
+    genericEnv = createPackExpansionEnvironment(packExpansion, shapeParam);
   }
 
   auto expectedElementTy =
-      elementEnv->mapContextualPackTypeIntoElementContext(patternType);
+      genericEnv->mapContextualPackTypeIntoElementContext(patternType);
   assert(!expectedElementTy->is<PackType>());
 
-  addConstraint(ConstraintKind::Equal, elementType, expectedElementTy,
-                locator);
+  addConstraint(ConstraintKind::Equal, elementType, expectedElementTy, locator);
   return SolutionKind::Solved;
 }
 

lib/Sema/CSTrail.cpp

@@ -199,6 +199,14 @@ SolverTrail::Change SolverTrail::Change::RecordedPackElementExpansion(
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RecordedPackExpansionEnvironment(PackExpansionExpr *expr) {
+  Change result;
+  result.Kind = ChangeKind::RecordedPackExpansionEnvironment;
+  result.TheExpansionExpr = expr;
+  return result;
+}
+
 SolverTrail::Change
 SolverTrail::Change::RecordedNodeType(ASTNode node, Type oldType) {
   Change result;
@@ -430,6 +438,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     cs.removePackElementExpansion(TheElement);
     break;
 
+  case ChangeKind::RecordedPackExpansionEnvironment:
+    cs.removePackExpansionEnvironment(TheExpansionExpr);
+    break;
+
   case ChangeKind::RecordedNodeType:
     cs.restoreType(Node.Node, Node.OldType);
     break;
@@ -702,6 +714,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     out << ")\n";
     break;
 
+  case ChangeKind::RecordedPackExpansionEnvironment:
+    out << "(RecordedPackExpansionEnvironment ";
+    dumpAnchor(TheExpansionExpr, &SM, out);
+    out << ")\n";
+    break;
+
   case ChangeKind::RecordedNodeType:
     out << "(RecordedNodeType at ";
     Node.Node.getStartLoc().print(out, cs.getASTContext().SourceMgr);

lib/Sema/ConstraintSystem.cpp

@@ -877,44 +877,35 @@ void ConstraintSystem::recordOpenedExistentialType(
 }
 
 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()) {
-    uuidAndShape = std::make_pair(UUID::fromTime(), shapeClass);
-    recordPackExpansionEnvironment(locator, uuidAndShape);
-  } else {
-    uuidAndShape = result->second;
-  }
-
-  if (!shapeClass->is<PackArchetypeType>() ||
-      !shapeClass->isEqual(uuidAndShape.second))
+ConstraintSystem::getPackExpansionEnvironment(PackExpansionExpr *expr) const {
+  auto result = PackExpansionEnvironments.find(expr);
+  if (result == PackExpansionEnvironments.end())
     return nullptr;
 
-  auto shapeParam = cast<GenericTypeParamType>(
-      shapeClass->mapTypeOutOfContext()->getCanonicalType());
+  return result->second;
+}
 
-  auto &ctx = getASTContext();
+GenericEnvironment *ConstraintSystem::createPackExpansionEnvironment(
+    PackExpansionExpr *expr, CanGenericTypeParamType shapeParam) {
   auto *contextEnv = PackElementGenericEnvironments.empty()
                          ? DC->getGenericEnvironmentOfContext()
                          : PackElementGenericEnvironments.back();
-  auto elementSig = ctx.getOpenedElementSignature(
+  auto elementSig = getASTContext().getOpenedElementSignature(
       contextEnv->getGenericSignature().getCanonicalSignature(), shapeParam);
   auto contextSubs = contextEnv->getForwardingSubstitutionMap();
-  return GenericEnvironment::forOpenedElement(elementSig, uuidAndShape.first,
-                                              shapeParam, contextSubs);
+  auto *env = GenericEnvironment::forOpenedElement(elementSig, UUID::fromTime(),
+                                                   shapeParam, contextSubs);
+  recordPackExpansionEnvironment(expr, env);
+  return env;
 }
 
-void ConstraintSystem::recordPackExpansionEnvironment(
-    ConstraintLocator *locator, std::pair<UUID, Type> uuidAndShape) {
-  bool inserted = PackExpansionEnvironments.insert({locator, uuidAndShape}).second;
+void ConstraintSystem::recordPackExpansionEnvironment(PackExpansionExpr *expr,
+                                                      GenericEnvironment *env) {
+  bool inserted = PackExpansionEnvironments.insert({expr, env}).second;
   ASSERT(inserted);
 
   if (solverState)
-    recordChange(SolverTrail::Change::RecordedPackExpansionEnvironment(locator));
+    recordChange(SolverTrail::Change::RecordedPackExpansionEnvironment(expr));
 }
 
 PackExpansionExpr *
@@ -4050,6 +4041,15 @@ ASTNode ConstraintSystem::includingParentApply(ASTNode node) {
   return node;
 }
 
+GenericEnvironment *
+Solution::getPackExpansionEnvironment(PackExpansionExpr *expr) const {
+  auto iter = PackExpansionEnvironments.find(expr);
+  if (iter == PackExpansionEnvironments.end())
+    return nullptr;
+
+  return iter->second;
+}
+
 std::optional<FunctionArgApplyInfo>
 Solution::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
   // It's only valid to use `&` in argument positions, but we need

lib/Sema/TypeCheckConstraints.cpp

@@ -1632,11 +1632,11 @@ void ConstraintSystem::print(raw_ostream &out) const {
 
   if (!PackExpansionEnvironments.empty()) {
     out.indent(indent) << "Pack Expansion Environments:\n";
-    for (const auto &env : PackExpansionEnvironments) {
+    for (const auto &[packExpansion, env] : PackExpansionEnvironments) {
       out.indent(indent + 2);
-      env.first->dump(&getASTContext().SourceMgr, out);
-      out << " = (" << env.second.first << ", "
-          << env.second.second->getString(PO) << ")" << '\n';
+      dumpAnchor(packExpansion, &getASTContext().SourceMgr, out);
+      out << " = (" << env->getOpenedElementShapeClass() << ", "
+      << env->getOpenedElementUUID() << ")" << '\n';
     }
   }