diff --git a/llvm/tools/llvm-xray/xray-graph.h b/llvm/tools/llvm-xray/xray-graph.h
index 8b0e208..fd8ac17 100644
--- a/llvm/tools/llvm-xray/xray-graph.h
+++ b/llvm/tools/llvm-xray/xray-graph.h
@@ -24,6 +24,7 @@
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/XRay/Graph.h"
 #include "llvm/XRay/Trace.h"
 #include "llvm/XRay/XRayRecord.h"
 
@@ -49,21 +50,22 @@
     std::string getAsString(StatType T) const;
     double compare(StatType T, const TimeStat &Other) const;
   };
+  typedef uint64_t TimestampT;
 
   /// An inner struct for storing edge attributes for our graph. Here the
   /// attributes are mainly function call statistics.
   ///
   /// FIXME: expand to contain more information eg call latencies.
-  struct EdgeAttribute {
+  struct CallStats {
     TimeStat S;
-    std::vector<uint64_t> Timings;
+    std::vector<TimestampT> Timings;
   };
 
   /// An Inner Struct for storing vertex attributes, at the moment just
   /// SymbolNames, however in future we could store bulk function statistics.
   ///
   /// FIXME: Store more attributes based on instrumentation map.
-  struct VertexAttribute {
+  struct FunctionStats {
     std::string SymbolName;
     TimeStat S;
   };
@@ -78,17 +80,15 @@
   typedef DenseMap<llvm::sys::ProcessInfo::ProcessId, FunctionStack>
       PerThreadFunctionStackMap;
 
-private:
-  /// The Graph stored in an edge-list like format, with the edges also having
-  /// An attached set of attributes.
-  DenseMap<int32_t, DenseMap<int32_t, EdgeAttribute>> Graph;
+  class GraphT : public Graph<FunctionStats, CallStats, int32_t> {
+  public:
+    TimeStat GraphEdgeMax = {};
+    TimeStat GraphVertexMax = {};
+  };
 
-  /// Graph Vertex Attributes. These are presently stored seperate from the
-  /// main graph.
-  DenseMap<int32_t, VertexAttribute> VertexAttrs;
-
-  TimeStat GraphEdgeMax;
-  TimeStat GraphVertexMax;
+  GraphT G;
+  typedef typename decltype(G)::VertexIdentifier VertexIdentifier;
+  typedef typename decltype(G)::EdgeIdentifier EdgeIdentifier;
 
   /// Use a Map to store the Function stack for each thread whilst building the
   /// graph.
@@ -99,7 +99,7 @@
   /// Usefull object for getting human readable Symbol Names.
   FuncIdConversionHelper &FuncIdHelper;
   bool DeduceSiblingCalls = false;
-  uint64_t CurrentMaxTSC = 0;
+  TimestampT CurrentMaxTSC = 0;
 
   /// A private function to help implement the statistic generation functions;
   template <typename U>
@@ -121,7 +121,9 @@
   /// Takes in a reference to a FuncIdHelper in order to have ready access to
   /// Symbol names.
   explicit GraphRenderer(FuncIdConversionHelper &FuncIdHelper, bool DSC)
-      : FuncIdHelper(FuncIdHelper), DeduceSiblingCalls(DSC) {}
+      : FuncIdHelper(FuncIdHelper), DeduceSiblingCalls(DSC) {
+    G[0] = {};
+  }
 
   /// Process an Xray record and expand the graph.
   ///
@@ -132,7 +134,7 @@
   /// FIXME: Make this more robust against small irregularities.
   Error accountRecord(const XRayRecord &Record);
 
-  const PerThreadFunctionStackMap getPerThreadFunctionStack() const {
+  const PerThreadFunctionStackMap &getPerThreadFunctionStack() const {
     return PerThreadFunctionStack;
   }
 
@@ -143,6 +145,13 @@
                         StatType EdgeColor = StatType::NONE,
                         StatType VertexLabel = StatType::NONE,
                         StatType VertexColor = StatType::NONE);
+
+  /// Get a reference to the internal graph.
+  const GraphT &getGraph() {
+    calculateEdgeStatistics();
+    calculateVertexStatistics();
+    return G;
+  }
 };
 }
 }