[clangd] Fix crash-bug in preamble indexing when using modules.

When preamble contains #undef, indexing code finds the matching #define
and uses that during indexing. However, it would only look for local
definitions. If the macro was defined in a module, MacroInfo
would be nullptr and clangd would crash.

This change makes clangd ignore any #undef without a matching #define
inside the same TU.

The indexing of macros happens for preamble only, so then #undef must be
in the preamble, which is why we need two .h files in a test.

Note that clangd is currently not ready for module support, but this
brings us one step closer.

This was previously attempted in
4061d9e42c, but had to be reverted due to
broken test. This version fixes that test-only bug by setting a custom module
cache path to avoid re-use of modules across test invocations.

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

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

@@ -1623,6 +1623,31 @@ TEST_F(SymbolCollectorTest, MacrosInHeaders) {
   EXPECT_THAT(Symbols,
               UnorderedElementsAre(AllOf(QName("X"), ForCodeCompletion(true))));
 }
+
+// Regression test for a crash-bug we used to have.
+TEST_F(SymbolCollectorTest, UndefOfModuleMacro) {
+  auto TU = TestTU::withCode(R"cpp(#include "bar.h")cpp");
+  TU.AdditionalFiles["bar.h"] = R"cpp(
+    #include "foo.h"
+    #undef X
+    )cpp";
+  TU.AdditionalFiles["foo.h"] = "#define X 1";
+  TU.AdditionalFiles["module.map"] = R"cpp(
+    module foo {
+     header "foo.h"
+     export *
+   }
+   )cpp";
+  TU.ExtraArgs.push_back("-fmodules");
+  TU.ExtraArgs.push_back("-fmodule-map-file=" + testPath("module.map"));
+  TU.OverlayRealFileSystemForModules = true;
+
+  TU.build();
+  // We mostly care about not crashing, but verify that we didn't insert garbage
+  // about X too.
+  EXPECT_THAT(TU.headerSymbols(), Not(Contains(QName("X"))));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/unittests/TestFS.h

@@ -34,12 +34,21 @@ buildTestFS(llvm::StringMap<std::string> const &Files,
 class MockFS : public ThreadsafeFS {
 public:
   IntrusiveRefCntPtr<llvm::vfs::FileSystem> viewImpl() const override {
-    return buildTestFS(Files, Timestamps);
+    auto MemFS = buildTestFS(Files, Timestamps);
+    if (!OverlayRealFileSystemForModules)
+      return MemFS;
+    llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> OverlayFileSystem =
+        new llvm::vfs::OverlayFileSystem(llvm::vfs::getRealFileSystem());
+    OverlayFileSystem->pushOverlay(MemFS);
+    return OverlayFileSystem;
   }
 
   // If relative paths are used, they are resolved with testPath().
   llvm::StringMap<std::string> Files;
   llvm::StringMap<time_t> Timestamps;
+  // If true, real file system will be used as fallback for the in-memory one.
+  // This is useful for testing module support.
+  bool OverlayRealFileSystemForModules = false;
 };
 
 // A Compilation database that returns a fixed set of compile flags.

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

@@ -16,6 +16,7 @@
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/Utils.h"
+#include "llvm/ADT/ScopeExit.h"
 
 namespace clang {
 namespace clangd {
@@ -54,6 +55,8 @@ ParseInputs TestTU::inputs(MockFS &FS) const {
   Inputs.CompileCommand.Filename = FullFilename;
   Inputs.CompileCommand.Directory = testRoot();
   Inputs.Contents = Code;
+  if (OverlayRealFileSystemForModules)
+    FS.OverlayRealFileSystemForModules = true;
   Inputs.TFS = &FS;
   Inputs.Opts = ParseOptions();
   Inputs.Opts.BuildRecoveryAST = true;
@@ -66,12 +69,31 @@ ParseInputs TestTU::inputs(MockFS &FS) const {
   return Inputs;
 }
 
+void initializeModuleCache(CompilerInvocation &CI) {
+  llvm::SmallString<128> ModuleCachePath;
+  ASSERT_FALSE(
+      llvm::sys::fs::createUniqueDirectory("module-cache", ModuleCachePath));
+  CI.getHeaderSearchOpts().ModuleCachePath = ModuleCachePath.c_str();
+  llvm::errs() << "MC: " << ModuleCachePath << "\n";
+  llvm::errs().flush();
+}
+
+void deleteModuleCache(const std::string ModuleCachePath) {
+  if (!ModuleCachePath.empty()) {
+    ASSERT_FALSE(llvm::sys::fs::remove_directories(ModuleCachePath));
+  }
+}
+
 std::shared_ptr<const PreambleData> TestTU::preamble() const {
   MockFS FS;
   auto Inputs = inputs(FS);
   IgnoreDiagnostics Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);
   assert(CI && "Failed to build compilation invocation.");
+  if (OverlayRealFileSystemForModules)
+    initializeModuleCache(*CI);
+  auto ModuleCacheDeleter = llvm::make_scope_exit(
+      std::bind(deleteModuleCache, CI->getHeaderSearchOpts().ModuleCachePath));
   return clang::clangd::buildPreamble(testPath(Filename), *CI, Inputs,
                                       /*StoreInMemory=*/true,
                                       /*PreambleCallback=*/nullptr);
@@ -83,6 +105,11 @@ ParsedAST TestTU::build() const {
   StoreDiags Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);
   assert(CI && "Failed to build compilation invocation.");
+  if (OverlayRealFileSystemForModules)
+    initializeModuleCache(*CI);
+  auto ModuleCacheDeleter = llvm::make_scope_exit(
+      std::bind(deleteModuleCache, CI->getHeaderSearchOpts().ModuleCachePath));
+
   auto Preamble = clang::clangd::buildPreamble(testPath(Filename), *CI, Inputs,
                                                /*StoreInMemory=*/true,
                                                /*PreambleCallback=*/nullptr);

clang-tools-extra/clangd/unittests/TestTU.h

@@ -66,6 +66,16 @@ struct TestTU {
   // Simulate a header guard of the header (using an #import directive).
   bool ImplicitHeaderGuard = true;
 
+  // Whether to use overlay the TestFS over the real filesystem. This is
+  // required for use of implicit modules.where the module file is written to
+  // disk and later read back.
+  // FIXME: Change the way reading/writing modules work to allow us to keep them
+  // in memory across multiple clang invocations, at least in tests, to
+  // eliminate the need for real file system here.
+  // Please avoid using this for things other than implicit modules. The plan is
+  // to eliminate this option some day.
+  bool OverlayRealFileSystemForModules = false;
+
   // By default, build() will report Error diagnostics as GTest errors.
   // Suppress this behavior by adding an 'error-ok' comment to the code.
   ParsedAST build() const;

clang/lib/Index/IndexingAction.cpp

@@ -165,11 +165,20 @@ static void indexTranslationUnit(ASTUnit &Unit, IndexingContext &IndexCtx) {
 static void indexPreprocessorMacros(const Preprocessor &PP,
                                     IndexDataConsumer &DataConsumer) {
   for (const auto &M : PP.macros())
-    if (MacroDirective *MD = M.second.getLatest())
+    if (MacroDirective *MD = M.second.getLatest()) {
+      auto *MI = MD->getMacroInfo();
+      // When using modules, it may happen that we find #undef of a macro that
+      // was defined in another module. In such case, MI may be nullptr, since
+      // we only look for macro definitions in the current TU. In that case,
+      // there is nothing to index.
+      if (!MI)
+        continue;
+
       DataConsumer.handleMacroOccurrence(
           M.first, MD->getMacroInfo(),
           static_cast<unsigned>(index::SymbolRole::Definition),
           MD->getLocation());
+    }
 }
 
 void index::indexASTUnit(ASTUnit &Unit, IndexDataConsumer &DataConsumer,