Merge pull request #40159 from xedin/rdar-85021348

[Diagnostics] Determine affected elements early in collection element…

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(