[Lex] Preambles should contain the global module fragment.

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.

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

Author: sam-mccall

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

@@ -8,6 +8,7 @@
 
 #include "TestFS.h"
 #include "TestTU.h"
+#include "llvm/ADT/StringRef.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
@@ -109,6 +110,34 @@ TEST(Modules, UnknownFormat) {
   // 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

@@ -705,6 +705,22 @@ PreambleBounds Lexer::ComputePreamble(StringRef Buffer,
       // 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

clang/unittests/Lex/CMakeLists.txt

@@ -25,5 +25,6 @@ clang_target_link_libraries(LexTests
 
 target_link_libraries(LexTests
   PRIVATE
+  LLVMTestingAnnotations
   LLVMTestingSupport
   )

clang/unittests/Lex/LexerTest.cpp

@@ -26,9 +26,11 @@
 #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 {
@@ -660,4 +662,36 @@ TEST_F(LexerTest, RawAndNormalLexSameForLineComments) {
   }
   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