1 files changed, 56 insertions, 20 deletions

diff --git a/llvm/lib/Transforms/Utils/ModuleUtils.cpp b/llvm/lib/Transforms/Utils/ModuleUtils.cpp
index fbdc3fd..a1e1370 100644
--- a/llvm/lib/Transforms/Utils/ModuleUtils.cpp
+++ b/llvm/lib/Transforms/Utils/ModuleUtils.cpp
@@ -74,35 +74,35 @@ void llvm::appendToGlobalDtors(Module &M, Function *F, int Priority, Constant *D
   appendToGlobalArray("llvm.global_dtors", M, F, Priority, Data);
 }
 
+static void collectUsedGlobals(GlobalVariable *GV,
+                               SmallSetVector<Constant *, 16> &Init) {
+  if (!GV || !GV->hasInitializer())
+    return;
+
+  auto *CA = cast<ConstantArray>(GV->getInitializer());
+  for (Use &Op : CA->operands())
+    Init.insert(cast<Constant>(Op));
+}
+
 static void appendToUsedList(Module &M, StringRef Name, ArrayRef<GlobalValue *> Values) {
   GlobalVariable *GV = M.getGlobalVariable(Name);
-  SmallPtrSet<Constant *, 16> InitAsSet;
-  SmallVector<Constant *, 16> Init;
-  if (GV) {
-    if (GV->hasInitializer()) {
-      auto *CA = cast<ConstantArray>(GV->getInitializer());
-      for (auto &Op : CA->operands()) {
-        Constant *C = cast_or_null<Constant>(Op);
-        if (InitAsSet.insert(C).second)
-          Init.push_back(C);
-      }
-    }
+
+  SmallSetVector<Constant *, 16> Init;
+  collectUsedGlobals(GV, Init);
+  if (GV)
     GV->eraseFromParent();
-  }
 
-  Type *Int8PtrTy = llvm::Type::getInt8PtrTy(M.getContext());
-  for (auto *V : Values) {
-    Constant *C = ConstantExpr::getPointerBitCastOrAddrSpaceCast(V, Int8PtrTy);
-    if (InitAsSet.insert(C).second)
-      Init.push_back(C);
-  }
+  Type *ArrayEltTy = llvm::Type::getInt8PtrTy(M.getContext());
+  for (auto *V : Values)
+    Init.insert(ConstantExpr::getPointerBitCastOrAddrSpaceCast(V, ArrayEltTy));
 
   if (Init.empty())
     return;
 
-  ArrayType *ATy = ArrayType::get(Int8PtrTy, Init.size());
+  ArrayType *ATy = ArrayType::get(ArrayEltTy, Init.size());
   GV = new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage,
-                                ConstantArray::get(ATy, Init), Name);
+                                ConstantArray::get(ATy, Init.getArrayRef()),
+                                Name);
   GV->setSection("llvm.metadata");
 }
 
@@ -114,6 +114,42 @@ void llvm::appendToCompilerUsed(Module &M, ArrayRef<GlobalValue *> Values) {
   appendToUsedList(M, "llvm.compiler.used", Values);
 }
 
+static void removeFromUsedList(Module &M, StringRef Name,
+                               function_ref<bool(Constant *)> ShouldRemove) {
+  GlobalVariable *GV = M.getNamedGlobal(Name);
+  if (!GV)
+    return;
+
+  SmallSetVector<Constant *, 16> Init;
+  collectUsedGlobals(GV, Init);
+
+  Type *ArrayEltTy = cast<ArrayType>(GV->getValueType())->getElementType();
+
+  SmallVector<Constant *, 16> NewInit;
+  for (Constant *MaybeRemoved : Init) {
+    if (!ShouldRemove(MaybeRemoved->stripPointerCasts()))
+      NewInit.push_back(MaybeRemoved);
+  }
+
+  if (!NewInit.empty()) {
+    ArrayType *ATy = ArrayType::get(ArrayEltTy, NewInit.size());
+    GlobalVariable *NewGV =
+        new GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage,
+                           ConstantArray::get(ATy, NewInit), "", GV,
+                           GV->getThreadLocalMode(), GV->getAddressSpace());
+    NewGV->setSection(GV->getSection());
+    NewGV->takeName(GV);
+  }
+
+  GV->eraseFromParent();
+}
+
+void llvm::removeFromUsedLists(Module &M,
+                               function_ref<bool(Constant *)> ShouldRemove) {
+  removeFromUsedList(M, "llvm.used", ShouldRemove);
+  removeFromUsedList(M, "llvm.compiler.used", ShouldRemove);
+}
+
 static void setKCFIType(Module &M, Function &F, StringRef MangledType) {
   if (!M.getModuleFlag("kcfi"))
     return;