[CS] Fix locator simplification with 'Member' path element

Previously, a `Member` path element in a `ConstraintLocator` was simplified to the base on which the member was accessed. This is incorrect.

Author: ahoppen

include/swift/Sema/ConstraintLocator.h

@@ -313,6 +313,11 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// branch, and if so, the kind of branch.
   Optional<SingleValueStmtBranchKind> isForSingleValueStmtBranch() const;
 
+  /// Returns true if \p locator is ending with either of the following
+  ///  - Member
+  ///  - Member -> KeyPathDynamicMember
+  bool isMemberRef() const;
+
   /// Determine whether this locator points directly to a given expression.
   template <typename E> bool directlyAt() const {
     if (auto *expr = getAnchor().dyn_cast<Expr *>())

lib/Sema/CSDiagnostics.cpp

@@ -77,13 +77,6 @@ ASTNode FailureDiagnostic::getAnchor() const {
       return locator->getAnchor();
   }
 
-  // FIXME: Work around an odd locator representation that doesn't separate the
-  // base of a subscript member from the member access.
-  if (locator->isLastElement<LocatorPathElt::SubscriptMember>()) {
-    if (auto subscript = getAsExpr<SubscriptExpr>(anchor))
-      anchor = subscript->getBase();
-  }
-
   return anchor;
 }
 
@@ -1345,6 +1338,15 @@ bool MissingExplicitConversionFailure::diagnoseAsError() {
   return true;
 }
 
+ASTNode MemberReferenceFailure::getAnchor() const {
+  auto anchor = FailureDiagnostic::getAnchor();
+  if (auto base = getBaseExprFor(getAsExpr(anchor))) {
+    return base;
+  } else {
+    return anchor;
+  }
+}
+
 SourceRange MemberAccessOnOptionalBaseFailure::getSourceRange() const {
   if (auto componentPathElt =
           getLocator()->getLastElementAs<LocatorPathElt::KeyPathComponent>()) {
@@ -3680,6 +3682,15 @@ bool MissingCallFailure::diagnoseAsError() {
   return true;
 }
 
+ASTNode PropertyWrapperReferenceFailure::getAnchor() const {
+  auto anchor = FailureDiagnostic::getAnchor();
+  if (getReferencedMember()) {
+    return getBaseExprFor(getAsExpr(anchor));
+  } else {
+    return anchor;
+  }
+}
+
 bool ExtraneousPropertyWrapperUnwrapFailure::diagnoseAsError() {
   auto newPrefix = usingProjection() ? "$" : "_";
 
@@ -3751,21 +3762,12 @@ bool SubscriptMisuseFailure::diagnoseAsError() {
   auto &sourceMgr = getASTContext().SourceMgr;
 
   auto *memberExpr = castToExpr<UnresolvedDotExpr>(getRawAnchor());
-
-  auto memberRange = getSourceRange();
-
-  {
-    auto rawAnchor = getRawAnchor();
-    auto path = locator->getPath();
-    simplifyLocator(rawAnchor, path, memberRange);
-  }
-
-  auto nameLoc = DeclNameLoc(memberRange.Start);
+  auto *base = memberExpr->getBase();
 
   auto diag = emitDiagnostic(diag::could_not_find_subscript_member_did_you_mean,
-                             getType(getAnchor()));
+                             getType(base));
 
-  diag.highlight(memberRange).highlight(nameLoc.getSourceRange());
+  diag.highlight(memberExpr->getNameLoc().getSourceRange());
 
   if (auto *parentExpr = dyn_cast_or_null<ApplyExpr>(findParentExpr(memberExpr))) {
     auto *args = parentExpr->getArgs();
@@ -3775,7 +3777,7 @@ bool SubscriptMisuseFailure::diagnoseAsError() {
 
     diag.fixItReplace(SourceRange(args->getStartLoc()),
                       getTokenText(tok::l_square));
-    diag.fixItRemove(nameLoc.getSourceRange());
+    diag.fixItRemove(memberExpr->getNameLoc().getSourceRange());
     diag.fixItRemove(SourceRange(memberExpr->getDotLoc()));
 
     if (sourceMgr.extractText(lastArgSymbol) == getTokenText(tok::r_paren))

lib/Sema/CSDiagnostics.h

@@ -534,20 +534,29 @@ class LabelingFailure final : public FailureDiagnostic {
   bool diagnoseAsNote() override;
 };
 
+/// A diagnostic that will be emitted on the base if its locator points to a
+/// member access.
+class MemberReferenceFailure : public FailureDiagnostic {
+public:
+  MemberReferenceFailure(const Solution &solution, ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator) {}
+
+  ASTNode getAnchor() const override;
+};
+
 /// Diagnose failures related to attempting member access on optional base
 /// type without optional chaining or force-unwrapping it first.
-class MemberAccessOnOptionalBaseFailure final : public FailureDiagnostic {
+class MemberAccessOnOptionalBaseFailure final : public MemberReferenceFailure {
   DeclNameRef Member;
   Type MemberBaseType;
   bool ResultTypeIsOptional;
 
 public:
   MemberAccessOnOptionalBaseFailure(const Solution &solution,
                                     ConstraintLocator *locator,
-                                    DeclNameRef memberName,
-                                    Type memberBaseType,
+                                    DeclNameRef memberName, Type memberBaseType,
                                     bool resultOptional)
-      : FailureDiagnostic(solution, locator), Member(memberName),
+      : MemberReferenceFailure(solution, locator), Member(memberName),
         MemberBaseType(resolveType(memberBaseType)),
         ResultTypeIsOptional(resultOptional) {}
 
@@ -1127,6 +1136,8 @@ class PropertyWrapperReferenceFailure : public ContextualFailure {
       : ContextualFailure(solution, base, wrapper, locator), Property(property),
         UsingProjection(usingProjection) {}
 
+  ASTNode getAnchor() const override;
+
   VarDecl *getProperty() const { return Property; }
 
   Identifier getPropertyName() const { return Property->getName(); }
@@ -1201,14 +1212,14 @@ class SubscriptMisuseFailure final : public FailureDiagnostic {
   bool diagnoseAsNote() override;
 };
 
-class InvalidMemberRefFailure : public FailureDiagnostic {
+class InvalidMemberRefFailure : public MemberReferenceFailure {
   Type BaseType;
   DeclNameRef Name;
 
 public:
   InvalidMemberRefFailure(const Solution &solution, Type baseType,
                           DeclNameRef memberName, ConstraintLocator *locator)
-      : FailureDiagnostic(solution, locator),
+      : MemberReferenceFailure(solution, locator),
         BaseType(baseType->getRValueType()), Name(memberName) {}
 
 protected:
@@ -1334,7 +1345,7 @@ class InvalidMemberRefOnExistential final : public InvalidMemberRefFailure {
 ///
 /// }
 /// ```
-class AllowTypeOrInstanceMemberFailure final : public FailureDiagnostic {
+class AllowTypeOrInstanceMemberFailure final : public MemberReferenceFailure {
   Type BaseType;
   ValueDecl *Member;
   DeclNameRef Name;
@@ -1343,7 +1354,7 @@ class AllowTypeOrInstanceMemberFailure final : public FailureDiagnostic {
   AllowTypeOrInstanceMemberFailure(const Solution &solution, Type baseType,
                                    ValueDecl *member, DeclNameRef name,
                                    ConstraintLocator *locator)
-      : FailureDiagnostic(solution, locator),
+      : MemberReferenceFailure(solution, locator),
         BaseType(baseType->getRValueType()), Member(member), Name(name) {
     assert(member);
   }

lib/Sema/CSSimplify.cpp

@@ -10538,7 +10538,9 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
         // let's return, otherwise let's fall-through and report
         // this problem as a missing member.
         if (result == SolutionKind::Solved)
-          return recordFix(InsertExplicitCall::create(*this, locator))
+          return recordFix(InsertExplicitCall::create(
+                     *this, getConstraintLocator(
+                                locator, ConstraintLocator::MemberRefBase)))
                      ? SolutionKind::Error
                      : SolutionKind::Solved;
       }

lib/Sema/ConstraintLocator.cpp

@@ -662,6 +662,19 @@ ConstraintLocator::isForSingleValueStmtBranch() const {
   return SingleValueStmtBranchKind::Regular;
 }
 
+bool ConstraintLocator::isMemberRef() const {
+  if (isLastElement<LocatorPathElt::Member>()) {
+    return true;
+  } else if (isLastElement<LocatorPathElt::KeyPathDynamicMember>()) {
+    auto path = getPath();
+    if (path.size() >= 2 &&
+        path[path.size() - 2].is<LocatorPathElt::Member>()) {
+      return true;
+    }
+  }
+  return false;
+}
+
 GenericTypeParamType *ConstraintLocator::getGenericParameter() const {
   // Check whether we have a path that terminates at a generic parameter.
   return isForGenericParameter() ?

lib/Sema/ConstraintSystem.cpp

@@ -5469,6 +5469,15 @@ void constraints::simplifyLocator(ASTNode &anchor,
       LLVM_FALLTHROUGH;
 
     case ConstraintLocator::Member:
+      if (auto UDE = getAsExpr<UnresolvedDotExpr>(anchor)) {
+        path = path.slice(1);
+        continue;
+      }
+      if (anchor.is<Pattern *>()) {
+        path = path.slice(1);
+        continue;
+      }
+      break;
     case ConstraintLocator::MemberRefBase:
       if (auto UDE = getAsExpr<UnresolvedDotExpr>(anchor)) {
         range = UDE->getNameLoc().getSourceRange();

test/Constraints/rdar46544601.swift

@@ -22,9 +22,9 @@ extension P {
 func crash(_ p: P, payload: [UInt8]) throws {
   p.foo(data: payload).then { _ in
     return Future<(D, [D])>()
-  }.then { (id, arr) in // expected-error {{generic parameter 'U' could not be inferred}}
+  }.then { (id, arr) in
     p.foo(arr: arr, data: []).and(result: (id, arr))
-  }.then { args0 in
+  }.then { args0 in // expected-error {{generic parameter 'U' could not be inferred}}
     let (parentID, args1) = args0
     p.bar(root: parentID, from: p).and(result: args1)
   }.whenFailure { _ in }

test/Misc/serialized-diagnostics-prettyprint.swift

@@ -6,7 +6,7 @@
 // RUN: %FileCheck --input-file=%t.deserialized_diagnostics.txt %s
 
 var x = String.init // expected-error{{ambiguous use of 'init'}}
-// CHECK: {{.*[/\\]}}serialized-diagnostics-prettyprint.swift:[[@LINE-1]]:9: error: ambiguous use of 'init'
+// CHECK: {{.*[/\\]}}serialized-diagnostics-prettyprint.swift:[[@LINE-1]]:16: error: ambiguous use of 'init'
 
 // CHECK: Swift.String:2:23: note: found this candidate
 // CHECK: CONTENTS OF FILE Swift.String: