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

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> SplitSection;
 
   // Section matchers
   NameMatcher KeepSection;

llvm/lib/ObjCopy/ConfigManager.cpp

@@ -27,7 +27,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.SplitSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for COFF");
 
@@ -48,7 +48,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.SplitSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for MachO");
 
@@ -69,7 +69,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.SplitSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "only flags for section dumping, removal, and "
                              "addition are supported");
@@ -99,7 +99,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.SplitSection.empty()) {
     return createStringError(
         llvm::errc::invalid_argument,
         "no flags are supported yet, only basic copying is allowed");
@@ -126,12 +126,11 @@ 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 directx");
   }
 
   // If a flag is listed here, then it has support for DXContainer:
-  // Common.PreserveDates, Common.ToRemove
+  // Common.PreserveDates, Common.ToRemove, Common.SplitSection
   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 splitPartAsObject(StringRef PartName, StringRef OutFilename,
+                               StringRef InputFilename, const Object &Obj) {
+  for (const Part &P : Obj.Parts)
+    if (P.Name == PartName) {
+      auto PartObj = std::make_unique<Object>();
+      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,32 @@ static Error handleArgs(const CommonConfig &Config, Object &Obj) {
       return Config.ToRemove.matches(P.Name);
     };
 
+  if (!Config.SplitSection.empty()) {
+    NameMatcher SectionMatches;
+    for (StringRef Flag : Config.SplitSection) {
+      StringRef SectionName;
+      StringRef FileName;
+      std::tie(SectionName, FileName) = Flag.split('=');
+
+      if (Error E = splitPartAsObject(SectionName, FileName,
+                                      Config.InputFilename, Obj))
+        return E;
+    }
+
+    RemovePred = [&Config, RemovePred](const Part &P) {
+      if (RemovePred(P))
+        return true;
+
+      for (StringRef Flag : Config.SplitSection) {
+        StringRef SectionName = Flag.take_front(Flag.find('='));
+        if (P.Name == SectionName)
+          return true;
+      }
+
+      return false;
+    };
+  }
+
   if (auto E = Obj.removeParts(RemovePred))
     return E;
 

llvm/lib/ObjCopy/DXContainer/DXContainerObject.h

@@ -42,8 +42,6 @@ struct Object {
   }
 
   Error removeParts(PartPred ToRemove);
-
-private:
   void recomputeHeader();
 };