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-exception
6//
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 auto *PartIter = llvm::find_if(
27 Range: Obj.Parts, P: [&PartName](const Part &P) { return P.Name == PartName; });
28 if (PartIter == Obj.Parts.end())
29 return createFileError(F: InputFilename,
30 EC: std::make_error_code(e: std::errc::invalid_argument),
31 Fmt: "part '%s' not found", Vals: PartName.str().c_str());
32
33 Object PartObj;
34 PartObj.Header = Obj.Header;
35 PartObj.Parts.push_back(Elt: {.Name: PartIter->Name, .Data: PartIter->Data});
36 PartObj.recomputeHeader();
37
38 auto Write = [&OutFilename, &PartObj](raw_ostream &Out) -> Error {
39 DXContainerWriter Writer(PartObj, Out);
40 if (Error E = Writer.write())
41 return createFileError(F: OutFilename, E: std::move(E));
42 return Error::success();
43 };
44
45 return writeToOutput(OutputFileName: OutFilename, Write);
46}
47
48static Error dumpPartToFile(StringRef PartName, StringRef Filename,
49 StringRef InputFilename, Object &Obj) {
50 auto *PartIter = llvm::find_if(
51 Range&: Obj.Parts, P: [&PartName](const Part &P) { return P.Name == PartName; });
52 if (PartIter == Obj.Parts.end())
53 return createFileError(F: Filename,
54 EC: std::make_error_code(e: std::errc::invalid_argument),
55 Fmt: "part '%s' not found", Vals: PartName.str().c_str());
56
57 ArrayRef<uint8_t> Contents = PartIter->Data;
58 // The DXContainer format is a bit odd because the part-specific headers are
59 // contained inside the part data itself. For parts that contain LLVM bitcode
60 // when we dump the part we want to skip the part-specific header so that we
61 // get a valid .bc file that we can inspect. All the data contained inside the
62 // program header is pulled out of the bitcode, so the header can be
63 // reconstructed if needed from the bitcode itself. More comprehensive
64 // documentation on the DXContainer format can be found at
65 // https://llvm.org/docs/DirectX/DXContainer.html.
66
67 if (PartName == "DXIL" || PartName == "STAT")
68 Contents = Contents.drop_front(N: sizeof(llvm::dxbc::ProgramHeader));
69 if (Contents.empty())
70 return createFileError(F: Filename, EC: object_error::parse_failed,
71 Fmt: "part '%s' is empty", Vals: PartName.str().c_str());
72 Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
73 FileOutputBuffer::create(FilePath: Filename, Size: Contents.size());
74 if (!BufferOrErr)
75 return createFileError(F: Filename, E: BufferOrErr.takeError());
76 std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);
77 llvm::copy(Range&: Contents, Out: Buf->getBufferStart());
78 if (Error E = Buf->commit())
79 return createFileError(F: Filename, E: std::move(E));
80 return Error::success();
81}
82
83static Error handleArgs(const CommonConfig &Config, Object &Obj) {
84 for (StringRef Flag : Config.DumpSection) {
85 auto [SectionName, FileName] = Flag.split(Separator: "=");
86 if (Error E =
87 dumpPartToFile(PartName: SectionName, Filename: FileName, InputFilename: Config.InputFilename, Obj))
88 return E;
89 }
90
91 // Extract all sections before any modifications.
92 for (StringRef Flag : Config.ExtractSection) {
93 auto [SectionName, FileName] = Flag.split(Separator: '=');
94 if (Error E = extractPartAsObject(PartName: SectionName, OutFilename: FileName,
95 InputFilename: Config.InputFilename, Obj))
96 return E;
97 }
98
99 std::function<bool(const Part &)> RemovePred = [](const Part &) {
100 return false;
101 };
102
103 if (!Config.ToRemove.empty())
104 RemovePred = [&Config](const Part &P) {
105 return Config.ToRemove.matches(S: P.Name);
106 };
107
108 if (!Config.OnlySection.empty())
109 RemovePred = [&Config](const Part &P) {
110 // Explicitly keep these sections regardless of previous removes and
111 // remove everything else.
112 return !Config.OnlySection.matches(S: P.Name);
113 };
114
115 if (auto E = Obj.removeParts(ToRemove: RemovePred))
116 return E;
117
118 Obj.recomputeHeader();
119 return Error::success();
120}
121
122Error executeObjcopyOnBinary(const CommonConfig &Config,
123 const DXContainerConfig &,
124 DXContainerObjectFile &In, raw_ostream &Out) {
125 DXContainerReader Reader(In);
126 Expected<std::unique_ptr<Object>> ObjOrErr = Reader.create();
127 if (!ObjOrErr)
128 return createFileError(F: Config.InputFilename, E: ObjOrErr.takeError());
129 Object *Obj = ObjOrErr->get();
130 assert(Obj && "Unable to deserialize DXContainer object");
131
132 if (Error E = handleArgs(Config, Obj&: *Obj))
133 return E;
134
135 DXContainerWriter Writer(*Obj, Out);
136 if (Error E = Writer.write())
137 return createFileError(F: Config.OutputFilename, E: std::move(E));
138 return Error::success();
139}
140
141} // end namespace dxbc
142} // end namespace objcopy
143} // end namespace llvm
144