Index: lib/Tooling/Inclusions/HeaderIncludes.cpp =================================================================== --- lib/Tooling/Inclusions/HeaderIncludes.cpp +++ lib/Tooling/Inclusions/HeaderIncludes.cpp @@ -9,6 +9,7 @@ #include "clang/Tooling/Inclusions/HeaderIncludes.h" #include "clang/Basic/SourceManager.h" #include "clang/Lex/Lexer.h" +#include "llvm/ADT/Optional.h" #include "llvm/Support/FormatVariadic.h" namespace clang { @@ -50,12 +51,16 @@ // Check if a sequence of tokens is like "# ". If it is, // \p Tok will be the token after this directive; otherwise, it can be any token -// after the given \p Tok (including \p Tok). -bool checkAndConsumeDirectiveWithName(Lexer &Lex, StringRef Name, Token &Tok) { +// after the given \p Tok (including \p Tok). If \p RawIDName is provided, the +// (second) raw_identifier name is checked. +bool checkAndConsumeDirectiveWithName( + Lexer &Lex, StringRef Name, Token &Tok, + llvm::Optional RawIDName = llvm::None) { bool Matched = Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) && Tok.is(tok::raw_identifier) && Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) && - Tok.is(tok::raw_identifier); + Tok.is(tok::raw_identifier) && + (!RawIDName || Tok.getRawIdentifier() == *RawIDName); if (Matched) Lex.LexFromRawLexer(Tok); return Matched; @@ -68,24 +73,45 @@ } // Returns the offset after header guard directives and any comments -// before/after header guards. If no header guard presents in the code, this -// will returns the offset after skipping all comments from the start of the -// code. +// before/after header guards (e.g. #ifndef/#define pair, #pragma once). If no +// header guard presents in the code, this will returns the offset after +// skipping all comments from the start of the code. unsigned getOffsetAfterHeaderGuardsAndComments(StringRef FileName, StringRef Code, const IncludeStyle &Style) { - return getOffsetAfterTokenSequence( - FileName, Code, Style, - [](const SourceManager &SM, Lexer &Lex, Token Tok) { - skipComments(Lex, Tok); - unsigned InitialOffset = SM.getFileOffset(Tok.getLocation()); - if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) { - skipComments(Lex, Tok); - if (checkAndConsumeDirectiveWithName(Lex, "define", Tok)) - return SM.getFileOffset(Tok.getLocation()); - } - return InitialOffset; - }); + // \p Consume returns location after header guard or 0 if no header guard is + // found. + auto ConsumeHeaderGuardAndComment = + [&](std::function + Consume) { + return getOffsetAfterTokenSequence( + FileName, Code, Style, + [&Consume](const SourceManager &SM, Lexer &Lex, Token Tok) { + skipComments(Lex, Tok); + unsigned InitialOffset = SM.getFileOffset(Tok.getLocation()); + return std::max(InitialOffset, Consume(SM, Lex, Tok)); + }); + }; + return std::max( + // #ifndef/#define + ConsumeHeaderGuardAndComment( + [](const SourceManager &SM, Lexer &Lex, Token Tok) -> unsigned { + if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) { + skipComments(Lex, Tok); + if (checkAndConsumeDirectiveWithName(Lex, "define", Tok)) + return SM.getFileOffset(Tok.getLocation()); + } + return 0; + }), + // #pragma once + ConsumeHeaderGuardAndComment( + [](const SourceManager &SM, Lexer &Lex, Token Tok) -> unsigned { + if (checkAndConsumeDirectiveWithName(Lex, "pragma", Tok, + StringRef("once"))) + return SM.getFileOffset(Tok.getLocation()); + return 0; + })); } // Check if a sequence of tokens is like Index: unittests/Tooling/HeaderIncludesTest.cpp =================================================================== --- unittests/Tooling/HeaderIncludesTest.cpp +++ unittests/Tooling/HeaderIncludesTest.cpp @@ -14,9 +14,6 @@ #include "gtest/gtest.h" -using clang::tooling::ReplacementTest; -using clang::tooling::toReplacements; - namespace clang { namespace tooling { namespace { @@ -315,6 +312,17 @@ EXPECT_EQ(Expected, insert(Code, "")); } +TEST_F(HeaderIncludesTest, PragmaOnce) { + std::string Code = "// comment \n" + "#pragma once\n" + "int x;\n"; + std::string Expected = "// comment \n" + "#pragma once\n" + "#include \n" + "int x;\n"; + EXPECT_EQ(Expected, insert(Code, "")); +} + TEST_F(HeaderIncludesTest, IfNDefWithNoDefine) { std::string Code = "// comment \n" "#ifndef X\n"