[Diagnostics] Determine affected elements early in collection element mismatches

Diagnostics cannot assume that solution would always be applied
to the constraint system, so all of the elements affected by the
mismatch have to be determined by the fix.

Resolves: rdar://85021348

Author: xedin

include/swift/Sema/CSFix.h

@@ -1858,18 +1858,35 @@ class RemoveReturn final : public ContextualMismatch {
   }
 };
 
-class CollectionElementContextualMismatch final : public ContextualMismatch {
-  CollectionElementContextualMismatch(ConstraintSystem &cs, Type srcType,
-                                      Type dstType, ConstraintLocator *locator)
+class CollectionElementContextualMismatch final
+    : public ContextualMismatch,
+      private llvm::TrailingObjects<CollectionElementContextualMismatch,
+                                    Expr *> {
+  friend TrailingObjects;
+
+  unsigned NumElements;
+
+  CollectionElementContextualMismatch(ConstraintSystem &cs,
+                                      ArrayRef<Expr *> affectedElements,
+                                      Type srcType, Type dstType,
+                                      ConstraintLocator *locator)
       : ContextualMismatch(cs,
                            FixKind::IgnoreCollectionElementContextualMismatch,
-                           srcType, dstType, locator) {}
+                           srcType, dstType, locator),
+        NumElements(affectedElements.size()) {
+    std::uninitialized_copy(affectedElements.begin(), affectedElements.end(),
+                            getElementBuffer().begin());
+  }
 
 public:
   std::string getName() const override {
     return "fix collection element contextual mismatch";
   }
 
+  ArrayRef<Expr *> getElements() const {
+    return {getTrailingObjects<Expr *>(), NumElements};
+  }
+
   bool diagnose(const Solution &solution, bool asNote = false) const override;
 
   static CollectionElementContextualMismatch *
@@ -1879,6 +1896,11 @@ class CollectionElementContextualMismatch final : public ContextualMismatch {
   static bool classof(ConstraintFix *fix) {
     return fix->getKind() == FixKind::IgnoreCollectionElementContextualMismatch;
   }
+
+private:
+  MutableArrayRef<Expr *> getElementBuffer() {
+    return {getTrailingObjects<Expr *>(), NumElements};
+  }
 };
 
 class DefaultGenericArgument final : public ConstraintFix {

lib/Sema/CSDiagnostics.cpp

@@ -5441,6 +5441,13 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
   auto eltType = getFromType();
   auto contextualType = getToType();
 
+  auto diagnoseAllOccurances = [&](Diag<Type, Type> diagnostic) {
+    assert(AffectedElements.size() > 1);
+    for (auto *element : AffectedElements) {
+      emitDiagnosticAt(element->getLoc(), diagnostic, eltType, contextualType);
+    }
+  };
+
   auto isFixedToDictionary = [&](ArrayExpr *anchor) {
     return llvm::any_of(getSolution().Fixes, [&](ConstraintFix *fix) {
       auto *fixAnchor = getAsExpr<ArrayExpr>(fix->getAnchor());
@@ -5453,8 +5460,10 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
   Optional<InFlightDiagnostic> diagnostic;
   if (auto *AE = getAsExpr<ArrayExpr>(anchor)) {
     if (!(treatAsDictionary = isFixedToDictionary(AE))) {
-      if (diagnoseMergedLiteralElements())
+      if (AffectedElements.size() > 1) {
+        diagnoseAllOccurances(diag::cannot_convert_array_element);
         return true;
+      }
 
       diagnostic.emplace(emitDiagnostic(diag::cannot_convert_array_element,
                                         eltType, contextualType));
@@ -5464,15 +5473,27 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
   if (treatAsDictionary || isExpr<DictionaryExpr>(anchor)) {
     auto eltLoc = locator->castLastElementTo<LocatorPathElt::TupleElement>();
     switch (eltLoc.getIndex()) {
-    case 0: // key
+    case 0: { // key
+      if (AffectedElements.size() > 1) {
+        diagnoseAllOccurances(diag::cannot_convert_dict_key);
+        return true;
+      }
+
       diagnostic.emplace(emitDiagnostic(diag::cannot_convert_dict_key, eltType,
                                         contextualType));
       break;
+    }
+
+    case 1: { // value
+      if (AffectedElements.size() > 1) {
+        diagnoseAllOccurances(diag::cannot_convert_dict_value);
+        return true;
+      }
 
-    case 1: // value
       diagnostic.emplace(emitDiagnostic(diag::cannot_convert_dict_value,
                                         eltType, contextualType));
       break;
+    }
 
     default:
       break;
@@ -5508,30 +5529,6 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
   return true;
 }
 
-bool CollectionElementContextualFailure::diagnoseMergedLiteralElements() {
-  auto elementAnchor = simplifyLocatorToAnchor(getLocator());
-  if (!elementAnchor)
-    return false;
-
-  auto *typeVar = getRawType(elementAnchor)->getAs<TypeVariableType>();
-  if (!typeVar || !typeVar->getImpl().getAtomicLiteralKind())
-    return false;
-
-  // This element is a literal whose type variable could have been merged with others,
-  // but the conversion constraint to the array element type was only placed on one
-  // of them. So, we want to emit the error for each element whose type variable is in
-  // this equivalence class.
-  auto &cs = getConstraintSystem();
-  auto node = cs.getRepresentative(typeVar)->getImpl().getGraphNode();
-  for (const auto *typeVar : node->getEquivalenceClass()) {
-    auto anchor = typeVar->getImpl().getLocator()->getAnchor();
-    emitDiagnosticAt(constraints::getLoc(anchor), diag::cannot_convert_array_element,
-                     getFromType(), getToType());
-  }
-
-  return true;
-}
-
 bool MissingContextualConformanceFailure::diagnoseAsError() {
   auto anchor = getAnchor();
   auto path = getLocator()->getPath();

lib/Sema/CSDiagnostics.h

@@ -1761,15 +1761,18 @@ class ExtraneousReturnFailure final : public FailureDiagnostic {
 /// let _: [Int] = ["hello"]
 /// ```
 class CollectionElementContextualFailure final : public ContextualFailure {
+  SmallVector<Expr *, 4> AffectedElements;
+
 public:
-  CollectionElementContextualFailure(const Solution &solution, Type eltType,
-                                     Type contextualType,
+  CollectionElementContextualFailure(const Solution &solution,
+                                     ArrayRef<Expr *> affectedElements,
+                                     Type eltType, Type contextualType,
                                      ConstraintLocator *locator)
-      : ContextualFailure(solution, eltType, contextualType, locator) {}
+      : ContextualFailure(solution, eltType, contextualType, locator) {
+    AffectedElements.append(affectedElements.begin(), affectedElements.end());
+  }
 
   bool diagnoseAsError() override;
-
-  bool diagnoseMergedLiteralElements();
 };
 
 class MissingContextualConformanceFailure final : public ContextualFailure {

lib/Sema/CSFix.cpp

@@ -1128,17 +1128,40 @@ RemoveReturn *RemoveReturn::create(ConstraintSystem &cs, Type resultTy,
 
 bool CollectionElementContextualMismatch::diagnose(const Solution &solution,
                                                    bool asNote) const {
-  CollectionElementContextualFailure failure(solution, getFromType(),
-                                             getToType(), getLocator());
+  CollectionElementContextualFailure failure(
+      solution, getElements(), getFromType(), getToType(), getLocator());
   return failure.diagnose(asNote);
 }
 
 CollectionElementContextualMismatch *
 CollectionElementContextualMismatch::create(ConstraintSystem &cs, Type srcType,
                                             Type dstType,
                                             ConstraintLocator *locator) {
-  return new (cs.getAllocator())
-      CollectionElementContextualMismatch(cs, srcType, dstType, locator);
+  // It's common for a single literal element to represent types of other
+  // literal elements of the same kind, let's check whether that is the case
+  // here and record all of the affected positions.
+
+  SmallVector<Expr *, 4> affected;
+  {
+    if (auto *elementLoc = getAsExpr(simplifyLocatorToAnchor(locator))) {
+      auto *typeVar = cs.getType(elementLoc)->getAs<TypeVariableType>();
+      if (typeVar && typeVar->getImpl().getAtomicLiteralKind()) {
+        const auto *node =
+            cs.getRepresentative(typeVar)->getImpl().getGraphNode();
+        for (auto *typeVar : node->getEquivalenceClass()) {
+          auto *locator = typeVar->getImpl().getLocator();
+          if (auto *eltLoc = getAsExpr(simplifyLocatorToAnchor(locator)))
+            affected.push_back(eltLoc);
+        }
+      }
+    }
+  }
+
+  unsigned size = totalSizeToAlloc<Expr *>(affected.size());
+  void *mem = cs.getAllocator().Allocate(
+      size, alignof(CollectionElementContextualMismatch));
+  return new (mem) CollectionElementContextualMismatch(cs, affected, srcType,
+                                                       dstType, locator);
 }
 
 bool DefaultGenericArgument::coalesceAndDiagnose(