[Tooling] Handle #pragma once header guard in include insertion.

Reviewers: ilya-biryukov

Subscribers: cfe-commits

Differential Revision: https://reviews.llvm.org/D57223

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@352503 91177308-0d34-0410-b5e6-96231b3b80d8

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 @@ unsigned getOffsetAfterTokenSequence(
 
 // Check if a sequence of tokens is like "#<Name> <raw_identifier>". 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<StringRef> 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 @@ void skipComments(Lexer &Lex, Token &Tok) {
 }
 
 // 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 is present in the code, this will return 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<unsigned(const SourceManager &SM, Lexer &Lex,
+                                 Token Tok)>
+              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

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 @@ TEST_F(HeaderIncludesTest, RealHeaderGuardAfterComments) {
   EXPECT_EQ(Expected, insert(Code, "<vector>"));
 }
 
+TEST_F(HeaderIncludesTest, PragmaOnce) {
+  std::string Code = "// comment \n"
+                     "#pragma once\n"
+                     "int x;\n";
+  std::string Expected = "// comment \n"
+                         "#pragma once\n"
+                         "#include <vector>\n"
+                         "int x;\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
 TEST_F(HeaderIncludesTest, IfNDefWithNoDefine) {
   std::string Code = "// comment \n"
                      "#ifndef X\n"