[llvm] 1e46eb7 - [Attributor][FIX] Avoid dangling value pointers during code modification

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2020-01-08T19:32:37-06:00
New Revision: 1e46eb74be6527377e47090bbe0fc9298f7de2c5

URL: https://github.com/llvm/llvm-project/commit/1e46eb74be6527377e47090bbe0fc9298f7de2c5
DIFF: https://github.com/llvm/llvm-project/commit/1e46eb74be6527377e47090bbe0fc9298f7de2c5.diff

LOG: [Attributor][FIX] Avoid dangling value pointers during code modification

When we replace instructions with unreachable we delete instructions. We
now avoid dangling pointers to those deleted instructions in the
`ToBeChangedToUnreachableInsts` set. Other modification collections
might need to be updated in the future as well.

Added: 
    

Modified: 
    llvm/include/llvm/IR/ValueHandle.h
    llvm/include/llvm/Transforms/IPO/Attributor.h
    llvm/lib/Transforms/IPO/Attributor.cpp
    llvm/test/Transforms/Attributor/undefined_behavior.ll

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/ValueHandle.h b/llvm/include/llvm/IR/ValueHandle.h
index 11ac2a8608d8..50b7701f6716 100644
--- a/llvm/include/llvm/IR/ValueHandle.h
+++ b/llvm/include/llvm/IR/ValueHandle.h
@@ -171,6 +171,25 @@ template <> struct simplify_type<const WeakVH> {
   static SimpleType getSimplifiedValue(const WeakVH &WVH) { return WVH; }
 };
 
+// Specialize DenseMapInfo to allow WeakVH to participate in DenseMap.
+template <> struct DenseMapInfo<WeakVH> {
+  static inline WeakVH getEmptyKey() {
+    return WeakVH(DenseMapInfo<Value *>::getEmptyKey());
+  }
+
+  static inline WeakVH getTombstoneKey() {
+    return WeakVH(DenseMapInfo<Value *>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const WeakVH &Val) {
+    return DenseMapInfo<Value *>::getHashValue(Val);
+  }
+
+  static bool isEqual(const WeakVH &LHS, const WeakVH &RHS) {
+    return DenseMapInfo<Value *>::isEqual(LHS, RHS);
+  }
+};
+
 /// Value handle that is nullable, but tries to track the Value.
 ///
 /// This is a value handle that tries hard to point to a Value, even across

diff  --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index 53236b54ff0f..4edfec2e982d 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -1174,7 +1174,7 @@ struct Attributor {
   DenseMap<Use *, Value *> ToBeChangedUses;
 
   /// Instructions we replace with `unreachable` insts after manifest is done.
-  SmallPtrSet<Instruction *, 8> ToBeChangedToUnreachableInsts;
+  SmallDenseSet<WeakVH, 16> ToBeChangedToUnreachableInsts;
 
   /// Functions, blocks, and instructions we delete after manifest is done.
   ///

diff  --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 13fcf6aa7247..b3b9e1e185a6 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -5668,8 +5668,9 @@ ChangeStatus Attributor::run(Module &M) {
         }
       }
     }
-    for (Instruction *I : ToBeChangedToUnreachableInsts)
-      changeToUnreachable(I, /* UseLLVMTrap */ false);
+    for (auto &V : ToBeChangedToUnreachableInsts)
+      if (Instruction *I = dyn_cast_or_null<Instruction>(V))
+        changeToUnreachable(I, /* UseLLVMTrap */ false);
     for (Instruction *I : TerminatorsToFold)
       ConstantFoldTerminator(I->getParent());
 

diff  --git a/llvm/test/Transforms/Attributor/undefined_behavior.ll b/llvm/test/Transforms/Attributor/undefined_behavior.ll
index e9b782452182..fd0ddb1ebb85 100644
--- a/llvm/test/Transforms/Attributor/undefined_behavior.ll
+++ b/llvm/test/Transforms/Attributor/undefined_behavior.ll
@@ -16,6 +16,16 @@ define void @load_wholly_unreachable() {
   ret void
 }
 
+define void @loads_wholly_unreachable() {
+; ATTRIBUTOR-LABEL: @loads_wholly_unreachable(
+; ATTRIBUTOR-NEXT:    unreachable
+;
+  %a = load i32, i32* null
+  %b = load i32, i32* null
+  ret void
+}
+
+
 define void @load_single_bb_unreachable(i1 %cond) {
 ; ATTRIBUTOR-LABEL: @load_single_bb_unreachable(
 ; ATTRIBUTOR-NEXT:    br i1 [[COND:%.*]], label [[T:%.*]], label [[E:%.*]]