From 11404f81dba96c449a5d8b05d1844035f849b9a8 Mon Sep 17 00:00:00 2001
From: Alex Lorenz <arphaman@gmail.com>
Date: Wed, 30 Aug 2017 13:24:37 +0000
Subject: [PATCH] [refactor] Examine the whole range for ObjC @implementation
 decls when computing the AST selection

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@312121 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Tooling/Refactoring/ASTSelection.cpp | 18 ++++++++++++++++--
 unittests/Tooling/ASTSelectionTest.cpp   | 19 +++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/lib/Tooling/Refactoring/ASTSelection.cpp b/lib/Tooling/Refactoring/ASTSelection.cpp
index 0bf5d95624..4fa0086697 100644
--- a/lib/Tooling/Refactoring/ASTSelection.cpp
+++ b/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/unittests/Tooling/ASTSelectionTest.cpp b/unittests/Tooling/ASTSelectionTest.cpp
index 95326dded8..bde052e849 100644
--- a/unittests/Tooling/ASTSelectionTest.cpp
+++ b/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
-- 
2.40.0