From 37e13a87ec632a1079f81b6b5e0df27b974009eb Mon Sep 17 00:00:00 2001
From: Chris B <cbieneman@microsoft.com>
Date: Mon, 23 May 2022 11:16:08 -0500
Subject: [PATCH] Teach llvm-dis to read DXIL out of DXBC files (#4451)

* Teach llvm-dis to read DXIL out of DXBC files

This change has no impact on the DXC codebase since we don't use
llvm-dis here, but in LLVM upstream I'm using llvm-dis from DXC to
verify bitcode compatability.

The change here is to intercept the file passed into llvm-dis to detect
DXIL Container files, and read the bitcode out of the DXIL part.

* Updating error handling code

Thanks for the review @tex3d!
---
 tools/llvm-dis/CMakeLists.txt |  1 +
 tools/llvm-dis/llvm-dis.cpp   | 51 ++++++++++++++++++++++++-----------
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/tools/llvm-dis/CMakeLists.txt b/tools/llvm-dis/CMakeLists.txt
index b05f77bc8..accd68cd0 100644
--- a/tools/llvm-dis/CMakeLists.txt
+++ b/tools/llvm-dis/CMakeLists.txt
@@ -1,6 +1,7 @@
 set(LLVM_LINK_COMPONENTS
   BitReader
   Core
+  DxilContainer
   Support
   MSSupport # HLSL Change
   )
diff --git a/tools/llvm-dis/llvm-dis.cpp b/tools/llvm-dis/llvm-dis.cpp
index 29c45c89b..f6cc4fe56 100644
--- a/tools/llvm-dis/llvm-dis.cpp
+++ b/tools/llvm-dis/llvm-dis.cpp
@@ -34,6 +34,7 @@
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include "dxc/DxilContainer/DxilContainer.h" // HLSL Change
 #include <system_error>
 using namespace llvm;
 
@@ -162,24 +163,44 @@ int __cdecl main(int argc, char **argv) { // HLSL Change - __cdecl
   std::string ErrorMessage;
   std::unique_ptr<Module> M;
 
-  // Use the bitcode streaming interface
-  std::unique_ptr<DataStreamer> Streamer =
-      getDataFileStreamer(InputFilename, &ErrorMessage);
-  if (Streamer) {
-    std::string DisplayFilename;
-    if (InputFilename == "-")
-      DisplayFilename = "<stdin>";
-    else
-      DisplayFilename = InputFilename;
-    ErrorOr<std::unique_ptr<Module>> MOrErr =
-        getStreamedBitcodeModule(DisplayFilename, std::move(Streamer), Context);
-    M = std::move(*MOrErr);
-    M->materializeAllPermanently();
-  } else {
-    errs() << argv[0] << ": " << ErrorMessage << '\n';
+  ErrorOr<std::unique_ptr<MemoryBuffer>> FileOrErr =
+      MemoryBuffer::getFileOrSTDIN(InputFilename);
+  if (std::error_code EC = FileOrErr.getError()) {
+    errs() << argv[0] << ": "
+           << "Could not open file '" << InputFilename << "': " << EC.message()
+           << '\n';
     return 1;
   }
 
+  std::unique_ptr<MemoryBuffer> &Buf = FileOrErr.get();
+  MemoryBufferRef BitcodeData = Buf->getMemBufferRef();
+  if (Buf->getBuffer().startswith("DXBC")) {
+    // move along until I get to the bitcode
+    const hlsl::DxilContainerHeader *Header =
+        reinterpret_cast<const hlsl::DxilContainerHeader *>(Buf->getBufferStart());
+    const hlsl::DxilProgramHeader *DXILHeader =
+        hlsl::GetDxilProgramHeader(Header, hlsl::DFCC_DXIL);
+    if (!DXILHeader) {
+      errs() << argv[0] << ": DXBC file '" << InputFilename
+             << "': Does not contain DXIL part\n";
+      return 1;
+    }
+    StringRef DXILData = StringRef(hlsl::GetDxilBitcodeData(DXILHeader),
+                                   hlsl::GetDxilBitcodeSize(DXILHeader));
+    BitcodeData = MemoryBufferRef(DXILData, "");
+  }
+
+  ErrorOr<std::unique_ptr<Module>> MOrErr =
+      parseBitcodeFile(BitcodeData, Context);
+  if (std::error_code EC = MOrErr.getError()) {
+    errs() << argv[0] << ": "
+           << "Could not load bitcode from file '" << InputFilename
+           << "': " << EC.message() << '\n';
+    return 1;
+  }
+  M = std::move(*MOrErr);
+  M->materializeAllPermanently();
+
   // Just use stdout.  We won't actually print anything on it.
   if (DontPrint)
     OutputFilename = "-";