Merge pull request #68709 from hamishknight/loc-picking-lawyer

Author: hamishknight

lib/IDE/CursorInfo.cpp

@@ -182,22 +182,22 @@ class NodeFinder : ASTWalker {
       DeclContextStack.push_back(newDC);
     }
 
-    if (D->getLoc() != LocToResolve) {
+    auto *VD = dyn_cast<ValueDecl>(D);
+    if (!VD)
       return Action::Continue();
-    }
 
-    if (auto VD = dyn_cast<ValueDecl>(D)) {
-      // FIXME: ParamDecls might be closure parameters that can have ambiguous
-      // types. The current infrastructure of just asking for the VD's type
-      // doesn't work here. We need to inspect the constraints system solution.
-      if (VD->hasName() && !isa<ParamDecl>(D)) {
-        assert(Result == nullptr);
-        Result = std::make_unique<NodeFinderDeclResult>(VD);
-        return Action::Stop();
-      }
-    }
+    // FIXME: ParamDecls might be closure parameters that can have ambiguous
+    // types. The current infrastructure of just asking for the VD's type
+    // doesn't work here. We need to inspect the constraints system solution.
+    if (isa<ParamDecl>(VD))
+      return Action::Continue();
 
-    return Action::Continue();
+    if (!VD->hasName() || VD->getNameLoc() != LocToResolve)
+      return Action::Continue();
+
+    assert(Result == nullptr);
+    Result = std::make_unique<NodeFinderDeclResult>(VD);
+    return Action::Stop();
   }
 
   PostWalkAction walkToDeclPost(Decl *D) override {
@@ -208,6 +208,23 @@ class NodeFinder : ASTWalker {
     return Action::Continue();
   }
 
+  /// Retrieve the name location for an expression that supports cursor info.
+  DeclNameLoc getExprNameLoc(Expr *E) {
+    if (auto *DRE = dyn_cast<DeclRefExpr>(E))
+      return DRE->getNameLoc();
+    
+    if (auto *UDRE = dyn_cast<UnresolvedDeclRefExpr>(E))
+      return UDRE->getNameLoc();
+
+    if (auto *ODRE = dyn_cast<OverloadedDeclRefExpr>(E))
+      return ODRE->getNameLoc();
+
+    if (auto *UDE = dyn_cast<UnresolvedDotExpr>(E))
+      return UDE->getNameLoc();
+
+    return DeclNameLoc();
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     if (auto closure = dyn_cast<ClosureExpr>(E)) {
       DeclContextStack.push_back(closure);
@@ -224,25 +241,13 @@ class NodeFinder : ASTWalker {
       }
     }
 
-    if (E->getLoc() != LocToResolve) {
+    if (getExprNameLoc(E).getBaseNameLoc() != LocToResolve)
       return Action::Continue(E);
-    }
 
-    switch (E->getKind()) {
-    case ExprKind::DeclRef:
-    case ExprKind::UnresolvedDot:
-    case ExprKind::UnresolvedDeclRef:
-    case ExprKind::OverloadedDeclRef: {
-      assert(Result == nullptr);
-      Result =
-          std::make_unique<NodeFinderExprResult>(E, getCurrentDeclContext());
-      return Action::Stop();
-    }
-    default:
-      break;
-    }
-
-    return Action::Continue(E);
+    assert(Result == nullptr);
+    Result =
+        std::make_unique<NodeFinderExprResult>(E, getCurrentDeclContext());
+    return Action::Stop();
   }
 
   PostWalkResult<Expr *> walkToExprPost(Expr *E) override {