143 lines · cpp
1//===- DXContainerObjcopy.cpp ---------------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/ObjCopy/DXContainer/DXContainerObjcopy.h"10#include "DXContainerReader.h"11#include "DXContainerWriter.h"12#include "llvm/BinaryFormat/DXContainer.h"13#include "llvm/ObjCopy/CommonConfig.h"14#include "llvm/ObjCopy/DXContainer/DXContainerConfig.h"15#include "llvm/Support/FileOutputBuffer.h"16#include "llvm/Support/raw_ostream.h"17 18namespace llvm {19namespace objcopy {20namespace dxbc {21 22using namespace object;23 24static Error extractPartAsObject(StringRef PartName, StringRef OutFilename,25 StringRef InputFilename, const Object &Obj) {26 for (const Part &P : Obj.Parts)27 if (P.Name == PartName) {28 Object PartObj;29 PartObj.Header = Obj.Header;30 PartObj.Parts.push_back({P.Name, P.Data});31 PartObj.recomputeHeader();32 33 auto Write = [&OutFilename, &PartObj](raw_ostream &Out) -> Error {34 DXContainerWriter Writer(PartObj, Out);35 if (Error E = Writer.write())36 return createFileError(OutFilename, std::move(E));37 return Error::success();38 };39 40 return writeToOutput(OutFilename, Write);41 }42 43 return createFileError(InputFilename, object_error::parse_failed,44 "part '%s' not found", PartName.str().c_str());45}46 47static Error dumpPartToFile(StringRef PartName, StringRef Filename,48 StringRef InputFilename, Object &Obj) {49 auto PartIter = llvm::find_if(50 Obj.Parts, [&PartName](const Part &P) { return P.Name == PartName; });51 if (PartIter == Obj.Parts.end())52 return createFileError(Filename,53 std::make_error_code(std::errc::invalid_argument),54 "part '%s' not found", PartName.str().c_str());55 ArrayRef<uint8_t> Contents = PartIter->Data;56 // The DXContainer format is a bit odd because the part-specific headers are57 // contained inside the part data itself. For parts that contain LLVM bitcode58 // when we dump the part we want to skip the part-specific header so that we59 // get a valid .bc file that we can inspect. All the data contained inside the60 // program header is pulled out of the bitcode, so the header can be61 // reconstructed if needed from the bitcode itself. More comprehensive62 // documentation on the DXContainer format can be found at63 // https://llvm.org/docs/DirectX/DXContainer.html.64 65 if (PartName == "DXIL" || PartName == "STAT")66 Contents = Contents.drop_front(sizeof(llvm::dxbc::ProgramHeader));67 if (Contents.empty())68 return createFileError(Filename, object_error::parse_failed,69 "part '%s' is empty", PartName.str().c_str());70 Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =71 FileOutputBuffer::create(Filename, Contents.size());72 if (!BufferOrErr)73 return createFileError(Filename, BufferOrErr.takeError());74 std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);75 llvm::copy(Contents, Buf->getBufferStart());76 if (Error E = Buf->commit())77 return createFileError(Filename, std::move(E));78 return Error::success();79}80 81static Error handleArgs(const CommonConfig &Config, Object &Obj) {82 for (StringRef Flag : Config.DumpSection) {83 auto [SecName, FileName] = Flag.split("=");84 if (Error E = dumpPartToFile(SecName, FileName, Config.InputFilename, Obj))85 return E;86 }87 88 // Extract all sections before any modifications.89 for (StringRef Flag : Config.ExtractSection) {90 StringRef SectionName;91 StringRef FileName;92 std::tie(SectionName, FileName) = Flag.split('=');93 if (Error E = extractPartAsObject(SectionName, FileName,94 Config.InputFilename, Obj))95 return E;96 }97 98 std::function<bool(const Part &)> RemovePred = [](const Part &) {99 return false;100 };101 102 if (!Config.ToRemove.empty())103 RemovePred = [&Config](const Part &P) {104 return Config.ToRemove.matches(P.Name);105 };106 107 if (!Config.OnlySection.empty())108 RemovePred = [&Config](const Part &P) {109 // Explicitly keep these sections regardless of previous removes and110 // remove everything else.111 return !Config.OnlySection.matches(P.Name);112 };113 114 if (auto E = Obj.removeParts(RemovePred))115 return E;116 117 Obj.recomputeHeader();118 return Error::success();119}120 121Error executeObjcopyOnBinary(const CommonConfig &Config,122 const DXContainerConfig &,123 DXContainerObjectFile &In, raw_ostream &Out) {124 DXContainerReader Reader(In);125 Expected<std::unique_ptr<Object>> ObjOrErr = Reader.create();126 if (!ObjOrErr)127 return createFileError(Config.InputFilename, ObjOrErr.takeError());128 Object *Obj = ObjOrErr->get();129 assert(Obj && "Unable to deserialize DXContainer object");130 131 if (Error E = handleArgs(Config, *Obj))132 return E;133 134 DXContainerWriter Writer(*Obj, Out);135 if (Error E = Writer.write())136 return createFileError(Config.OutputFilename, std::move(E));137 return Error::success();138}139 140} // end namespace dxbc141} // end namespace objcopy142} // end namespace llvm143