[Lex] Preambles should contain the global module fragment.

[sammccall]

Summary

For applications like clangd, the preamble remains an important optimization
when editing a module definition. The global module fragment is a good fit for
it as it by definition contains only preprocessor directives.
Before this patch, we would terminate the preamble immediately at the "module"
keyword.

[ChuanqiXu]

LGTM. Thanks!

clang-tools-extra/clangd/unittests/ModulesTests.cpp

 //===-- ModulesTests.cpp  ---------------------------------------*- 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 "TestFS.h"
 #include "TestTU.h"
+#include "llvm/ADT/StringRef.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 #include <memory>
 #include <string>
 
 namespace clang {
 namespace clangd {
TEST(Modules, UnknownFormat) {
   TU.AdditionalFiles["m.modulemap"] = R"modulemap(
     module M {
       header "modular.h"
     })modulemap";
 
   // Test that we do not crash.
   TU.build();
 }
+
+// Test that we can build and use a preamble for a module unit.
+TEST(Modules, ModulePreamble) {
+  TestTU TU = TestTU::withCode(R"cpp(
+    module;
+    #define PREAMBLE_MACRO 1
+    export module foo;
+    #define MODULE_MACRO 2
+    module :private;
+    #define PRIVATE_MACRO 3
+  )cpp");
+  TU.ExtraArgs.push_back("-std=c++20");
+  TU.ExtraArgs.push_back("--precompile");
+
+  auto AST = TU.build();
+  auto &SM = AST.getSourceManager();
+  auto GetMacroFile = [&](llvm::StringRef Name) -> FileID {
+    if (auto *MI = AST.getPreprocessor().getMacroInfo(
+            &AST.getASTContext().Idents.get(Name)))
+      return SM.getFileID(MI->getDefinitionLoc());
+    return {};
+  };
+
+  EXPECT_EQ(GetMacroFile("PREAMBLE_MACRO"), SM.getPreambleFileID());
+  EXPECT_EQ(GetMacroFile("MODULE_MACRO"), SM.getMainFileID());
+  EXPECT_EQ(GetMacroFile("PRIVATE_MACRO"), SM.getMainFileID());
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

clang/lib/Lex/Lexer.cpp

if (TheTok.isAtStartOfLine() && TheTok.getKind() == tok::hash) {
           break;
         }
       }
 
       // We only end up here if we didn't recognize the preprocessor
       // directive or it was one that can't occur in the preamble at this
       // point. Roll back the current token to the location of the '#'.
       TheTok = HashTok;
+    } else if (TheTok.isAtStartOfLine() &&
+               TheTok.getKind() == tok::raw_identifier &&
+               TheTok.getRawIdentifier() == "module" &&
+               LangOpts.CPlusPlusModules) {
+      // The initial global module fragment introducer "module;" is part of
+      // the preamble, which runs up to the module declaration "module foo;".
+      Token ModuleTok = TheTok;
+      do {
+        TheLexer.LexFromRawLexer(TheTok);
+      } while (TheTok.getKind() == tok::comment);
+      if (TheTok.getKind() != tok::semi) {
+        // Not global module fragment, roll back.
+        TheTok = ModuleTok;
+        break;
+      }
+      continue;
     }
 
     // We hit a token that we don't recognize as being in the
     // "preprocessing only" part of the file, so we're no longer in
     // the preamble.
     break;
   } while (true);
 

clang/unittests/Lex/CMakeLists.txt

clang_target_link_libraries(LexTests
   clangBasic
   clangLex
   clangParse
   clangSema
   )
 
 target_link_libraries(LexTests
   PRIVATE
+  LLVMTestingAnnotations
   LLVMTestingSupport
   )

clang/unittests/Lex/LexerTest.cpp

 #include "clang/Lex/LiteralSupport.h"
 #include "clang/Lex/MacroArgs.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/ModuleLoader.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Testing/Annotations/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <memory>
+#include <string>
 #include <vector>
 
 namespace {
 using namespace clang;
 using testing::ElementsAre;
 
 // The test fixture.
 class LexerTest : public ::testing::Test {
TEST_F(LexerTest, RawAndNormalLexSameForLineComments) {
   EXPECT_FALSE(ToksView.empty());
   while (!L.LexFromRawLexer(T)) {
     ASSERT_TRUE(!ToksView.empty());
     EXPECT_EQ(T.getKind(), ToksView.front().getKind());
     ToksView = ToksView.drop_front();
   }
   EXPECT_TRUE(ToksView.empty());
 }
+
+TEST(LexerPreambleTest, PreambleBounds) {
+  std::vector<std::string> Cases = {
+      R"cc([[
+        #include <foo>
+        ]]int bar;
+      )cc",
+      R"cc([[
+        #include <foo>
+      ]])cc",
+      R"cc([[
+        // leading comment
+        #include <foo>
+        ]]// trailing comment
+        int bar;
+      )cc",
+      R"cc([[
+        module;
+        #include <foo>
+        ]]module bar;
+        int x;
+      )cc",
+  };
+  for (const auto& Case : Cases) {
+    llvm::Annotations A(Case);
+    clang::LangOptions LangOpts;
+    LangOpts.CPlusPlusModules = true;
+    auto Bounds = Lexer::ComputePreamble(A.code(), LangOpts);
+    EXPECT_EQ(Bounds.Size, A.range().End) << Case;
+  }
+}
+
 } // anonymous namespace