5 files changed, 23 insertions, 5 deletions
diff --git a/include/llvm/Analysis/DSGraph.h b/include/llvm/Analysis/DSGraph.h
index 568a90bc60..1c10851848 100644
--- a/include/llvm/Analysis/DSGraph.h
+++ b/include/llvm/Analysis/DSGraph.h
@@ -182,6 +182,7 @@ public:
   /// addNode - Add a new node to the graph.
   ///
   void addNode(DSNode *N) { Nodes.push_back(N); }
+  void unlinkNode(DSNode *N) { Nodes.remove(N); }
 
   /// getScalarMap - Get a map that describes what the nodes the scalars in this
   /// function point to...
diff --git a/include/llvm/Analysis/DSNode.h b/include/llvm/Analysis/DSNode.h
index e76a48090f..611293132b 100644
--- a/include/llvm/Analysis/DSNode.h
+++ b/include/llvm/Analysis/DSNode.h
@@ -159,12 +159,18 @@ public:
   DSNode *getForwardNode() const { return ForwardNH.getNode(); }
 
   /// isForwarding - Return true if this node is forwarding to another.
+  ///
   bool isForwarding() const { return !ForwardNH.isNull(); }
 
+  /// stopForwarding - When the last reference to this forwarding node has been
+  /// dropped, delete the node.
   void stopForwarding() {
     assert(isForwarding() &&
            "Node isn't forwarding, cannot stopForwarding!");
     ForwardNH.setNode(0);
+    assert(ParentGraph == 0 &&
+           "Forwarding nodes must have been removed from graph!");
+    delete this;
   }
 
   /// hasLink - Return true if this memory object has a link in slot #LinkNo
@@ -377,8 +383,8 @@ inline DSNode *DSNodeHandle::getNode() const {
 }
 
 inline void DSNodeHandle::setNode(DSNode *n) const {
-  assert(!n || !n->getForwardNode() && "Cannot set node to a forwarded node!");
-  if (N) N->NumReferrers--;
+  assert(!n || !n->isForwarding() && "Cannot set node to a forwarded node!");
+  if (N) getNode()->NumReferrers--;
   N = n;
   if (N) {
     N->NumReferrers++;
diff --git a/include/llvm/Analysis/DataStructure/DSGraph.h b/include/llvm/Analysis/DataStructure/DSGraph.h
index 568a90bc60..1c10851848 100644
--- a/include/llvm/Analysis/DataStructure/DSGraph.h
+++ b/include/llvm/Analysis/DataStructure/DSGraph.h
@@ -182,6 +182,7 @@ public:
   /// addNode - Add a new node to the graph.
   ///
   void addNode(DSNode *N) { Nodes.push_back(N); }
+  void unlinkNode(DSNode *N) { Nodes.remove(N); }
 
   /// getScalarMap - Get a map that describes what the nodes the scalars in this
   /// function point to...
diff --git a/include/llvm/Analysis/DataStructure/DSNode.h b/include/llvm/Analysis/DataStructure/DSNode.h
index e76a48090f..611293132b 100644
--- a/include/llvm/Analysis/DataStructure/DSNode.h
+++ b/include/llvm/Analysis/DataStructure/DSNode.h
@@ -159,12 +159,18 @@ public:
   DSNode *getForwardNode() const { return ForwardNH.getNode(); }
 
   /// isForwarding - Return true if this node is forwarding to another.
+  ///
   bool isForwarding() const { return !ForwardNH.isNull(); }
 
+  /// stopForwarding - When the last reference to this forwarding node has been
+  /// dropped, delete the node.
   void stopForwarding() {
     assert(isForwarding() &&
            "Node isn't forwarding, cannot stopForwarding!");
     ForwardNH.setNode(0);
+    assert(ParentGraph == 0 &&
+           "Forwarding nodes must have been removed from graph!");
+    delete this;
   }
 
   /// hasLink - Return true if this memory object has a link in slot #LinkNo
@@ -377,8 +383,8 @@ inline DSNode *DSNodeHandle::getNode() const {
 }
 
 inline void DSNodeHandle::setNode(DSNode *n) const {
-  assert(!n || !n->getForwardNode() && "Cannot set node to a forwarded node!");
-  if (N) N->NumReferrers--;
+  assert(!n || !n->isForwarding() && "Cannot set node to a forwarded node!");
+  if (N) getNode()->NumReferrers--;
   N = n;
   if (N) {
     N->NumReferrers++;
diff --git a/lib/Analysis/DataStructure/DataStructure.cpp b/lib/Analysis/DataStructure/DataStructure.cpp
index 57651d78e1..68530f043e 100644
--- a/lib/Analysis/DataStructure/DataStructure.cpp
+++ b/lib/Analysis/DataStructure/DataStructure.cpp
@@ -46,7 +46,7 @@ namespace {
 using namespace DS;
 
 DSNode *DSNodeHandle::HandleForwarding() const {
-  assert(!N->ForwardNH.isNull() && "Can only be invoked if forwarding!");
+  assert(N->isForwarding() && "Can only be invoked if forwarding!");
 
   // Handle node forwarding here!
   DSNode *Next = N->ForwardNH.getNode();  // Cause recursive shrinkage
@@ -124,6 +124,10 @@ void DSNode::forwardNode(DSNode *To, unsigned Offset) {
   NodeType = DEAD;
   Size = 0;
   Ty = Type::VoidTy;
+
+  // Remove this node from the parent graph's Nodes list.
+  ParentGraph->unlinkNode(this);  
+  ParentGraph = 0;
 }
 
 // addGlobal - Add an entry for a global value to the Globals list.  This also