[clang] [llvm] [clang][Driver] Don't ignore -gmodules .gch files (PR #77711)

[Michael Spencer via cfe-commits]

https://github.com/Bigcheese created https://github.com/llvm/llvm-project/pull/77711

A previous commit (82f75ed) made clang ignore .gch files that were not Clang AST files. This broke `-gmodules`, which embeds the Clang AST into an object file containing debug info.

This changes the probing to detect any file format recognized by `llvm::identify_magic()` as potentially containing a Clang AST.

Previous PR: https://github.com/llvm/llvm-project/pull/69204

>From e2cd3ee8741eabc1f251f64d389ed5dcd07810d5 Mon Sep 17 00:00:00 2001
From: Michael Spencer <bigcheesegs at gmail.com>
Date: Wed, 10 Jan 2024 16:58:59 -0800
Subject: [PATCH] [clang][Driver] Don't ignore -gmodules .gch files

A previous commit (82f75ed) made clang ignore .gch files that were not
Clang AST files. This broke `-gmodules`, which embeds the Clang AST
into an object file containing debug info.

This changes the probing to detect any file format recognized by
`llvm::identify_magic()` as potentially containing a Clang AST.
---
 clang/lib/Driver/ToolChains/Clang.cpp  | 12 ++++++++----
 clang/test/PCH/gch-probe.c             |  4 ++++
 llvm/include/llvm/BinaryFormat/Magic.h |  1 +
 llvm/lib/BinaryFormat/Magic.cpp        |  2 ++
 llvm/lib/Object/Binary.cpp             |  1 +
 llvm/lib/Object/ObjectFile.cpp         |  1 +
 6 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index 2d8ef841d4f6be..36f1545e285e14 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -43,6 +43,7 @@
 #include "clang/Driver/XRayArgs.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/BinaryFormat/Magic.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Support/CodeGen.h"
@@ -948,10 +949,13 @@ static void handleAMDGPUCodeObjectVersionOptions(const Driver &D,
   }
 }
 
-static bool hasClangPchSignature(const Driver &D, StringRef Path) {
+static bool maybeHasClangPchSignature(const Driver &D, StringRef Path) {
   if (llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> MemBuf =
           D.getVFS().getBufferForFile(Path))
-    return (*MemBuf)->getBuffer().starts_with("CPCH");
+    // Return true for both raw Clang AST files and object files which may
+    // contain a __clangast section.
+    return llvm::identify_magic((*MemBuf)->getBuffer()) !=
+           llvm::file_magic::unknown;
   return false;
 }
 
@@ -964,14 +968,14 @@ static bool gchProbe(const Driver &D, StringRef Path) {
     std::error_code EC;
     for (llvm::vfs::directory_iterator DI = D.getVFS().dir_begin(Path, EC), DE;
          !EC && DI != DE; DI = DI.increment(EC)) {
-      if (hasClangPchSignature(D, DI->path()))
+      if (maybeHasClangPchSignature(D, DI->path()))
         return true;
     }
     D.Diag(diag::warn_drv_pch_ignoring_gch_dir) << Path;
     return false;
   }
 
-  if (hasClangPchSignature(D, Path))
+  if (maybeHasClangPchSignature(D, Path))
     return true;
   D.Diag(diag::warn_drv_pch_ignoring_gch_file) << Path;
   return false;
diff --git a/clang/test/PCH/gch-probe.c b/clang/test/PCH/gch-probe.c
index 8b1e1fab5ad97b..0905c9baebdae0 100644
--- a/clang/test/PCH/gch-probe.c
+++ b/clang/test/PCH/gch-probe.c
@@ -2,6 +2,10 @@
 // RUN: %clang -x c-header -c %s -o %t.h.gch
 // RUN: %clang -fsyntax-only -include %t.h %s
 
+// -gmodules embeds the Clang AST file in an object file.
+// RUN: %clang -x c-header -c %s -gmodules -o %t.h.gch
+// RUN: %clang -fsyntax-only -include %t.h %s
+
 // gch probing should ignore files which are not clang pch files.
 // RUN: %clang -fsyntax-only -include %S/Inputs/gch-probe.h %s 2>&1 | FileCheck %s
 // CHECK: warning: precompiled header '{{.*}}gch-probe.h.gch' was ignored because it is not a clang PCH file
diff --git a/llvm/include/llvm/BinaryFormat/Magic.h b/llvm/include/llvm/BinaryFormat/Magic.h
index c635a269576587..6978c066bda468 100644
--- a/llvm/include/llvm/BinaryFormat/Magic.h
+++ b/llvm/include/llvm/BinaryFormat/Magic.h
@@ -21,6 +21,7 @@ struct file_magic {
   enum Impl {
     unknown = 0,       ///< Unrecognized file
     bitcode,           ///< Bitcode file
+    clang_ast,         ///< Clang PCH or PCM
     archive,           ///< ar style archive file
     elf,               ///< ELF Unknown type
     elf_relocatable,   ///< ELF Relocatable object file
diff --git a/llvm/lib/BinaryFormat/Magic.cpp b/llvm/lib/BinaryFormat/Magic.cpp
index 45a0b7e11452b4..bd378337ed3338 100644
--- a/llvm/lib/BinaryFormat/Magic.cpp
+++ b/llvm/lib/BinaryFormat/Magic.cpp
@@ -98,6 +98,8 @@ file_magic llvm::identify_magic(StringRef Magic) {
   case 'C':
     if (startswith(Magic, "CCOB"))
       return file_magic::offload_bundle_compressed;
+    if (startswith(Magic, "CPCH"))
+      return file_magic::clang_ast;
     break;
   case '!':
     if (startswith(Magic, "!<arch>\n") || startswith(Magic, "!<thin>\n"))
diff --git a/llvm/lib/Object/Binary.cpp b/llvm/lib/Object/Binary.cpp
index 0b9d95485287dc..2dfae8ab5d3c64 100644
--- a/llvm/lib/Object/Binary.cpp
+++ b/llvm/lib/Object/Binary.cpp
@@ -84,6 +84,7 @@ Expected<std::unique_ptr<Binary>> object::createBinary(MemoryBufferRef Buffer,
     // PDB does not support the Binary interface.
     return errorCodeToError(object_error::invalid_file_type);
   case file_magic::unknown:
+  case file_magic::clang_ast:
   case file_magic::cuda_fatbinary:
   case file_magic::coff_cl_gl_object:
   case file_magic::dxcontainer_object:
diff --git a/llvm/lib/Object/ObjectFile.cpp b/llvm/lib/Object/ObjectFile.cpp
index c05eb0a0468e2b..6a226a3bbdbca3 100644
--- a/llvm/lib/Object/ObjectFile.cpp
+++ b/llvm/lib/Object/ObjectFile.cpp
@@ -155,6 +155,7 @@ ObjectFile::createObjectFile(MemoryBufferRef Object, file_magic Type,
   switch (Type) {
   case file_magic::unknown:
   case file_magic::bitcode:
+  case file_magic::clang_ast:
   case file_magic::coff_cl_gl_object:
   case file_magic::archive:
   case file_magic::macho_universal_binary: