[DirectX] Add support for remove-section of DXContainer for llvm-objcopy (#153246)

This pr implements the `remove-section` option for a `DXContainer`
object in `llvm-objcopy`.

It implements a base `removeParts` to the minimal `object`
representation of a `DXContainerObject`.

This is the second step to implement
#150275 as a compiler actions
that invokes llvm-objcopy for functionality.

Author: inbelic

llvm/lib/ObjCopy/CMakeLists.txt

@@ -44,6 +44,7 @@ source_group("Source Files\\DXContainer" REGULAR_EXPRESSION
 add_llvm_component_library(LLVMObjCopy
   Archive.cpp
   DXContainer/DXContainerObjcopy.cpp
+  DXContainer/DXContainerObject.cpp
   DXContainer/DXContainerReader.cpp
   DXContainer/DXContainerWriter.cpp
   CommonConfig.cpp

llvm/lib/ObjCopy/ConfigManager.cpp

@@ -116,14 +116,14 @@ ConfigManager::getDXContainerConfig() const {
       !Common.AllocSectionsPrefix.empty() ||
       Common.DiscardMode != DiscardType::None || !Common.AddSection.empty() ||
       !Common.DumpSection.empty() || !Common.KeepSection.empty() ||
-      !Common.OnlySection.empty() || !Common.ToRemove.empty() ||
-      !Common.SectionsToRename.empty() || !Common.SetSectionAlignment.empty() ||
-      !Common.SetSectionFlags.empty() || !Common.SetSectionType.empty() ||
-      Common.ExtractDWO || Common.OnlyKeepDebug || Common.StripAllGNU ||
-      Common.StripDWO || Common.StripDebug || Common.StripNonAlloc ||
-      Common.StripSections || Common.StripUnneeded ||
-      Common.DecompressDebugSections || Common.GapFill != 0 ||
-      Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
+      !Common.OnlySection.empty() || !Common.SectionsToRename.empty() ||
+      !Common.SetSectionAlignment.empty() || !Common.SetSectionFlags.empty() ||
+      !Common.SetSectionType.empty() || Common.ExtractDWO ||
+      Common.OnlyKeepDebug || Common.StripAllGNU || Common.StripDWO ||
+      Common.StripDebug || Common.StripNonAlloc || Common.StripSections ||
+      Common.StripUnneeded || Common.DecompressDebugSections ||
+      Common.GapFill != 0 || Common.PadTo != 0 ||
+      Common.ChangeSectionLMAValAll != 0 ||
       !Common.ChangeSectionAddress.empty()) {
     return createStringError(
         llvm::errc::invalid_argument,

llvm/lib/ObjCopy/DXContainer/DXContainerObjcopy.cpp

@@ -19,6 +19,19 @@ namespace dxbc {
 using namespace object;
 
 static Error handleArgs(const CommonConfig &Config, Object &Obj) {
+  std::function<bool(const Part &)> RemovePred = [](const Part &) {
+    return false;
+  };
+
+  if (!Config.ToRemove.empty())
+    RemovePred = [&Config](const Part &P) {
+      return Config.ToRemove.matches(P.Name);
+    };
+
+  if (auto E = Obj.removeParts(RemovePred))
+    return E;
+
+  Obj.recomputeHeader();
   return Error::success();
 }
 

llvm/lib/ObjCopy/DXContainer/DXContainerObject.cpp

@@ -0,0 +1,29 @@
+//===- DXContainerObject.cpp ----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "DXContainerObject.h"
+
+namespace llvm {
+namespace objcopy {
+namespace dxbc {
+
+Error Object::removeParts(PartPred ToRemove) {
+  erase_if(Parts, ToRemove);
+  return Error::success();
+}
+
+void Object::recomputeHeader() {
+  Header.FileSize = headerSize();
+  Header.PartCount = Parts.size();
+  for (const Part &P : Parts)
+    Header.FileSize += P.size();
+}
+
+} // end namespace dxbc
+} // end namespace objcopy
+} // end namespace llvm

llvm/lib/ObjCopy/DXContainer/DXContainerObject.h

@@ -30,6 +30,8 @@ struct Part {
   }
 };
 
+using PartPred = llvm::function_ref<bool(const Part &)>;
+
 struct Object {
   ::llvm::dxbc::Header Header;
   SmallVector<Part> Parts;
@@ -38,6 +40,9 @@ struct Object {
     return sizeof(::llvm::dxbc::Header)       // base header
            + sizeof(uint32_t) * Parts.size(); // part offset values
   }
+
+  Error removeParts(PartPred ToRemove);
+  void recomputeHeader();
 };
 
 } // end namespace dxbc

llvm/test/tools/llvm-objcopy/DXContainer/remove-headers.test

@@ -0,0 +1,51 @@
+## Tests that the copied DXContainer correctly has the specified headers
+## removed.
+
+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy --remove-section=FKE1 --remove-section=FKE4 %t %t.out
+# RUN: obj2yaml %t.out | FileCheck %s
+
+--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+## FileSize = 1996 - 8 (FKE1 content) - 1688 (FKE4 content)
+##          - 8 (2 part offsets) - 16 (2 part headers)
+##          = 276
+# CHECK: FileSize:       276
+  FileSize:        1996
+# CHECK-NEXT: PartCount:      5
+  PartCount:       7
+# CHECK-NEXT: PartOffsets:    [ 52, 68, 84, 212, 240 ]
+  PartOffsets:     [ 60, 76, 92, 108, 236, 1932, 1960 ]
+Parts:
+# CHECK-NEXT: Parts
+# CHECK-NOT: FKE1
+# CHECK-NOT: FKE4
+  - Name:            FKE0
+    Size:            8
+  - Name:            FKE1
+    Size:            8
+  - Name:            FKE2
+    Size:            8
+  - Name:            FKE3
+    Size:            120
+  - Name:            FKE4
+    Size:            1688
+  - Name:            FKE5
+    Size:            20
+  - Name:            DXIL
+    Size:            28
+    Program:
+      MajorVersion:    6
+      MinorVersion:    5
+      ShaderKind:      5
+      Size:            8
+      DXILMajorVersion: 1
+      DXILMinorVersion: 5
+      DXILSize:        4
+      DXIL:            [ 0x42, 0x43, 0xC0, 0xDE, ]
+...