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

[Alex Lorenz via cfe-commits]

Author: arphaman
Date: Wed Aug 30 06:24:37 2017
New Revision: 312121

URL: http://llvm.org/viewvc/llvm-project?rev=312121&view=rev
Log:
[refactor] Examine the whole range for ObjC @implementation decls
when computing the AST selection

Modified:
    cfe/trunk/lib/Tooling/Refactoring/ASTSelection.cpp
    cfe/trunk/unittests/Tooling/ASTSelectionTest.cpp

Modified: cfe/trunk/lib/Tooling/Refactoring/ASTSelection.cpp
URL: http://llvm.org/viewvc/llvm-project/cfe/trunk/lib/Tooling/Refactoring/ASTSelection.cpp?rev=312121&r1=312120&r2=312121&view=diff
==============================================================================
--- cfe/trunk/lib/Tooling/Refactoring/ASTSelection.cpp (original)
+++ cfe/trunk/lib/Tooling/Refactoring/ASTSelection.cpp Wed Aug 30 06:24:37 2017
@@ -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);

Modified: cfe/trunk/unittests/Tooling/ASTSelectionTest.cpp
URL: http://llvm.org/viewvc/llvm-project/cfe/trunk/unittests/Tooling/ASTSelectionTest.cpp?rev=312121&r1=312120&r2=312121&view=diff
==============================================================================
--- cfe/trunk/unittests/Tooling/ASTSelectionTest.cpp (original)
+++ cfe/trunk/unittests/Tooling/ASTSelectionTest.cpp Wed Aug 30 06:24:37 2017
@@ -494,4 +494,23 @@ void foo() {
       });
 }
 
+TEST(ASTSelectionFinder, CorrectEndForObjectiveCImplementation) {
+  StringRef Source = R"(
+ at interface I
+ at end
+ at 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