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

[ilya-biryukov]

LGTM

lib/Tooling/Inclusions/HeaderIncludes.cpp

77

typo NITs:
s/header guards presents/header guard is present
s/will returns/will return

84

This function is getting really complex after this change. The comments help, but still.
I don't have a good recipe to improve it, though, especially without changing the whole file, so happy to keep as is.

[ioeric]

• fix typo

lib/Tooling/Inclusions/HeaderIncludes.cpp

 //===--- HeaderIncludes.cpp - Insert/Delete #includes --*- C++ -*----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
 namespace tooling {
 namespace {
 
 LangOptions createLangOpts() {
   LangOptions LangOpts;
unsigned getOffsetAfterTokenSequence(
   Token Tok;
   // Get the first token.
   Lex.LexFromRawLexer(Tok);
   return GetOffsetAfterSequence(SM, Lex, Tok);
 }
 
 // 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;
 }
 
 void skipComments(Lexer &Lex, Token &Tok) {
   while (Tok.is(tok::comment))
     if (Lex.LexFromRawLexer(Tok))
       return;
 }
 
 // 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) {
+  // \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,
-      [](const SourceManager &SM, Lexer &Lex, Token Tok) {
+            [&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 InitialOffset;
-      });
+            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
 //    "#include ("header.h" | <header.h>)".
 // 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 checkAndConsumeInclusiveDirective(Lexer &Lex, Token &Tok) {
   auto Matched = [&]() {

unittests/Tooling/HeaderIncludesTest.cpp

 //===- unittest/Tooling/CleanupTest.cpp - Include insertion/deletion tests ===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/Tooling/Inclusions/HeaderIncludes.h"
 #include "../Tooling/ReplacementTest.h"
 #include "../Tooling/RewriterTestContext.h"
 #include "clang/Format/Format.h"
 #include "clang/Tooling/Core/Replacement.h"
 
 #include "gtest/gtest.h"
 
-using clang::tooling::ReplacementTest;
-using clang::tooling::toReplacements;
-
 namespace clang {
 namespace tooling {
 namespace {
 
 class HeaderIncludesTest : public ::testing::Test {
 protected:
   std::string insert(llvm::StringRef Code, llvm::StringRef Header) {
     HeaderIncludes Includes(FileName, Code, Style);
std::string Expected = "// comment \n"
                          "#ifndef X\n"
                          "#define X\n"
                          "#include <vector>\n"
                          "int x;\n"
                          "#define Y 1\n";
   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"
                      "int x;\n"
                      "#define Y 1\n";
   std::string Expected = "// comment \n"
                          "#include <vector>\n"
                          "#ifndef X\n"