[DirectX] Add extract-section to llvm-objcopy and implement it for DXContainer

llvm/docs/CommandGuide/llvm-objcopy.rst

@@ -79,6 +79,18 @@ multiple file formats.
  Enable deterministic mode when copying archives, i.e. use 0 for archive member
  header UIDs, GIDs and timestamp fields. On by default.
 
+.. option:: --extract-section <section>=<file>
+
+ Extract the specified section ``<section>`` into the file ``<file>`` as a
+ seperate object. Can be specified multiple times to extract multiple sections.
+ ``<file>`` is unrelated to the input and output files provided to
+ :program:`llvm-objcopy` and as such the normal copying and editing
+ operations will still be performed. No operations are performed on the sections
+ prior to dumping them.
+
+ For MachO objects, ``<section>`` must be formatted as
+ ``<segment name>,<section name>``.
+
 .. option:: --globalize-symbol <symbol>
 
  Mark any defined symbols named ``<symbol>`` as global symbols in the output.

llvm/include/llvm/ObjCopy/CommonConfig.h

@@ -233,6 +233,7 @@ struct CommonConfig {
   SmallVector<StringRef, 0> DumpSection;
   SmallVector<NewSectionInfo, 0> UpdateSection;
   SmallVector<SectionPatternAddressUpdate, 0> ChangeSectionAddress;
+  SmallVector<StringRef, 0> ExtractSection;
 
   // Section matchers
   NameMatcher KeepSection;

llvm/include/llvm/ObjCopy/ConfigManager.h

@@ -27,7 +27,7 @@ struct LLVM_ABI ConfigManager : public MultiFormatConfig {
 
   const CommonConfig &getCommonConfig() const override { return Common; }
 
-  Expected<const ELFConfig &> getELFConfig() const override { return ELF; }
+  Expected<const ELFConfig &> getELFConfig() const override;
 
   Expected<const COFFConfig &> getCOFFConfig() const override;
 

llvm/lib/ObjCopy/ConfigManager.cpp

@@ -13,6 +13,13 @@
 using namespace llvm;
 using namespace llvm::objcopy;
 
+Expected<const ELFConfig &> ConfigManager::getELFConfig() const {
+  if (!Common.ExtractSection.empty())
+    return createStringError(llvm::errc::invalid_argument,
+                             "option is not supported for ELF");
+  return ELF;
+}
+
 Expected<const COFFConfig &> ConfigManager::getCOFFConfig() const {
   if (!Common.SplitDWO.empty() || !Common.SymbolsPrefix.empty() ||
       !Common.SymbolsPrefixRemove.empty() || !Common.SymbolsToSkip.empty() ||
@@ -27,7 +34,7 @@ Expected<const COFFConfig &> ConfigManager::getCOFFConfig() const {
       Common.DiscardMode == DiscardType::Locals ||
       !Common.SymbolsToAdd.empty() || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty())
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for COFF");
 
@@ -48,7 +55,7 @@ Expected<const MachOConfig &> ConfigManager::getMachOConfig() const {
       Common.DiscardMode == DiscardType::Locals ||
       !Common.SymbolsToAdd.empty() || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty())
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for MachO");
 
@@ -69,7 +76,7 @@ Expected<const WasmConfig &> ConfigManager::getWasmConfig() const {
       !Common.SetSectionFlags.empty() || !Common.SetSectionType.empty() ||
       !Common.SymbolsToRename.empty() || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty())
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "only flags for section dumping, removal, and "
                              "addition are supported");
@@ -99,7 +106,7 @@ Expected<const XCOFFConfig &> ConfigManager::getXCOFFConfig() const {
       Common.Weaken || Common.StripUnneeded || Common.DecompressDebugSections ||
       Common.GapFill != 0 || Common.PadTo != 0 ||
       Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty()) {
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty()) {
     return createStringError(
         llvm::errc::invalid_argument,
         "no flags are supported yet, only basic copying is allowed");
@@ -125,9 +132,8 @@ ConfigManager::getDXContainerConfig() const {
       Common.GapFill != 0 || Common.PadTo != 0 ||
       Common.ChangeSectionLMAValAll != 0 ||
       !Common.ChangeSectionAddress.empty()) {
-    return createStringError(
-        llvm::errc::invalid_argument,
-        "no flags are supported yet, only basic copying is allowed");
+    return createStringError(llvm::errc::invalid_argument,
+                             "option is not supported for DXContainer");
   }
   return DXContainer;
 }

llvm/lib/ObjCopy/DXContainer/DXContainerObjcopy.cpp

@@ -11,13 +11,37 @@
 #include "DXContainerWriter.h"
 #include "llvm/ObjCopy/CommonConfig.h"
 #include "llvm/ObjCopy/DXContainer/DXContainerConfig.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
 namespace objcopy {
 namespace dxbc {
 
 using namespace object;
 
+static Error extractPartAsObject(StringRef PartName, StringRef OutFilename,
+                                 StringRef InputFilename, const Object &Obj) {
+  for (const Part &P : Obj.Parts)
+    if (P.Name == PartName) {
+      Object PartObj;
+      PartObj.Header = Obj.Header;
+      PartObj.Parts.push_back({P.Name, P.Data});
+      PartObj.recomputeHeader();
+
+      auto Write = [&OutFilename, &PartObj](raw_ostream &Out) -> Error {
+        DXContainerWriter Writer(PartObj, Out);
+        if (Error E = Writer.write())
+          return createFileError(OutFilename, std::move(E));
+        return Error::success();
+      };
+
+      return writeToOutput(OutFilename, Write);
+    }
+
+  return createFileError(InputFilename, object_error::parse_failed,
+                         "part '%s' not found", PartName.str().c_str());
+}
+
 static Error handleArgs(const CommonConfig &Config, Object &Obj) {
   std::function<bool(const Part &)> RemovePred = [](const Part &) {
     return false;
@@ -28,6 +52,15 @@ static Error handleArgs(const CommonConfig &Config, Object &Obj) {
       return Config.ToRemove.matches(P.Name);
     };
 
+  for (StringRef Flag : Config.ExtractSection) {
+    StringRef SectionName;
+    StringRef FileName;
+    std::tie(SectionName, FileName) = Flag.split('=');
+    if (Error E = extractPartAsObject(SectionName, FileName,
+                                      Config.InputFilename, Obj))
+      return E;
+  }
+
   if (auto E = Obj.removeParts(RemovePred))
     return E;