[DirectX] Add suport to remove specified sections (parts)

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/DXContainer/DXContainerObjcopy.cpp

@@ -19,6 +19,18 @@ 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;
+
   return Error::success();
 }
 

llvm/lib/ObjCopy/DXContainer/DXContainerObject.cpp

@@ -0,0 +1,30 @@
+//===- 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);
+  recomputeHeader();
+  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,11 @@ struct Object {
     return sizeof(::llvm::dxbc::Header)       // base header
            + sizeof(uint32_t) * Parts.size(); // part offset values
   }
+
+  Error removeParts(PartPred ToRemove);
+
+private:
+  void recomputeHeader();
 };
 
 } // end namespace dxbc