[Diagnostics] Diagnose subscript operator misuse via fixes

Fix to use subscript operator instead of spelled out name helps
to produce a solution, that makes it much easier to diagnose
problems precisely and provide proper fix-its, it also helps to
diagnose ambiguous cases, and stacks up nicely with other errors.

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -64,6 +64,9 @@ ERROR(ambiguous_reference_to_decl,none,
 ERROR(ambiguous_subscript,none,
       "ambiguous subscript with base type %0 and index type %1",
       (Type, Type))
+NOTE(ambiguous_subscript_candidate_requires_operator,none,
+     "candidate requires use of the subscript operator", ())
+
 ERROR(could_not_find_value_subscript,none,
       "value of type %0 has no subscripts",
       (Type))

lib/Sema/CSDiag.cpp

@@ -572,10 +572,6 @@ class FailureDiagnosis :public ASTVisitor<FailureDiagnosis, /*exprresult*/bool>{
 
   bool diagnoseSubscriptErrors(SubscriptExpr *SE, bool performingSet);
 
-  /// Diagnose the usage of 'subscript' instead of the operator when calling
-  /// a subscript and offer a fixit if the inputs are compatible.
-  bool diagnoseSubscriptMisuse(ApplyExpr *callExpr);
-
   /// Try to diagnose common errors involving implicitly non-escaping parameters
   /// of function type, giving more specific and simpler diagnostics, attaching
   /// notes on the parameter, and offering fixits to insert @escaping. Returns
@@ -5089,105 +5085,6 @@ bool FailureDiagnosis::diagnoseCallContextualConversionErrors(
                                            exprType);
 }
 
-bool FailureDiagnosis::diagnoseSubscriptMisuse(ApplyExpr *callExpr) {
-  auto UDE = dyn_cast<UnresolvedDotExpr>(callExpr->getFn());
-  if (!UDE)
-    return false;
-
-  auto baseExpr = UDE->getBase();
-  if (!baseExpr || UDE->getName().getBaseName() != getTokenText(tok::kw_subscript))
-    return false;
-
-  auto baseType = CS.getType(baseExpr);
-  // Look up subscript declarations.
-  auto lookup = CS.lookupMember(baseType->getRValueType(),
-                                DeclName(DeclBaseName::createSubscript()));
-  auto nonSubscrLookup = CS.lookupMember(baseType->getRValueType(),
-                                         UDE->getName());
-  // Make sure we only found subscript declarations. If not, the problem
-  // is different - return.
-  if (lookup.empty() || !nonSubscrLookup.empty())
-    return false;
-  // Try to resolve a type of the argument expression.
-  auto argExpr = typeCheckChildIndependently(callExpr->getArg(),
-                                             Type(), CTP_CallArgument);
-  if (!argExpr)
-    return CS.TC.Diags.hadAnyError();
-
-  SmallVector<Identifier, 2> scratch;
-  ArrayRef<Identifier> argLabels = callExpr->getArgumentLabels(scratch);
-  SmallVector<OverloadChoice, 2> choices;
-
-  for (auto candidate : lookup)
-    choices.push_back(OverloadChoice(baseType, candidate.getValueDecl(),
-                                     UDE->getFunctionRefKind()));
-  CalleeCandidateInfo candidateInfo(baseType, choices,
-                                    callArgHasTrailingClosure(argExpr),
-                                    CS, true);
-
-  auto params = decomposeArgType(CS.getType(argExpr), argLabels);
-  using ClosenessPair = CalleeCandidateInfo::ClosenessResultTy;
-
-  candidateInfo.filterList([&](OverloadCandidate cand) -> ClosenessPair {
-    auto candFuncType = cand.getFunctionType();
-    if (!candFuncType)
-      return {CC_GeneralMismatch, {}};
-
-    auto candParams = candFuncType->getParams();
-    if (params.size() != candParams.size())
-      return {CC_GeneralMismatch, {}};
-
-    for (unsigned i = 0, e = params.size(); i < e; i ++) {
-      if (CS.TC.isConvertibleTo(params[i].getOldType(),
-                                candParams[i].getOldType(), CS.DC) ||
-          candParams[i].getOldType()->is<GenericTypeParamType>())
-        continue;
-      return {CC_GeneralMismatch, {}};
-    }
-    return {CC_ExactMatch, {}};
-  });
-
-  auto *locator = CS.getConstraintLocator(UDE, ConstraintLocator::Member);
-  auto memberRange = baseExpr->getSourceRange();
-  if (locator)
-    locator = simplifyLocator(CS, locator, memberRange);
-  auto nameLoc = DeclNameLoc(memberRange.Start);
-
-  auto diag = diagnose(baseExpr->getLoc(),
-                       diag::could_not_find_subscript_member_did_you_mean,
-                       baseType);
-  diag.highlight(memberRange).highlight(nameLoc.getSourceRange());
-
-  auto showNote = [&]() {
-    diag.flush();
-    if (candidateInfo.size() == 1)
-      diagnose(candidateInfo.candidates.front().getDecl(),
-               diag::kind_declared_here, DescriptiveDeclKind::Subscript);
-  };
-  if (candidateInfo.closeness != CC_ExactMatch) {
-    showNote();
-    return true;
-  }
-  auto toCharSourceRange = Lexer::getCharSourceRangeFromSourceRange;
-  auto lastArgSymbol = toCharSourceRange(CS.TC.Context.SourceMgr,
-                                         argExpr->getEndLoc());
-
-  diag.fixItReplace(SourceRange(argExpr->getStartLoc()),
-                    getTokenText(tok::l_square));
-  diag.fixItRemove(nameLoc.getSourceRange());
-  diag.fixItRemove(SourceRange(UDE->getDotLoc()));
-  if (CS.TC.Context.SourceMgr.extractText(lastArgSymbol) ==
-      getTokenText(tok::r_paren))
-    diag.fixItReplace(SourceRange(argExpr->getEndLoc()),
-                      getTokenText(tok::r_square));
-  else
-    diag.fixItInsertAfter(argExpr->getEndLoc(),
-                          getTokenText(tok::r_square));
-  showNote();
-
-  return true;
-}
-
 // Check if there is a structural problem in the function expression
 // by performing type checking with the option to allow unresolved
 // type variables. If that is going to produce a function type with
@@ -5262,12 +5159,6 @@ bool FailureDiagnosis::visitApplyExpr(ApplyExpr *callExpr) {
   auto originalFnType = CS.getType(callExpr->getFn());
 
   if (shouldTypeCheckFunctionExpr(*this, CS.DC, fnExpr)) {
-
-    // If we are misusing a subscript, diagnose that and provide a fixit.
-    // We diagnose this here to have enough context to offer an appropriate fixit.
-    if (diagnoseSubscriptMisuse(callExpr)) {
-      return CS.TC.Diags.hadAnyError();
-    }
     // Type check the function subexpression to resolve a type for it if
     // possible.
     fnExpr = typeCheckChildIndependently(callExpr->getFn());

lib/Sema/CSDiagnostics.cpp

@@ -78,6 +78,10 @@ FailureDiagnostic::emitDiagnostic(ArgTypes &&... Args) const {
   return cs.TC.diagnose(std::forward<ArgTypes>(Args)...);
 }
 
+Expr *FailureDiagnostic::findParentExpr(Expr *subExpr) const {
+  return E ? E->getParentMap()[subExpr] : nullptr;
+}
+
 Type RequirementFailure::getOwnerType() const {
   return getType(getAnchor())->getInOutObjectType()->getMetatypeInstanceType();
 }
@@ -1275,3 +1279,57 @@ bool MissingCallFailure::diagnoseAsError() {
       .fixItInsertAfter(insertLoc, "()");
   return true;
 }
+
+bool SubscriptMisuseFailure::diagnoseAsError() {
+  auto &sourceMgr = getASTContext().SourceMgr;
+
+  auto *memberExpr = cast<UnresolvedDotExpr>(getRawAnchor());
+  auto *baseExpr = getAnchor();
+
+  auto memberRange = baseExpr->getSourceRange();
+  (void)simplifyLocator(getConstraintSystem(), getLocator(), memberRange);
+
+  auto nameLoc = DeclNameLoc(memberRange.Start);
+
+  auto diag = emitDiagnostic(baseExpr->getLoc(),
+                             diag::could_not_find_subscript_member_did_you_mean,
+                             getType(baseExpr));
+
+  diag.highlight(memberRange).highlight(nameLoc.getSourceRange());
+
+  auto *parentExpr = findParentExpr(memberExpr);
+  assert(parentExpr && "Couldn't find a parent expression for a member call?!");
+
+  auto *argExpr = cast<ApplyExpr>(parentExpr)->getArg();
+
+  auto toCharSourceRange = Lexer::getCharSourceRangeFromSourceRange;
+  auto lastArgSymbol = toCharSourceRange(sourceMgr, argExpr->getEndLoc());
+
+  diag.fixItReplace(SourceRange(argExpr->getStartLoc()),
+                    getTokenText(tok::l_square));
+  diag.fixItRemove(nameLoc.getSourceRange());
+  diag.fixItRemove(SourceRange(memberExpr->getDotLoc()));
+
+  if (sourceMgr.extractText(lastArgSymbol) == getTokenText(tok::r_paren))
+    diag.fixItReplace(SourceRange(argExpr->getEndLoc()),
+                      getTokenText(tok::r_square));
+  else
+    diag.fixItInsertAfter(argExpr->getEndLoc(), getTokenText(tok::r_square));
+
+  diag.flush();
+  if (auto overload = getOverloadChoiceIfAvailable(getLocator())) {
+    emitDiagnostic(overload->choice.getDecl(), diag::kind_declared_here,
+                   DescriptiveDeclKind::Subscript);
+  }
+
+  return true;
+}
+
+bool SubscriptMisuseFailure::diagnoseAsNote() {
+  if (auto overload = getOverloadChoiceIfAvailable(getLocator())) {
+    emitDiagnostic(overload->choice.getDecl(),
+                   diag::ambiguous_subscript_candidate_requires_operator);
+    return true;
+  }
+  return false;
+}

lib/Sema/CSDiagnostics.h

@@ -143,6 +143,10 @@ class FailureDiagnostic {
   /// \returns true is locator hasn't been simplified down to expression.
   bool hasComplexLocator() const { return HasComplexLocator; }
 
+  /// \returns A parent expression if sub-expression is contained anywhere
+  /// in the root expression or `nullptr` otherwise.
+  Expr *findParentExpr(Expr *subExpr) const;
+
 private:
   /// Compute anchor expression associated with current diagnostic.
   std::pair<Expr *, bool> computeAnchor() const;
@@ -182,15 +186,8 @@ class RequirementFailure : public FailureDiagnostic {
     if (!expr)
       return;
 
-    auto *anchor = getAnchor();
-    expr->forEachChildExpr([&](Expr *subExpr) -> Expr * {
-      auto *AE = dyn_cast<ApplyExpr>(subExpr);
-      if (!AE || AE->getFn() != anchor)
-        return subExpr;
-
-      Apply = AE;
-      return nullptr;
-    });
+    if (auto *parentExpr = findParentExpr(getAnchor()))
+      Apply = dyn_cast<ApplyExpr>(parentExpr);
   }
 
   unsigned getRequirementIndex() const {
@@ -639,6 +636,16 @@ class MissingCallFailure final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+class SubscriptMisuseFailure final : public FailureDiagnostic {
+public:
+  SubscriptMisuseFailure(Expr *root, ConstraintSystem &cs,
+                         ConstraintLocator *locator)
+      : FailureDiagnostic(root, cs, locator) {}
+
+  bool diagnoseAsError() override;
+  bool diagnoseAsNote() override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -237,7 +237,8 @@ InsertExplicitCall *InsertExplicitCall::create(ConstraintSystem &cs,
 }
 
 bool UseSubscriptOperator::diagnose(Expr *root, bool asNote) const {
-  return false;
+  auto failure = SubscriptMisuseFailure(root, getConstraintSystem(), getLocator());
+  return failure.diagnose(asNote);
 }
 
 UseSubscriptOperator *UseSubscriptOperator::create(ConstraintSystem &cs,

test/decl/subscript/subscripting.swift

@@ -265,8 +265,12 @@ func testGenSubscriptFixit(_ s0: GenSubscriptFixitTest) {
 }
 
 struct SubscriptTest1 {
-  subscript(keyword:String) -> Bool { return true }  // expected-note 2 {{found this candidate}}
-  subscript(keyword:String) -> String? {return nil }  // expected-note 2 {{found this candidate}}
+  subscript(keyword:String) -> Bool { return true }
+  // expected-note@-1 2 {{found this candidate}}
+  // expected-note@-2 2 {{candidate requires use of the subscript operator}}
+  subscript(keyword:String) -> String? {return nil }
+  // expected-note@-1 2 {{found this candidate}}
+  // expected-note@-2 2 {{candidate requires use of the subscript operator}}
 
   subscript(arg: SubClass) -> Bool { return true } // expected-note {{declared here}}
   subscript(arg: Protocol) -> Bool { return true } // expected-note 2 {{declared here}}
@@ -285,21 +289,23 @@ func testSubscript1(_ s1 : SubscriptTest1) {
   _ = s1.subscript((true, (5, baz: SubSubClass())), "hello")
   // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}} {{9-10=}} {{10-19=}} {{19-20=[}} {{60-61=]}}
   _ = s1.subscript((fo: true, (5, baz: SubClass())), "hello")
-  // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}}
+  // expected-error@-1 {{cannot convert value of type '(fo: Bool, (Int, baz: SubClass))' to expected argument type '(foo: Bool, bar: (Int, baz: SubClass))'}}
   _ = s1.subscript(SubSubClass())
   // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}} {{9-10=}} {{10-19=}} {{19-20=[}} {{33-34=]}}
   _ = s1.subscript(ClassConformingToProtocol())
   // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}} {{9-10=}} {{10-19=}} {{19-20=[}} {{47-48=]}}
   _ = s1.subscript(ClassConformingToRefinedProtocol())
   // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}} {{9-10=}} {{10-19=}} {{19-20=[}} {{54-55=]}}
   _ = s1.subscript(true)
-  // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}}
+  // expected-error@-1 {{cannot invoke 'subscript' with an argument list of type '(Bool)'}}
+  // expected-note@-2 {{overloads for 'subscript' exist with these partially matching parameter lists: (Protocol), (String), (SubClass)}}
   _ = s1.subscript(SuperClass())
-  // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}}
+  // expected-error@-1 {{cannot invoke 'subscript' with an argument list of type '(SuperClass)'}}
+  // expected-note@-2 {{overloads for 'subscript' exist with these partially matching parameter lists: (Protocol), (String), (SubClass)}}
   _ = s1.subscript("hello")
-  // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}} {{9-10=}} {{10-19=}} {{19-20=[}} {{27-28=]}}
+  // expected-error@-1 {{ambiguous reference to subscript 'subscript(_:)'}}
   _ = s1.subscript("hello"
-  // expected-error@-1 {{value of type 'SubscriptTest1' has no property or method named 'subscript'; did you mean to use the subscript operator?}} {{9-10=}} {{10-19=}} {{19-20=[}} {{27-27=]}}
+  // expected-error@-1 {{ambiguous reference to subscript 'subscript(_:)'}}
   // expected-note@-2 {{to match this opening '('}}
 
   let _ = s1["hello"]