[PPCMergeStringPool] Only replace constant once (llvm#92996)

nikic · web-flow · commit 9f85bc834b07 · 2024-05-27T08:54:11.000+02:00
In llvm#88846 I changed this code to use RAUW to perform the replacement instead of manual updates -- but kept the outer loop, which means we try to perform RAUW once per user. However, some of the users might be freed by the RAUW operation, resulting in use-after-free. The case where this happens is constant users where the replacement might result in the destruction of the original constant. Fixes llvm#92991.
diff --git a/llvm/lib/Target/PowerPC/PPCMergeStringPool.cpp b/llvm/lib/Target/PowerPC/PPCMergeStringPool.cpp
@@ -302,13 +302,6 @@ bool PPCMergeStringPool::mergeModuleStringPool(Module &M) {
   return true;
 }
 
-static bool userHasOperand(User *TheUser, GlobalVariable *GVOperand) {
-  for (Value *Op : TheUser->operands())
-    if (Op == GVOperand)
-      return true;
-  return false;
-}
-
 // For pooled strings we need to add the offset into the pool for each string.
 // This is done by adding a Get Element Pointer (GEP) before each user. This
 // function adds the GEP.
@@ -319,29 +312,13 @@ void PPCMergeStringPool::replaceUsesWithGEP(GlobalVariable *GlobalToReplace,
   Indices.push_back(ConstantInt::get(Type::getInt32Ty(*Context), 0));
   Indices.push_back(ConstantInt::get(Type::getInt32Ty(*Context), ElementIndex));
 
-  // Need to save a temporary copy of each user list because we remove uses
-  // as we replace them.
-  SmallVector<User *> Users;
-  for (User *CurrentUser : GlobalToReplace->users())
-    Users.push_back(CurrentUser);
-
-  for (User *CurrentUser : Users) {
-    // The user was not found so it must have been replaced earlier.
-    if (!userHasOperand(CurrentUser, GlobalToReplace))
-      continue;
-
-    // We cannot replace operands in globals so we ignore those.
-    if (isa<GlobalValue>(CurrentUser))
-      continue;
-
-    Constant *ConstGEP = ConstantExpr::getInBoundsGetElementPtr(
-        PooledStructType, GPool, Indices);
-    LLVM_DEBUG(dbgs() << "Replacing this global:\n");
-    LLVM_DEBUG(GlobalToReplace->dump());
-    LLVM_DEBUG(dbgs() << "with this:\n");
-    LLVM_DEBUG(ConstGEP->dump());
-    GlobalToReplace->replaceAllUsesWith(ConstGEP);
-  }
+  Constant *ConstGEP =
+      ConstantExpr::getInBoundsGetElementPtr(PooledStructType, GPool, Indices);
+  LLVM_DEBUG(dbgs() << "Replacing this global:\n");
+  LLVM_DEBUG(GlobalToReplace->dump());
+  LLVM_DEBUG(dbgs() << "with this:\n");
+  LLVM_DEBUG(ConstGEP->dump());
+  GlobalToReplace->replaceAllUsesWith(ConstGEP);
 }
 
 } // namespace
diff --git a/llvm/test/CodeGen/PowerPC/mergeable-string-pool-pr92991.ll b/llvm/test/CodeGen/PowerPC/mergeable-string-pool-pr92991.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=powerpc64le-unknown-linux-gnu < %s | FileCheck %s
+
+@g = private constant [4 x i32] [i32 122, i32 67, i32 35, i32 56]
+@g2 = private constant [1 x i64] [i64 1], align 8
+
+define void @test(ptr %p, ptr %p2) {
+; CHECK-LABEL: test:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addis 5, 2, .L__ModuleStringPool@toc@ha
+; CHECK-NEXT:    addi 5, 5, .L__ModuleStringPool@toc@l
+; CHECK-NEXT:    addi 6, 5, 12
+; CHECK-NEXT:    std 6, 0(3)
+; CHECK-NEXT:    addi 3, 5, 16
+; CHECK-NEXT:    std 3, 0(4)
+; CHECK-NEXT:    blr
+  store ptr getelementptr inbounds ([4 x i32], ptr @g, i64 0, i64 1), ptr %p
+  store ptr getelementptr inbounds ([4 x i32], ptr @g, i64 0, i64 2), ptr %p2
+  ret void
+}
