diff --git a/clang/lib/Tooling/Refactoring/ASTSelection.cpp b/clang/lib/Tooling/Refactoring/ASTSelection.cpp index 0bf5d95624eb..4fa008669781 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(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 95326dded88b..bde052e849ca 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 Node) { + EXPECT_TRUE(Node); + EXPECT_EQ(Node->Children.size(), 1u); + checkNode( + Node->Children[0], + SourceSelectionKind::ContainsSelection); + }, + SelectionFinderVisitor::Lang_OBJC); +} + } // end anonymous namespace