[Sema] Allow TreatArrayLiteralAsDictionary fix to handle literals with more than one element

Author: LucianoPAlmeida

include/swift/Sema/CSFix.h

@@ -708,10 +708,13 @@ class TreatArrayLiteralAsDictionary final : public ContextualMismatch {
   }
 
   bool diagnose(const Solution &solution, bool asNote = false) const override;
+  bool diagnoseForAmbiguity(CommonFixesArray commonFixes) const override {
+    return diagnose(*commonFixes.front().first);
+  }
 
-  static TreatArrayLiteralAsDictionary *create(ConstraintSystem &cs,
-                                               Type dictionaryTy, Type arrayTy,
-                                               ConstraintLocator *loc);
+  static TreatArrayLiteralAsDictionary *attempt(ConstraintSystem &cs,
+                                                Type dictionaryTy, Type arrayTy,
+                                                ConstraintLocator *loc);
 
   static bool classof(ConstraintFix *fix) {
     return fix->getKind() == FixKind::TreatArrayLiteralAsDictionary;

lib/Sema/CSDiagnostics.cpp

@@ -920,8 +920,21 @@ bool ArrayLiteralToDictionaryConversionFailure::diagnoseAsError() {
                  CTP == CTP_Initialization);
 
   auto diagnostic = emitDiagnostic(diag::meant_dictionary_lit);
-  if (AE->getNumElements() == 1)
+  const auto numElements = AE->getNumElements();
+  if (numElements == 1) {
     diagnostic.fixItInsertAfter(AE->getElement(0)->getEndLoc(), ": <#value#>");
+  } else {
+    // If there is an even number of elements in the array, let's produce
+    // a fix-it which suggests to replace "," with ":" to form a dictionary
+    // literal.
+    if ((numElements & 1) == 0) {
+      const auto commaLocs = AE->getCommaLocs();
+      if (commaLocs.size() == numElements - 1) {
+        for (unsigned i = 0, e = numElements / 2; i != e; ++i)
+          diagnostic.fixItReplace(commaLocs[i * 2], ":");
+      }
+    }
+  }
   return true;
 }
 

lib/Sema/CSFix.cpp

@@ -173,12 +173,37 @@ bool TreatArrayLiteralAsDictionary::diagnose(const Solution &solution,
 }
 
 TreatArrayLiteralAsDictionary *
-TreatArrayLiteralAsDictionary::create(ConstraintSystem &cs,
-                                      Type dictionaryTy, Type arrayTy,
-                                      ConstraintLocator *locator) {
-  assert(getAsExpr<ArrayExpr>(locator->getAnchor())->getNumElements() <= 1);
+TreatArrayLiteralAsDictionary::attempt(ConstraintSystem &cs, Type dictionaryTy,
+                                       Type arrayTy,
+                                       ConstraintLocator *locator) {
+  if (!cs.isArrayType(arrayTy))
+    return nullptr;
+
+  // Determine the ArrayExpr from the locator.
+  auto *expr = getAsExpr(simplifyLocatorToAnchor(locator));
+  if (!expr)
+    return nullptr;
+
+  if (auto *AE = dyn_cast<AssignExpr>(expr))
+    expr = AE->getSrc();
+
+  auto *arrayExpr = dyn_cast<ArrayExpr>(expr);
+  if (!arrayExpr)
+    return nullptr;
+
+  // This fix only applies if the array is used as a dictionary.
+  auto unwrappedDict = dictionaryTy->lookThroughAllOptionalTypes();
+  if (unwrappedDict->isTypeVariableOrMember())
+    return nullptr;
+
+  if (!TypeChecker::conformsToKnownProtocol(
+          unwrappedDict, KnownProtocolKind::ExpressibleByDictionaryLiteral,
+          cs.DC->getParentModule()))
+    return nullptr;
+
+  auto arrayLoc = cs.getConstraintLocator(arrayExpr);
   return new (cs.getAllocator())
-      TreatArrayLiteralAsDictionary(cs, dictionaryTy, arrayTy, locator);
+      TreatArrayLiteralAsDictionary(cs, dictionaryTy, arrayTy, arrayLoc);
 }
 
 bool MarkExplicitlyEscaping::diagnose(const Solution &solution,

lib/Sema/CSSimplify.cpp

@@ -4225,52 +4225,25 @@ static bool repairArrayLiteralUsedAsDictionary(
     ConstraintKind matchKind,
     SmallVectorImpl<RestrictionOrFix> &conversionsOrFixes,
     ConstraintLocator *loc) {
+  if (auto *fix = TreatArrayLiteralAsDictionary::attempt(cs, dictType,
+                                                         arrayType, loc)) {
+    // Ignore any attempts at promoting the value to an optional as even after
+    // stripping off all optionals above the underlying types don't match (array
+    // vs dictionary).
+    conversionsOrFixes.erase(
+        llvm::remove_if(conversionsOrFixes,
+                        [&](RestrictionOrFix &E) {
+                          if (auto restriction = E.getRestriction())
+                            return *restriction ==
+                                   ConversionRestrictionKind::ValueToOptional;
+                          return false;
+                        }),
+        conversionsOrFixes.end());
 
-  if (!cs.isArrayType(arrayType))
-    return false;
-
-  // Determine the ArrayExpr from the locator.
-  auto *expr = getAsExpr(simplifyLocatorToAnchor(loc));
-  if (!expr)
-    return false;
-
-  if (auto *AE = dyn_cast<AssignExpr>(expr))
-    expr = AE->getSrc();
-
-  auto *arrayExpr = dyn_cast<ArrayExpr>(expr);
-  if (!arrayExpr)
-    return false;
-
-  // This fix currently only handles empty and single-element arrays:
-  //   [] => [:] and [1] => [1:_]
-  if (arrayExpr->getNumElements() > 1)
-    return false;
-
-  // This fix only applies if the array is used as a dictionary.
-  auto unwrappedDict = dictType->lookThroughAllOptionalTypes();
-  if (unwrappedDict->isTypeVariableOrMember())
-    return false;
-
-  if (!TypeChecker::conformsToKnownProtocol(
-          unwrappedDict,
-          KnownProtocolKind::ExpressibleByDictionaryLiteral,
-          cs.DC->getParentModule()))
-    return false;
-
-  // Ignore any attempts at promoting the value to an optional as even after
-  // stripping off all optionals above the underlying types don't match (array
-  // vs dictionary).
-  conversionsOrFixes.erase(llvm::remove_if(conversionsOrFixes,
-                                           [&](RestrictionOrFix &E) {
-    if (auto restriction = E.getRestriction())
-      return *restriction == ConversionRestrictionKind::ValueToOptional;
-    return false;
-  }), conversionsOrFixes.end());
-
-  auto argLoc = cs.getConstraintLocator(arrayExpr);
-  conversionsOrFixes.push_back(TreatArrayLiteralAsDictionary::create(
-      cs, dictType, arrayType, argLoc));
-  return true;
+    conversionsOrFixes.push_back(fix);
+    return true;
+  }
+  return false;
 }
 
 /// Let's check whether this is an out-of-order argument in binary
@@ -5563,6 +5536,17 @@ bool ConstraintSystem::repairFailures(
       if (hasConversionOrRestriction(ConversionRestrictionKind::DeepEquality))
         break;
 
+      // We already have a fix for trying to initialize/assign an array literal
+      // to a dictionary type. In this case elements mismatch only add extra
+      // verbosity to the diagnostic. So let's skip the fix and only increase
+      // the score to focus on suggesting using dictionary literal instead.
+      path.pop_back();
+      auto loc = getConstraintLocator(anchor, path);
+      if (hasFixFor(loc, FixKind::TreatArrayLiteralAsDictionary)) {
+        increaseScore(SK_Fix);
+        return true;
+      }
+
       conversionsOrFixes.push_back(CollectionElementContextualMismatch::create(
           *this, lhs, rhs, getConstraintLocator(locator)));
       break;

test/Constraints/dictionary_literal.swift

@@ -85,28 +85,24 @@ var _: [Float: Int] = [1] // expected-error {{dictionary of type '[Float : Int]'
 
 var _: [Int: Int] = ["foo"] // expected-error {{dictionary of type '[Int : Int]' cannot be initialized with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{27-27=: <#value#>}}
-// expected-error@-2 {{cannot convert value of type 'String' to expected dictionary key type 'Int'}}
 
 var _ = useDictStringInt(["Key"]) // expected-error {{dictionary of type 'DictStringInt' cannot be used with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{32-32=: <#value#>}}
 
 var _ = useDictStringInt([4]) // expected-error {{dictionary of type 'DictStringInt' cannot be used with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{28-28=: <#value#>}}
-// expected-error@-2 {{cannot convert value of type 'Int' to expected dictionary key type 'DictStringInt.Key' (aka 'String')}}
 
 var _: [[Int: Int]] = [[5]] // expected-error {{dictionary of type '[Int : Int]' cannot be used with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{26-26=: <#value#>}}
 
 var _: [[Int: Int]] = [["bar"]] // expected-error {{dictionary of type '[Int : Int]' cannot be used with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{30-30=: <#value#>}}
-// expected-error@-2 {{cannot convert value of type 'String' to expected dictionary key type 'Int'}}
 
 assignDict = [1] // expected-error {{dictionary of type '[Int : Int]' cannot be used with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{16-16=: <#value#>}}
 
 assignDict = [""] // expected-error {{dictionary of type '[Int : Int]' cannot be used with array literal}}
 // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{17-17=: <#value#>}}
-// expected-error@-2 {{cannot convert value of type 'String' to expected dictionary key type 'Int'}}
 
 func arrayLiteralDictionaryMismatch<T>(a: inout T) where T: ExpressibleByDictionaryLiteral, T.Key == Int, T.Value == Int {
   a = [] // expected-error {{use [:] to get an empty dictionary literal}} {{8-8=:}}
@@ -116,7 +112,6 @@ func arrayLiteralDictionaryMismatch<T>(a: inout T) where T: ExpressibleByDiction
 
   a = [""] // expected-error {{dictionary of type 'T' cannot be used with array literal}}
   // expected-note@-1 {{did you mean to use a dictionary literal instead?}} {{10-10=: <#value#>}}
-  // expected-error@-2 {{cannot convert value of type 'String' to expected dictionary key type 'Int'}}
 }
 
 
@@ -167,3 +162,33 @@ func rdar32330004_2() -> [String: Any] {
   // expected-error@-1 {{dictionary of type '[String : Any]' cannot be used with array literal}}
   // expected-note@-2 {{did you mean to use a dictionary literal instead?}} {{14-15=:}} {{24-25=:}} {{34-35=:}} {{46-47=:}}
 }
+
+// https://github.com/apple/swift/issues/59215
+class S59215 {
+  var m: [String: [String: String]] = [:]
+  init() {
+    m["a"] = ["", 1] //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+    // expected-note@-1{{did you mean to use a dictionary literal instead?}} {{17-18=:}}
+    m["a"] = [1 , ""] //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+    // expected-note@-1{{did you mean to use a dictionary literal instead?}} {{17-18=:}}
+    m["a"] = ["", ""] //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+    // expected-note@-1{{did you mean to use a dictionary literal instead?}} {{17-18=:}}
+    m["a"] = [1 , 1] //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+    // expected-note@-1{{did you mean to use a dictionary literal instead?}} {{17-18=:}}
+  }
+}
+
+func f59215(_ a: [String: String]) {}
+f59215(["", 1]) //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+// expected-note@-1{{did you mean to use a dictionary literal instead?}} {{11-12=:}}
+f59215([1 , ""]) //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+// expected-note@-1{{did you mean to use a dictionary literal instead?}} {{11-12=:}}
+f59215([1 , 1]) //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+// expected-note@-1{{did you mean to use a dictionary literal instead?}} {{11-12=:}}
+f59215(["", ""]) //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+// expected-note@-1{{did you mean to use a dictionary literal instead?}} {{11-12=:}}
+
+f59215(["", "", "", ""]) //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+// expected-note@-1{{did you mean to use a dictionary literal instead?}} {{11-12=:}} {{19-20=:}}
+f59215(["", "", "", ""]) //expected-error{{dictionary of type '[String : String]' cannot be used with array literal}}
+// expected-note@-1{{did you mean to use a dictionary literal instead?}} {{11-12=:}} {{19-20=:}}

test/Constraints/optional.swift

@@ -506,7 +506,7 @@ func rdar85166519() {
 
   var _: [Int: AnyObject] = [ // expected-error {{dictionary of type '[Int : AnyObject]' cannot be initialized with array literal}}
     // expected-note@-1 {{did you mean to use a dictionary literal instead?}}
-    v?.addingReportingOverflow(0) // expected-error {{cannot convert value of type '(partialValue: Int, overflow: Bool)?' to expected dictionary key type 'Int'}}
+    v?.addingReportingOverflow(0)
   ]
 }