[refactor] Examine the whole range for ObjC @implementation decls

when computing the AST selection

llvm-svn: 312121

2 files changed, 35 insertions, 2 deletions

diff --git a/clang/lib/Tooling/Refactoring/ASTSelection.cpp b/clang/lib/Tooling/Refactoring/ASTSelection.cpp
index 0bf5d95624e..4fa00866978 100644
--- a/clang/lib/Tooling/Refactoring/ASTSelection.cpp
+++ b/clang/lib/Tooling/Refactoring/ASTSelection.cpp
@@ -17,6 +17,21 @@ using ast_type_traits::DynTypedNode;
 
 namespace {
 
+CharSourceRange getLexicalDeclRange(Decl *D, const SourceManager &SM,
+                                    const LangOptions &LangOpts) {
+  if (!isa<ObjCImplDecl>(D))
+    return CharSourceRange::getTokenRange(D->getSourceRange());
+  // Objective-C implementation declarations end at the '@' instead of the 'end'
+  // keyword. Use the lexer to find the location right after 'end'.
+  SourceRange R = D->getSourceRange();
+  SourceLocation LocAfterEnd = Lexer::findLocationAfterToken(
+      R.getEnd(), tok::raw_identifier, SM, LangOpts,
+      /*SkipTrailingWhitespaceAndNewLine=*/false);
+  return LocAfterEnd.isValid()
+             ? CharSourceRange::getCharRange(R.getBegin(), LocAfterEnd)
+             : CharSourceRange::getTokenRange(R);
+}
+
 /// Constructs the tree of selected AST nodes that either contain the location
 /// of the cursor or overlap with the selection range.
 class ASTSelectionFinder
@@ -62,9 +77,8 @@ public:
     if (SM.getFileID(FileLoc) != TargetFile)
       return true;
 
-    // FIXME (Alex Lorenz): Add location adjustment for ObjCImplDecls.
     SourceSelectionKind SelectionKind =
-        selectionKindFor(CharSourceRange::getTokenRange(D->getSourceRange()));
+        selectionKindFor(getLexicalDeclRange(D, SM, Context.getLangOpts()));
     SelectionStack.push_back(
         SelectedASTNode(DynTypedNode::create(*D), SelectionKind));
     LexicallyOrderedRecursiveASTVisitor::TraverseDecl(D);
diff --git a/clang/unittests/Tooling/ASTSelectionTest.cpp b/clang/unittests/Tooling/ASTSelectionTest.cpp
index 95326dded88..bde052e849c 100644
--- a/clang/unittests/Tooling/ASTSelectionTest.cpp
+++ b/clang/unittests/Tooling/ASTSelectionTest.cpp
@@ -494,4 +494,23 @@ void foo() {
       });
 }
 
+TEST(ASTSelectionFinder, CorrectEndForObjectiveCImplementation) {
+  StringRef Source = R"(
+@interface I
+@end
+@implementation I
+@ end
+)";
+  // Just after '@ end'
+  findSelectedASTNodes(Source, {5, 6}, None,
+                       [](Optional<SelectedASTNode> Node) {
+                         EXPECT_TRUE(Node);
+                         EXPECT_EQ(Node->Children.size(), 1u);
+                         checkNode<ObjCImplementationDecl>(
+                             Node->Children[0],
+                             SourceSelectionKind::ContainsSelection);
+                       },
+                       SelectionFinderVisitor::Lang_OBJC);
+}
+
 } // end anonymous namespace