3 files changed, 191 insertions, 3 deletions

diff --git a/include/llvm/Support/ValueHandle.h b/include/llvm/Support/ValueHandle.h
index d003b02..f9b30fc 100644
--- a/include/llvm/Support/ValueHandle.h
+++ b/include/llvm/Support/ValueHandle.h
@@ -41,7 +41,7 @@ class ValueHandleBase {
 protected:
   /// HandleBaseKind - This indicates what sub class the handle actually is.
   /// This is to avoid having a vtable for the light-weight handle pointers. The
-  /// fully generally Callback version does have a vtable.
+  /// fully general Callback version does have a vtable.
   enum HandleBaseKind {
     Assert,
     Weak,
@@ -187,6 +187,50 @@ public:
   ValueTy &operator*() const { return *getValPtr(); }
 };
 
+/// CallbackVH - This is a value handle that allows subclasses to define
+/// callbacks that run when the underlying Value has RAUW called on it or is
+/// destroyed.  This class can be used as the key of a map, as long as the user
+/// takes it out of the map before calling setValPtr() (since the map has to
+/// rearrange itself when the pointer changes).  Unlike ValueHandleBase, this
+/// class has a vtable and a virtual destructor.
+class CallbackVH : public ValueHandleBase {
+protected:
+  CallbackVH(const CallbackVH &RHS)
+    : ValueHandleBase(Callback, RHS) {}
+
+  virtual ~CallbackVH();
+
+  void setValPtr(Value *P) {
+    ValueHandleBase::operator=(P);
+  }
+
+public:
+  CallbackVH() : ValueHandleBase(Callback) {}
+  CallbackVH(Value *P) : ValueHandleBase(Callback, P) {}
+
+  operator Value*() const {
+    return getValPtr();
+  }
+
+  /// Called when this->getValPtr() is destroyed, inside ~Value(), so you may
+  /// call any non-virtual Value method on getValPtr(), but no subclass methods.
+  /// If WeakVH were implemented as a CallbackVH, it would use this method to
+  /// call setValPtr(NULL).  AssertingVH would use this method to cause an
+  /// assertion failure.
+  ///
+  /// All implementations must remove the reference from this object to the
+  /// Value that's being destroyed.
+  virtual void deleted() {
+    setValPtr(NULL);
+  }
+
+  /// Called when this->getValPtr()->replaceAllUsesWith(new_value) is called,
+  /// _before_ any of the uses have actually been replaced.  If WeakVH were
+  /// implemented as a CallbackVH, it would use this method to call
+  /// setValPtr(new_value).  AssertingVH would do nothing in this method.
+  virtual void allUsesReplacedWith(Value *new_value) {}
+};
+
 } // End llvm namespace
 
 #endif
diff --git a/lib/VMCore/Value.cpp b/lib/VMCore/Value.cpp
index af3f4e7..35c8ccf 100644
--- a/lib/VMCore/Value.cpp
+++ b/lib/VMCore/Value.cpp
@@ -511,7 +511,9 @@ void ValueHandleBase::ValueIsDeleted(Value *V) {
       ThisNode->operator=(0);
       break;
     case Callback:
-      assert(0 && "Callback not implemented yet!");
+      // Forward to the subclass's implementation.
+      static_cast<CallbackVH*>(ThisNode)->deleted();
+      break;
     }
   }
   
@@ -543,11 +545,17 @@ void ValueHandleBase::ValueIsRAUWd(Value *Old, Value *New) {
       ThisNode->operator=(New);
       break;
     case Callback:
-      assert(0 && "Callback not implemented yet!");
+      // Forward to the subclass's implementation.
+      static_cast<CallbackVH*>(ThisNode)->allUsesReplacedWith(New);
+      break;
     }
   }
 }
 
+/// ~CallbackVH. Empty, but defined here to avoid emitting the vtable
+/// more than once.
+CallbackVH::~CallbackVH() {}
+
 
 //===----------------------------------------------------------------------===//
 //                                 User Class
diff --git a/unittests/Support/ValueHandleTest.cpp b/unittests/Support/ValueHandleTest.cpp
index 7980758..89c8250 100644
--- a/unittests/Support/ValueHandleTest.cpp
+++ b/unittests/Support/ValueHandleTest.cpp
@@ -30,6 +30,12 @@ protected:
   }
 };
 
+class ConcreteCallbackVH : public CallbackVH {
+public:
+  ConcreteCallbackVH() : CallbackVH() {}
+  ConcreteCallbackVH(Value *V) : CallbackVH(V) {}
+};
+
 TEST_F(ValueHandle, WeakVH_BasicOperation) {
   WeakVH WVH(BitcastV.get());
   EXPECT_EQ(BitcastV.get(), WVH);
@@ -178,4 +184,134 @@ TEST_F(ValueHandle, AssertingVH_Asserts) {
 
 #endif  // NDEBUG
 
+TEST_F(ValueHandle, CallbackVH_BasicOperation) {
+  ConcreteCallbackVH CVH(BitcastV.get());
+  EXPECT_EQ(BitcastV.get(), CVH);
+  CVH = ConstantV;
+  EXPECT_EQ(ConstantV, CVH);
+
+  // Make sure I can call a method on the underlying Value.  It
+  // doesn't matter which method.
+  EXPECT_EQ(Type::Int32Ty, CVH->getType());
+  EXPECT_EQ(Type::Int32Ty, (*CVH).getType());
+}
+
+TEST_F(ValueHandle, CallbackVH_Comparisons) {
+  ConcreteCallbackVH BitcastCVH(BitcastV.get());
+  ConcreteCallbackVH ConstantCVH(ConstantV);
+
+  EXPECT_TRUE(BitcastCVH == BitcastCVH);
+  EXPECT_TRUE(BitcastV.get() == BitcastCVH);
+  EXPECT_TRUE(BitcastCVH == BitcastV.get());
+  EXPECT_FALSE(BitcastCVH == ConstantCVH);
+
+  EXPECT_TRUE(BitcastCVH != ConstantCVH);
+  EXPECT_TRUE(BitcastV.get() != ConstantCVH);
+  EXPECT_TRUE(BitcastCVH != ConstantV);
+  EXPECT_FALSE(BitcastCVH != BitcastCVH);
+
+  // Cast to Value* so comparisons work.
+  Value *BV = BitcastV.get();
+  Value *CV = ConstantV;
+  EXPECT_EQ(BV < CV, BitcastCVH < ConstantCVH);
+  EXPECT_EQ(BV <= CV, BitcastCVH <= ConstantCVH);
+  EXPECT_EQ(BV > CV, BitcastCVH > ConstantCVH);
+  EXPECT_EQ(BV >= CV, BitcastCVH >= ConstantCVH);
+
+  EXPECT_EQ(BV < CV, BitcastV.get() < ConstantCVH);
+  EXPECT_EQ(BV <= CV, BitcastV.get() <= ConstantCVH);
+  EXPECT_EQ(BV > CV, BitcastV.get() > ConstantCVH);
+  EXPECT_EQ(BV >= CV, BitcastV.get() >= ConstantCVH);
+
+  EXPECT_EQ(BV < CV, BitcastCVH < ConstantV);
+  EXPECT_EQ(BV <= CV, BitcastCVH <= ConstantV);
+  EXPECT_EQ(BV > CV, BitcastCVH > ConstantV);
+  EXPECT_EQ(BV >= CV, BitcastCVH >= ConstantV);
+}
+
+TEST_F(ValueHandle, CallbackVH_CallbackOnDeletion) {
+  class RecordingVH : public CallbackVH {
+  public:
+    int DeletedCalls;
+    int AURWCalls;
+
+    RecordingVH() : DeletedCalls(0), AURWCalls(0) {}
+    RecordingVH(Value *V) : CallbackVH(V), DeletedCalls(0), AURWCalls(0) {}
+
+  private:
+    virtual void deleted() { DeletedCalls++; CallbackVH::deleted(); }
+    virtual void allUsesReplacedWith(Value *) { AURWCalls++; }
+  };
+
+  RecordingVH RVH;
+  RVH = BitcastV.get();
+  EXPECT_EQ(0, RVH.DeletedCalls);
+  EXPECT_EQ(0, RVH.AURWCalls);
+  BitcastV.reset();
+  EXPECT_EQ(1, RVH.DeletedCalls);
+  EXPECT_EQ(0, RVH.AURWCalls);
+}
+
+TEST_F(ValueHandle, CallbackVH_CallbackOnRAUW) {
+  class RecordingVH : public CallbackVH {
+  public:
+    int DeletedCalls;
+    Value *AURWArgument;
+
+    RecordingVH() : DeletedCalls(0), AURWArgument(NULL) {}
+    RecordingVH(Value *V)
+      : CallbackVH(V), DeletedCalls(0), AURWArgument(NULL) {}
+
+  private:
+    virtual void deleted() { DeletedCalls++; CallbackVH::deleted(); }
+    virtual void allUsesReplacedWith(Value *new_value) {
+      EXPECT_EQ(NULL, AURWArgument);
+      AURWArgument = new_value;
+    }
+  };
+
+  RecordingVH RVH;
+  RVH = BitcastV.get();
+  EXPECT_EQ(0, RVH.DeletedCalls);
+  EXPECT_EQ(NULL, RVH.AURWArgument);
+  BitcastV->replaceAllUsesWith(ConstantV);
+  EXPECT_EQ(0, RVH.DeletedCalls);
+  EXPECT_EQ(ConstantV, RVH.AURWArgument);
+}
+
+TEST_F(ValueHandle, CallbackVH_DeletionCanRAUW) {
+  class RecoveringVH : public CallbackVH {
+  public:
+    int DeletedCalls;
+    Value *AURWArgument;
+
+    RecoveringVH() : DeletedCalls(0), AURWArgument(NULL) {}
+    RecoveringVH(Value *V)
+      : CallbackVH(V), DeletedCalls(0), AURWArgument(NULL) {}
+
+  private:
+    virtual void deleted() {
+      getValPtr()->replaceAllUsesWith(Constant::getNullValue(Type::Int32Ty));
+      setValPtr(NULL);
+    }
+    virtual void allUsesReplacedWith(Value *new_value) {
+      ASSERT_TRUE(NULL != getValPtr());
+      EXPECT_EQ(1, getValPtr()->getNumUses());
+      EXPECT_EQ(NULL, AURWArgument);
+      AURWArgument = new_value;
+    }
+  };
+
+  // Normally, if a value has uses, deleting it will crash.  However, we can use
+  // a CallbackVH to remove the uses before the check for no uses.
+  RecoveringVH RVH;
+  RVH = BitcastV.get();
+  std::auto_ptr<BinaryOperator> BitcastUser(
+    BinaryOperator::CreateAdd(RVH, Constant::getNullValue(Type::Int32Ty)));
+  EXPECT_EQ(BitcastV.get(), BitcastUser->getOperand(0));
+  BitcastV.reset();  // Would crash without the ValueHandler.
+  EXPECT_EQ(Constant::getNullValue(Type::Int32Ty), RVH.AURWArgument);
+  EXPECT_EQ(Constant::getNullValue(Type::Int32Ty), BitcastUser->getOperand(0));
+}
+
 }