[CSFix] Creating object literal module import fix

Author: LucianoPAlmeida

lib/Sema/CSDiagnostics.cpp

@@ -6038,3 +6038,53 @@ bool UnableToInferClosureReturnType::diagnoseAsError() {
 
   return true;
 }
+
+static std::tuple<StringRef, StringRef>
+getImportModuleAndDefaultType(const ASTContext &ctx, ProtocolDecl *protocol) {
+  const auto &target = ctx.LangOpts.Target;
+  if (protocol ==
+      ctx.getProtocol(KnownProtocolKind::ExpressibleByColorLiteral)) {
+    if (target.isMacOSX()) {
+      return std::make_tuple("AppKit", "NSColor");
+    } else if (target.isiOS() || target.isTvOS()) {
+      return std::make_tuple("UIKit", "UIColor");
+    }
+  } else if (protocol ==
+             ctx.getProtocol(KnownProtocolKind::ExpressibleByImageLiteral)) {
+    if (target.isMacOSX()) {
+      return std::make_tuple("AppKit", "NSImage");
+    } else if (target.isiOS() || target.isTvOS()) {
+      return std::make_tuple("UIKit", "UIImage");
+    }
+  } else if (protocol ==
+             ctx.getProtocol(
+                 KnownProtocolKind::ExpressibleByFileReferenceLiteral)) {
+    return std::make_tuple("Foundation", "URL");
+  }
+  return std::make_tuple("", "");
+}
+
+bool UnableToInferProtocolLiteralType::diagnoseAsError() {
+  auto *locator = getLocator();
+  auto &cs = getConstraintSystem();
+  auto *expr = cast<ObjectLiteralExpr>(locator->getAnchor());
+
+  auto &ctx = cs.getASTContext();
+  auto protocol = TypeChecker::getLiteralProtocol(ctx, expr);
+  assert(protocol);
+  auto plainName = expr->getLiteralKindPlainName();
+
+  StringRef importModule;
+  StringRef importDefaultTypeName;
+  std::tie(importModule, importDefaultTypeName) =
+      getImportModuleAndDefaultType(ctx, protocol);
+
+  emitDiagnostic(expr->getLoc(), diag::object_literal_default_type_missing,
+                 plainName);
+  if (!importModule.empty()) {
+    emitDiagnostic(expr->getLoc(), diag::object_literal_resolve_import,
+                   importModule, importDefaultTypeName, plainName);
+  }
+
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -1926,6 +1926,15 @@ class UnableToInferClosureReturnType final : public FailureDiagnostic {
   bool diagnoseAsError();
 };
 
+class UnableToInferProtocolLiteralType final : public FailureDiagnostic {
+public:
+  UnableToInferProtocolLiteralType(ConstraintSystem &cs,
+                                   ConstraintLocator *locator)
+      : FailureDiagnostic(cs, locator) {}
+
+  bool diagnoseAsError();
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -1168,3 +1168,28 @@ SpecifyClosureReturnType::create(ConstraintSystem &cs,
                                  ConstraintLocator *locator) {
   return new (cs.getAllocator()) SpecifyClosureReturnType(cs, locator);
 }
+
+bool SpecifyObjectLiteralTypeImport::diagnose(bool asNote) const {
+  auto &cs = getConstraintSystem();
+  UnableToInferProtocolLiteralType failure(cs, getLocator());
+  return failure.diagnose(asNote);
+}
+
+SpecifyObjectLiteralTypeImport *
+SpecifyObjectLiteralTypeImport::attempt(ConstraintSystem &cs,
+                                        ConstraintLocator *locator) {
+  // FIXME: ArgumentMismatchFailure is currently used from CSDiag, meaning
+  // we could endup emiting this diagnostic duplicated.
+  if (cs.Options.contains(ConstraintSystemFlags::SubExpressionDiagnostics))
+    return nullptr;
+  
+  auto *anchor = locator->getAnchor();
+  if (!isa<ObjectLiteralExpr>(anchor))
+    return nullptr;
+
+  // * The object literal has no contextual type
+  if (cs.getContextualType())
+    return nullptr;
+  
+  return new (cs.getAllocator()) SpecifyObjectLiteralTypeImport(cs, locator);
+}

lib/Sema/CSFix.h

@@ -235,6 +235,11 @@ enum class FixKind : uint8_t {
   /// Closure return type has to be explicitly specified because it can't be
   /// inferred in current context e.g. because it's a multi-statement closure.
   SpecifyClosureReturnType,
+  
+  /// Object literal type coudn't be infered because the module where it's
+  /// protocol type was not imported.
+  SpecifyObjectLiteralTypeImport,
+
 };
 
 class ConstraintFix {
@@ -1633,6 +1638,22 @@ class SpecifyClosureReturnType final : public ConstraintFix {
                                           ConstraintLocator *locator);
 };
 
+class SpecifyObjectLiteralTypeImport final : public ConstraintFix {
+  SpecifyObjectLiteralTypeImport(ConstraintSystem &cs, ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::SpecifyObjectLiteralTypeImport, locator) {}
+
+public:
+  std::string getName() const {
+    return "specify object literal protocol type import";
+  }
+
+  bool diagnose(bool asNote = false) const;
+
+  static SpecifyObjectLiteralTypeImport *attempt(ConstraintSystem &cs,
+                                                 ConstraintLocator *locator);
+
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSSimplify.cpp

@@ -1027,7 +1027,18 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
             argsWithLabels, params, paramInfo, argInfo->HasTrailingClosure,
             cs.shouldAttemptFixes(), listener, parameterBindings))
       return cs.getTypeMatchFailure(locator);
-
+    
+    if (!callee) {
+      // If we couldn't find a callee, diagnose and maybe suggest
+      // import a module.
+      if (cs.shouldAttemptFixes()) {
+        if (auto *fix = SpecifyObjectLiteralTypeImport::attempt(cs, loc)) {
+          if (cs.recordFix(fix))
+            return cs.getTypeMatchFailure(locator);
+        }
+      }
+    }
+    
     auto extraArguments = listener.getExtraneousArguments();
     if (!extraArguments.empty()) {
       if (RemoveExtraneousArguments::isMinMaxNameShadowing(cs, locator))
@@ -1047,6 +1058,7 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
           return cs.getTypeMatchFailure(locator);
       }
     }
+    
   }
 
   // If this application is part of an operator, then we allow an implicit
@@ -8739,7 +8751,8 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::UseValueTypeOfRawRepresentative:
   case FixKind::ExplicitlyConstructRawRepresentable:
   case FixKind::SpecifyBaseTypeForContextualMember:
-  case FixKind::CoerceToCheckedCast: {
+  case FixKind::CoerceToCheckedCast:
+  case FixKind::SpecifyObjectLiteralTypeImport: {
     return recordFix(fix) ? SolutionKind::Error : SolutionKind::Solved;
   }