diff --git a/include/llvm/Analysis/DSGraphTraits.h b/include/llvm/Analysis/DSGraphTraits.h
index 4a99b738c2e..575c1e18306 100644
--- a/include/llvm/Analysis/DSGraphTraits.h
+++ b/include/llvm/Analysis/DSGraphTraits.h
@@ -12,8 +12,7 @@
 #include "llvm/Analysis/DSGraph.h"
 #include "Support/GraphTraits.h"
 #include "Support/iterator"
-
-#if 0
+#include "Support/STLExtras.h"
 
 class DSNodeIterator : public forward_iterator<DSNode, ptrdiff_t> {
   friend class DSNode;
@@ -22,13 +21,9 @@ class DSNodeIterator : public forward_iterator<DSNode, ptrdiff_t> {
   
   typedef DSNodeIterator _Self;
 
-  DSNodeIterator(DSNode *N) : Node(N), Link(0) {   // begin iterator
-    unsigned NumLinks = Node->getNumLinks();
-    while (Link < NumLinks && Node->getLink(Link) == 0)
-      ++Link;
-  }
+  DSNodeIterator(DSNode *N) : Node(N), Link(0) {}   // begin iterator
   DSNodeIterator(DSNode *N, bool)       // Create end iterator
-    : Node(N), Link(N->getNumLinks()) {
+    : Node(N), Link(N->getSize()) {
   }
 public:
 
@@ -36,38 +31,60 @@ public:
     return Link == x.Link;
   }
   bool operator!=(const _Self& x) const { return !operator==(x); }
+
+  const _Self &operator=(const _Self &I) {
+    assert(I.Node == Node && "Cannot assign iterators to two different nodes!");
+    Link = I.Link;
+    return *this;
+  }
   
   pointer operator*() const {
-    return Node->getLink(Link);
+    DSNodeHandle *NH = Node->getLink(Link);
+    return NH ? NH->getNode() : 0;
   }
   pointer operator->() const { return operator*(); }
   
   _Self& operator++() {                // Preincrement
-    unsigned NumLinks = Node->getNumLinks();
-    do {
-      ++Link;
-    } while (Link < NumLinks && Node->getLink(Link) != 0);
+    ++Link;
     return *this;
   }
   _Self operator++(int) { // Postincrement
     _Self tmp = *this; ++*this; return tmp; 
   }
-};
 
-
-template <> struct GraphTraits<DSNode*> {
-  typedef DSNode NodeType;
-  typedef DSNode::iterator ChildIteratorType;
-
-  static NodeType *getEntryNode(DSNode *N) { return N; }
-  static ChildIteratorType child_begin(NodeType *N) { return N->begin(); }
-  static ChildIteratorType child_end(NodeType *N) { return N->end(); }
+  unsigned getLink() const { return Link; }
+  DSNode *getNode() const { return Node; }
 };
 
 // Provide iterators for DSNode...
 inline DSNode::iterator DSNode::begin() { return DSNodeIterator(this); }
 inline DSNode::iterator DSNode::end()   { return DSNodeIterator(this, false); }
 
-#endif
+template <> struct GraphTraits<DSNode*> {
+  typedef DSNode NodeType;
+  typedef DSNode::iterator ChildIteratorType;
+
+  static NodeType *getEntryNode(NodeType *N) { return N; }
+  static ChildIteratorType child_begin(NodeType *N) { return N->begin(); }
+  static ChildIteratorType child_end(NodeType *N) { return N->end(); }
+};
+
+static DSNode &dereference(DSNode *N) { return *N; }
+
+template <> struct GraphTraits<DSGraph*> {
+  typedef DSNode NodeType;
+  typedef DSNode::iterator ChildIteratorType;
+
+  typedef std::pointer_to_unary_function<DSNode *, DSNode&> DerefFun;
+
+  // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
+  typedef mapped_iterator<std::vector<DSNode*>::iterator,
+                          DerefFun> nodes_iterator;
+  static nodes_iterator nodes_begin(DSGraph *G) { return map_iterator(G->getNodes().begin(), DerefFun(dereference));}
+  static nodes_iterator nodes_end  (DSGraph *G) { return map_iterator(G->getNodes().end(), DerefFun(dereference)); }
+
+  static ChildIteratorType child_begin(NodeType *N) { return N->begin(); }
+  static ChildIteratorType child_end(NodeType *N) { return N->end(); }
+};
 
 #endif
diff --git a/include/llvm/Analysis/DataStructure/DSGraphTraits.h b/include/llvm/Analysis/DataStructure/DSGraphTraits.h
index 4a99b738c2e..575c1e18306 100644
--- a/include/llvm/Analysis/DataStructure/DSGraphTraits.h
+++ b/include/llvm/Analysis/DataStructure/DSGraphTraits.h
@@ -12,8 +12,7 @@
 #include "llvm/Analysis/DSGraph.h"
 #include "Support/GraphTraits.h"
 #include "Support/iterator"
-
-#if 0
+#include "Support/STLExtras.h"
 
 class DSNodeIterator : public forward_iterator<DSNode, ptrdiff_t> {
   friend class DSNode;
@@ -22,13 +21,9 @@ class DSNodeIterator : public forward_iterator<DSNode, ptrdiff_t> {
   
   typedef DSNodeIterator _Self;
 
-  DSNodeIterator(DSNode *N) : Node(N), Link(0) {   // begin iterator
-    unsigned NumLinks = Node->getNumLinks();
-    while (Link < NumLinks && Node->getLink(Link) == 0)
-      ++Link;
-  }
+  DSNodeIterator(DSNode *N) : Node(N), Link(0) {}   // begin iterator
   DSNodeIterator(DSNode *N, bool)       // Create end iterator
-    : Node(N), Link(N->getNumLinks()) {
+    : Node(N), Link(N->getSize()) {
   }
 public:
 
@@ -36,38 +31,60 @@ public:
     return Link == x.Link;
   }
   bool operator!=(const _Self& x) const { return !operator==(x); }
+
+  const _Self &operator=(const _Self &I) {
+    assert(I.Node == Node && "Cannot assign iterators to two different nodes!");
+    Link = I.Link;
+    return *this;
+  }
   
   pointer operator*() const {
-    return Node->getLink(Link);
+    DSNodeHandle *NH = Node->getLink(Link);
+    return NH ? NH->getNode() : 0;
   }
   pointer operator->() const { return operator*(); }
   
   _Self& operator++() {                // Preincrement
-    unsigned NumLinks = Node->getNumLinks();
-    do {
-      ++Link;
-    } while (Link < NumLinks && Node->getLink(Link) != 0);
+    ++Link;
     return *this;
   }
   _Self operator++(int) { // Postincrement
     _Self tmp = *this; ++*this; return tmp; 
   }
-};
 
-
-template <> struct GraphTraits<DSNode*> {
-  typedef DSNode NodeType;
-  typedef DSNode::iterator ChildIteratorType;
-
-  static NodeType *getEntryNode(DSNode *N) { return N; }
-  static ChildIteratorType child_begin(NodeType *N) { return N->begin(); }
-  static ChildIteratorType child_end(NodeType *N) { return N->end(); }
+  unsigned getLink() const { return Link; }
+  DSNode *getNode() const { return Node; }
 };
 
 // Provide iterators for DSNode...
 inline DSNode::iterator DSNode::begin() { return DSNodeIterator(this); }
 inline DSNode::iterator DSNode::end()   { return DSNodeIterator(this, false); }
 
-#endif
+template <> struct GraphTraits<DSNode*> {
+  typedef DSNode NodeType;
+  typedef DSNode::iterator ChildIteratorType;
+
+  static NodeType *getEntryNode(NodeType *N) { return N; }
+  static ChildIteratorType child_begin(NodeType *N) { return N->begin(); }
+  static ChildIteratorType child_end(NodeType *N) { return N->end(); }
+};
+
+static DSNode &dereference(DSNode *N) { return *N; }
+
+template <> struct GraphTraits<DSGraph*> {
+  typedef DSNode NodeType;
+  typedef DSNode::iterator ChildIteratorType;
+
+  typedef std::pointer_to_unary_function<DSNode *, DSNode&> DerefFun;
+
+  // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
+  typedef mapped_iterator<std::vector<DSNode*>::iterator,
+                          DerefFun> nodes_iterator;
+  static nodes_iterator nodes_begin(DSGraph *G) { return map_iterator(G->getNodes().begin(), DerefFun(dereference));}
+  static nodes_iterator nodes_end  (DSGraph *G) { return map_iterator(G->getNodes().end(), DerefFun(dereference)); }
+
+  static ChildIteratorType child_begin(NodeType *N) { return N->begin(); }
+  static ChildIteratorType child_end(NodeType *N) { return N->end(); }
+};
 
 #endif