Index: lib/Format/Format.cpp =================================================================== --- lib/Format/Format.cpp +++ lib/Format/Format.cpp @@ -1436,6 +1436,24 @@ llvm::Regex(IncludeRegexPattern).match(Replace.getReplacementText()); } +// Skip the current token and all comments following the current token. +void skipComments(Lexer &Lex, Token &Tok) { + while (!Lex.LexFromRawLexer(Tok)) + if (Tok.isNot(tok::comment)) + return; +} + +// Check if a preprocessor directive is like "# ". +// \post If the preprocessor directive matches the given pattern, \p Tok will be +// the raw_identifier in the directive; otherwise, it can be any tokens after +// the hash token. +bool isDirectiveWithName(Lexer &Lex, StringRef Name, Token &Tok) { + assert(Tok.is(tok::hash) && "If-directive must start with hash!"); + return !Lex.LexFromRawLexer(Tok) && Tok.is(tok::raw_identifier) && + Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) && + Tok.is(tok::raw_identifier); +} + // FIXME: do not insert headers into conditional #include blocks, e.g. #includes // surrounded by compile condition "#if...". // FIXME: do not insert existing headers. @@ -1474,15 +1492,26 @@ const SourceManager &SourceMgr = Env->getSourceManager(); Lexer Lex(Env->getFileID(), SourceMgr.getBuffer(Env->getFileID()), SourceMgr, getFormattingLangOpts(Style)); - Token Tok; // All new headers should be inserted after this offset. int MinInsertOffset = Code.size(); - while (!Lex.LexFromRawLexer(Tok)) { - if (Tok.isNot(tok::comment)) { - MinInsertOffset = SourceMgr.getFileOffset(Tok.getLocation()); - break; + Token Tok; + skipComments(Lex, Tok); + Token TokAfterComments = Tok; + bool HeaderGuardFound = false; + if (Tok.is(tok::hash) && isDirectiveWithName(Lex, "ifndef", Tok) && + !Lex.LexFromRawLexer(Tok)) { + if (Tok.is(tok::comment)) + skipComments(Lex, Tok); + if (Tok.is(tok::hash) && isDirectiveWithName(Lex, "define", Tok)) { + HeaderGuardFound = true; + // Get the token after the header guard. + Lex.LexFromRawLexer(Tok); } } + if (!HeaderGuardFound) + Tok = TokAfterComments; + if (Tok.isNot(tok::eof)) + MinInsertOffset = SourceMgr.getFileOffset(Tok.getLocation()); // Record the offset of the end of the last include in each category. std::map CategoryEndOffsets; // All possible priorities. @@ -1491,7 +1520,6 @@ for (const auto &Category : Style.IncludeCategories) Priorities.insert(Category.Priority); int FirstIncludeOffset = -1; - bool HeaderGuardFound = false; StringRef TrimmedCode = Code.drop_front(MinInsertOffset); SmallVector Lines; TrimmedCode.split(Lines, '\n'); @@ -1506,11 +1534,6 @@ FirstIncludeOffset = Offset; } Offset += Line.size() + 1; - // FIXME: make header guard matching stricter, e.g. consider #ifndef. - if (!HeaderGuardFound && DefineRegex.match(Line)) { - HeaderGuardFound = true; - MinInsertOffset = Offset; - } } // Populate CategoryEndOfssets: Index: unittests/Format/CleanupTest.cpp =================================================================== --- unittests/Format/CleanupTest.cpp +++ unittests/Format/CleanupTest.cpp @@ -608,6 +608,70 @@ EXPECT_EQ(Expected, apply(Code, Replaces)); } +TEST_F(CleanUpReplacementsTest, FakeHeaderGuardIfDef) { + std::string Code = "// comment \n" + "#ifdef X\n" + "#define X\n"; + std::string Expected = "// comment \n" + "#include \n" + "#ifdef X\n" + "#define X\n"; + tooling::Replacements Replaces = {createInsertion("#include ")}; + EXPECT_EQ(Expected, apply(Code, Replaces)); +} + +TEST_F(CleanUpReplacementsTest, RealHeaderGuardAfterComments) { + std::string Code = "// comment \n" + "#ifndef X\n" + "#define X\n" + "int x;\n" + "#define Y 1\n"; + std::string Expected = "// comment \n" + "#ifndef X\n" + "#define X\n" + "#include \n" + "int x;\n" + "#define Y 1\n"; + tooling::Replacements Replaces = {createInsertion("#include ")}; + EXPECT_EQ(Expected, apply(Code, Replaces)); +} + +TEST_F(CleanUpReplacementsTest, IfNDefWithNoDefine) { + std::string Code = "// comment \n" + "#ifndef X\n" + "int x;\n" + "#define Y 1\n"; + std::string Expected = "// comment \n" + "#include \n" + "#ifndef X\n" + "int x;\n" + "#define Y 1\n"; + tooling::Replacements Replaces = {createInsertion("#include ")}; + EXPECT_EQ(Expected, apply(Code, Replaces)); +} + +TEST_F(CleanUpReplacementsTest, HeaderGuardWithComment) { + std::string Code = "// comment \n" + "#ifndef X // comment\n" + "// comment\n" + "/* comment\n" + "*/\n" + "/* comment */ #define X\n" + "int x;\n" + "#define Y 1\n"; + std::string Expected = "// comment \n" + "#ifndef X // comment\n" + "// comment\n" + "/* comment\n" + "*/\n" + "/* comment */ #define X\n" + "#include \n" + "int x;\n" + "#define Y 1\n"; + tooling::Replacements Replaces = {createInsertion("#include ")}; + EXPECT_EQ(Expected, apply(Code, Replaces)); +} + } // end namespace } // end namespace format } // end namespace clang