[PM/LCG] Teach the LazyCallGraph how to replace a function without disturbing the graph or having to update edges.

[chandlerc]

Summary

This is motivated by porting argument promotion to the new pass manager.
Because of how LLVM IR Function objects work, in order to change their
signature a new object needs to be created. This is efficient and
straight forward in the IR but previously was very hard to implement in
LCG. We could easily replace the function a node in the graph
represents. The challenging part is how to handle updating the edges in
the graph.

LCG previously used an edge to a raw function to represent a node that
had not yet been scanned for calls and references. This was the core
of its laziness. However, that model causes this kind of update to be
very hard:

1. The keys to lookup an edge need to be Function*s that would all need to be updated when we update the node.
2. There will be some unknown number of edges that haven't transitioned from Function* edges to Node* edges.

All of this complexity isn't necessary. Instead, we can always build
a node around any function, always pointing edges at it and always using
it as the key to lookup an edge. To maintain the laziness, we need to
sink the *edges* of a node into a secondary object and explicitly model
transitioning a node from empty to populated by scanning the function.
This design seems much cleaner in a number of ways, but importantly
there is now exactly *one* place where the Function* has to be
updated!

Some other cleanups that fall out of this include having something to
model the *entry* edges more accurately. Rather than hand rolling parts
of the node in the graph itself, we have an explicit EdgeSequence
object that gives us exactly the functionality needed. We also have
a consistent place to define the edge iterators and can use them for
both the entry edges and the internal edges of the graph.

The API used to model the separation between a node and its edges is
intentionally very thin as most clients are expected to deal with nodes
that have populated edges. We model this exactly as an optional does
with an additional method to populate the edges when that is
a reasonable thing for a client to do. This is based on API design
suggestions from Richard Smith and David Blaikie, credit goes to them
for helping pick how to model this without it being either too explicit
or too implicit.

The patch is somewhat noisy due to shifting around iterator types and
new syntax for walking the edges of a node, but most of the
functionality change is in the Edge, EdgeSequence, and Node types.

Depends on D29381.

[davide]

While this patch is large, a fair amount of it are just mechanical changes. I wasn't able to spot any obvious mistake, but I would appreciate another pair of eyes looking at it (@sanjoy ?)

include/llvm/Analysis/LazyCallGraph.h

330–331 ↗

(On Diff #87203)

Not sure if static_cast<bool> works here, but I would prefer that, if possible.

[chandlerc]

Rebase and address a comment from Davide.

[chandlerc]

include/llvm/Analysis/LazyCallGraph.h

330–331 ↗

(On Diff #87203)

Sure, but we can do better with Optional::hasValue.

[davide]

LGTM

[sanjoy]

lgtm, but I glossed over much of the mechanical changes.

[chandlerc]

Thanks all, landing!

llvm/trunk/include/llvm/Analysis/LazyCallGraph.h

 /// invalidates the information in this analysis.
 ///
 /// FIXME: This class is named LazyCallGraph in a lame attempt to distinguish
 /// it from the existing CallGraph. At some point, it is expected that this
 /// will be the only call graph and it will be renamed accordingly.
 class LazyCallGraph {
 public:
   class Node;
+  class EdgeSequence;
   class SCC;
   class RefSCC;
   class edge_iterator;
   class call_edge_iterator;
 
   /// A class used to represent edges in the call graph.
   ///
   /// The lazy call graph models both *call* edges and *reference* edges. Call
   /// edges are much what you would expect, and exist when there is a 'call' or
   /// 'invoke' instruction of some function. Reference edges are also tracked
   /// along side these, and exist whenever any instruction (transitively
   /// through its operands) references a function. All call edges are
   /// inherently reference edges, and so the reference graph forms a superset
   /// of the formal call graph.
   ///
-  /// Furthermore, edges also may point to raw \c Function objects when those
-  /// functions have not been scanned and incorporated into the graph (yet).
-  /// This is one of the primary ways in which the graph can be lazy. When
-  /// functions are scanned and fully incorporated into the graph, all of the
-  /// edges referencing them are updated to point to the graph \c Node objects
-  /// instead of to the raw \c Function objects. This class even provides
-  /// methods to trigger this scan on-demand by attempting to get the target
-  /// node of the graph and providing a reference back to the graph in order to
-  /// lazily build it if necessary.
-  ///
   /// All of these forms of edges are fundamentally represented as outgoing
   /// edges. The edges are stored in the source node and point at the target
   /// node. This allows the edge structure itself to be a very compact data
   /// structure: essentially a tagged pointer.
   class Edge {
   public:
     /// The kind of edge in the graph.
     enum Kind : bool { Ref = false, Call = true };
 
     Edge();
-    explicit Edge(Function &F, Kind K);
     explicit Edge(Node &N, Kind K);
 
     /// Test whether the edge is null.
     ///
     /// This happens when an edge has been deleted. We leave the edge objects
     /// around but clear them.
     explicit operator bool() const;
 
     /// Returnss the \c Kind of the edge.
     Kind getKind() const;
 
     /// Test whether the edge represents a direct call to a function.
     ///
     /// This requires that the edge is not null.
     bool isCall() const;
 
-    /// Get the function referenced by this edge.
+    /// Get the call graph node referenced by this edge.
     ///
-    /// This requires that the edge is not null, but will succeed whether we
-    /// have built a graph node for the function yet or not.
-    Function &getFunction() const;
+    /// This requires that the edge is not null.
+    Node &getNode() const;
 
-    /// Get the call graph node referenced by this edge if one exists.
+    /// Get the function referenced by this edge.
     ///
-    /// This requires that the edge is not null. If we have built a graph node
-    /// for the function this edge points to, this will return that node,
-    /// otherwise it will return null.
-    Node *getNode() const;
-
-    /// Get the call graph node for this edge, building it if necessary.
-    ///
-    /// This requires that the edge is not null. If we have not yet built
-    /// a graph node for the function this edge points to, this will first ask
-    /// the graph to build that node, inserting it into all the relevant
-    /// structures.
-    Node &getNode(LazyCallGraph &G);
+    /// This requires that the edge is not null.
+    Function &getFunction() const;
 
   private:
-    friend class LazyCallGraph::Node;
+    friend class LazyCallGraph::EdgeSequence;
     friend class LazyCallGraph::RefSCC;
 
-    PointerIntPair<PointerUnion<Function *, Node *>, 1, Kind> Value;
+    PointerIntPair<Node *, 1, Kind> Value;
 
     void setKind(Kind K) { Value.setInt(K); }
   };
 
-  typedef SmallVector<Edge, 4> EdgeVectorT;
-  typedef SmallVectorImpl<Edge> EdgeVectorImplT;
-
-  /// A node in the call graph.
+  /// The edge sequence object.
   ///
-  /// This represents a single node. It's primary roles are to cache the list of
-  /// callees, de-duplicate and provide fast testing of whether a function is
-  /// a callee, and facilitate iteration of child nodes in the graph.
-  class Node {
+  /// This typically exists entirely within the node but is exposed as
+  /// a separate type because a node doesn't initially have edges. An explicit
+  /// population step is required to produce this sequence at first and it is
+  /// then cached in the node. It is also used to represent edges entering the
+  /// graph from outside the module to model the graph's roots.
+  ///
+  /// The sequence itself both iterable and indexable. The indexes remain
+  /// stable even as the sequence mutates (including removal).
+  class EdgeSequence {
     friend class LazyCallGraph;
-    friend class LazyCallGraph::SCC;
+    friend class LazyCallGraph::Node;
     friend class LazyCallGraph::RefSCC;
 
-    LazyCallGraph *G;
-    Function &F;
-
-    // We provide for the DFS numbering and Tarjan walk lowlink numbers to be
-    // stored directly within the node. These are both '-1' when nodes are part
-    // of an SCC (or RefSCC), or '0' when not yet reached in a DFS walk.
-    int DFSNumber;
-    int LowLink;
-
-    mutable EdgeVectorT Edges;
-    DenseMap<Function *, int> EdgeIndexMap;
-
-    /// Basic constructor implements the scanning of F into Edges and
-    /// EdgeIndexMap.
-    Node(LazyCallGraph &G, Function &F);
-
-    /// Internal helper to insert an edge to a function.
-    void insertEdgeInternal(Function &ChildF, Edge::Kind EK);
-
-    /// Internal helper to insert an edge to a node.
-    void insertEdgeInternal(Node &ChildN, Edge::Kind EK);
-
-    /// Internal helper to change an edge kind.
-    void setEdgeKind(Function &ChildF, Edge::Kind EK);
-
-    /// Internal helper to remove the edge to the given function.
-    void removeEdgeInternal(Function &ChildF);
-
-    void clear() {
-      Edges.clear();
-      EdgeIndexMap.clear();
-    }
-
-    /// Print the name of this node's function.
-    friend raw_ostream &operator<<(raw_ostream &OS, const Node &N) {
-      return OS << N.F.getName();
-    }
-
-    /// Dump the name of this node's function to stderr.
-    void dump() const;
+    typedef SmallVector<Edge, 4> VectorT;
+    typedef SmallVectorImpl<Edge> VectorImplT;
 
   public:
-    LazyCallGraph &getGraph() const { return *G; }
-
-    Function &getFunction() const { return F; }
-
-    edge_iterator begin() const {
-      return edge_iterator(Edges.begin(), Edges.end());
-    }
-    edge_iterator end() const { return edge_iterator(Edges.end(), Edges.end()); }
-
-    const Edge &operator[](int i) const { return Edges[i]; }
-    const Edge &operator[](Function &F) const {
-      assert(EdgeIndexMap.find(&F) != EdgeIndexMap.end() && "No such edge!");
-      return Edges[EdgeIndexMap.find(&F)->second];
-    }
-    const Edge &operator[](Node &N) const { return (*this)[N.getFunction()]; }
-
-    const Edge *lookup(Function &F) const {
-      auto EI = EdgeIndexMap.find(&F);
-      return EI != EdgeIndexMap.end() ? &Edges[EI->second] : nullptr;
-    }
-
-    call_edge_iterator call_begin() const {
-      return call_edge_iterator(Edges.begin(), Edges.end());
-    }
-    call_edge_iterator call_end() const {
-      return call_edge_iterator(Edges.end(), Edges.end());
-    }
-
-    iterator_range<call_edge_iterator> calls() const {
-      return make_range(call_begin(), call_end());
-    }
-
-    /// Equality is defined as address equality.
-    bool operator==(const Node &N) const { return this == &N; }
-    bool operator!=(const Node &N) const { return !operator==(N); }
-  };
-
-  /// A lazy iterator used for both the entry nodes and child nodes.
-  ///
-  /// When this iterator is dereferenced, if not yet available, a function will
-  /// be scanned for "calls" or uses of functions and its child information
-  /// will be constructed. All of these results are accumulated and cached in
-  /// the graph.
-  class edge_iterator
-      : public iterator_adaptor_base<edge_iterator, EdgeVectorImplT::iterator,
+    /// An iterator used for the edges to both entry nodes and child nodes.
+    class iterator
+        : public iterator_adaptor_base<iterator, VectorImplT::iterator,
                                        std::forward_iterator_tag> {
       friend class LazyCallGraph;
       friend class LazyCallGraph::Node;
 
-    EdgeVectorImplT::iterator E;
+      VectorImplT::iterator E;
 
       // Build the iterator for a specific position in the edge list.
-    edge_iterator(EdgeVectorImplT::iterator BaseI,
-                  EdgeVectorImplT::iterator E)
+      iterator(VectorImplT::iterator BaseI, VectorImplT::iterator E)
           : iterator_adaptor_base(BaseI), E(E) {
         while (I != E && !*I)
           ++I;
       }
 
     public:
-    edge_iterator() {}
+      iterator() {}
 
       using iterator_adaptor_base::operator++;
-    edge_iterator &operator++() {
+      iterator &operator++() {
         do {
           ++I;
         } while (I != E && !*I);
         return *this;
       }
     };
 
-  /// A lazy iterator over specifically call edges.
+    /// An iterator over specifically call edges.
     ///
-  /// This has the same iteration properties as the \c edge_iterator, but
+    /// This has the same iteration properties as the \c iterator, but
     /// restricts itself to edges which represent actual calls.
-  class call_edge_iterator
-      : public iterator_adaptor_base<call_edge_iterator,
-                                     EdgeVectorImplT::iterator,
+    class call_iterator
+        : public iterator_adaptor_base<call_iterator, VectorImplT::iterator,
                                        std::forward_iterator_tag> {
       friend class LazyCallGraph;
       friend class LazyCallGraph::Node;
 
-    EdgeVectorImplT::iterator E;
+      VectorImplT::iterator E;
 
       /// Advance the iterator to the next valid, call edge.
       void advanceToNextEdge() {
         while (I != E && (!*I || !I->isCall()))
           ++I;
       }
 
       // Build the iterator for a specific position in the edge list.
-    call_edge_iterator(EdgeVectorImplT::iterator BaseI,
-                       EdgeVectorImplT::iterator E)
+      call_iterator(VectorImplT::iterator BaseI, VectorImplT::iterator E)
           : iterator_adaptor_base(BaseI), E(E) {
         advanceToNextEdge();
       }
 
     public:
-    call_edge_iterator() {}
+      call_iterator() {}
 
       using iterator_adaptor_base::operator++;
-    call_edge_iterator &operator++() {
+      call_iterator &operator++() {
         ++I;
         advanceToNextEdge();
         return *this;
       }
     };
 
+    iterator begin() { return iterator(Edges.begin(), Edges.end()); }
+    iterator end() { return iterator(Edges.end(), Edges.end()); }
+
+    Edge &operator[](int i) { return Edges[i]; }
+    Edge &operator[](Node &N) {
+      assert(EdgeIndexMap.find(&N) != EdgeIndexMap.end() && "No such edge!");
+      return Edges[EdgeIndexMap.find(&N)->second];
+    }
+    Edge *lookup(Node &N) {
+      auto EI = EdgeIndexMap.find(&N);
+      return EI != EdgeIndexMap.end() ? &Edges[EI->second] : nullptr;
+    }
+
+    call_iterator call_begin() {
+      return call_iterator(Edges.begin(), Edges.end());
+    }
+    call_iterator call_end() { return call_iterator(Edges.end(), Edges.end()); }
+
+    iterator_range<call_iterator> calls() {
+      return make_range(call_begin(), call_end());
+    }
+
+    bool empty() {
+      for (auto &E : Edges)
+        if (E)
+          return false;
+
+      return true;
+    }
+
+  private:
+    VectorT Edges;
+    DenseMap<Node *, int> EdgeIndexMap;
+
+    EdgeSequence() = default;
+
+    /// Internal helper to insert an edge to a node.
+    void insertEdgeInternal(Node &ChildN, Edge::Kind EK);
+
+    /// Internal helper to change an edge kind.
+    void setEdgeKind(Node &ChildN, Edge::Kind EK);
+
+    /// Internal helper to remove the edge to the given function.
+    bool removeEdgeInternal(Node &ChildN);
+
+    /// Internal helper to replace an edge key with a new one.
+    ///
+    /// This should be used when the function for a particular node in the
+    /// graph gets replaced and we are updating all of the edges to that node
+    /// to use the new function as the key.
+    void replaceEdgeKey(Function &OldTarget, Function &NewTarget);
+  };
+
+  /// A node in the call graph.
+  ///
+  /// This represents a single node. It's primary roles are to cache the list of
+  /// callees, de-duplicate and provide fast testing of whether a function is
+  /// a callee, and facilitate iteration of child nodes in the graph.
+  ///
+  /// The node works much like an optional in order to lazily populate the
+  /// edges of each node. Until populated, there are no edges. Once populated,
+  /// you can access the edges by dereferencing the node or using the `->`
+  /// operator as if the node was an `Optional<EdgeSequence>`.
+  class Node {
+    friend class LazyCallGraph;
+    friend class LazyCallGraph::RefSCC;
+
+  public:
+    LazyCallGraph &getGraph() const { return *G; }
+
+    Function &getFunction() const { return *F; }
+
+    StringRef getName() const { return F->getName(); }
+
+    /// Equality is defined as address equality.
+    bool operator==(const Node &N) const { return this == &N; }
+    bool operator!=(const Node &N) const { return !operator==(N); }
+
+    /// Tests whether the node has been populated with edges.
+    operator bool() const { return Edges.hasValue(); }
+
+    // We allow accessing the edges by dereferencing or using the arrow
+    // operator, essentially wrapping the internal optional.
+    EdgeSequence &operator*() const {
+      // Rip const off because the node itself isn't changing here.
+      return const_cast<EdgeSequence &>(*Edges);
+    }
+    EdgeSequence *operator->() const { return &**this; }
+
+    /// Populate the edges of this node if necessary.
+    ///
+    /// The first time this is called it will populate the edges for this node
+    /// in the graph. It does this by scanning the underlying function, so once
+    /// this is done, any changes to that function must be explicitly reflected
+    /// in updates to the graph.
+    ///
+    /// \returns the populated \c EdgeSequence to simplify walking it.
+    ///
+    /// This will not update or re-scan anything if called repeatedly. Instead,
+    /// the edge sequence is cached and returned immediately on subsequent
+    /// calls.
+    EdgeSequence &populate() {
+      if (Edges)
+        return *Edges;
+
+      return populateSlow();
+    }
+
+  private:
+    LazyCallGraph *G;
+    Function *F;
+
+    // We provide for the DFS numbering and Tarjan walk lowlink numbers to be
+    // stored directly within the node. These are both '-1' when nodes are part
+    // of an SCC (or RefSCC), or '0' when not yet reached in a DFS walk.
+    int DFSNumber;
+    int LowLink;
+
+    Optional<EdgeSequence> Edges;
+
+    /// Basic constructor implements the scanning of F into Edges and
+    /// EdgeIndexMap.
+    Node(LazyCallGraph &G, Function &F)
+        : G(&G), F(&F), DFSNumber(0), LowLink(0) {}
+
+    /// Implementation of the scan when populating.
+    EdgeSequence &populateSlow();
+
+    /// Internal helper to directly replace the function with a new one.
+    ///
+    /// This is used to facilitate tranfsormations which need to replace the
+    /// formal Function object but directly move the body and users from one to
+    /// the other.
+    void replaceFunction(Function &NewF);
+
+    void clear() { Edges.reset(); }
+
+    /// Print the name of this node's function.
+    friend raw_ostream &operator<<(raw_ostream &OS, const Node &N) {
+      return OS << N.F->getName();
+    }
+
+    /// Dump the name of this node's function to stderr.
+    void dump() const;
+  };
+
   /// An SCC of the call graph.
   ///
   /// This represents a Strongly Connected Component of the direct call graph
   /// -- ignoring indirect calls and function references. It stores this as
   /// a collection of call graph nodes. While the order of nodes in the SCC is
   /// stable, it is not any particular order.
   ///
   /// The SCCs are nested within a \c RefSCC, see below for details about that
public:
     /// This is trivial whenever the target is in the same RefSCC as the source
     /// or the edge is an outgoing edge to some descendant RefSCC. In these
     /// cases there is no change to the cyclic structure of the RefSCCs.
     ///
     /// To further make calling this convenient, it also handles inserting
     /// already existing edges.
     void insertTrivialRefEdge(Node &SourceN, Node &TargetN);
 
+    /// Directly replace a node's function with a new function.
+    ///
+    /// This should be used when moving the body and users of a function to
+    /// a new formal function object but not otherwise changing the call graph
+    /// structure in any way.
+    ///
+    /// It requires that the old function in the provided node have zero uses
+    /// and the new function must have calls and references to it establishing
+    /// an equivalent graph.
+    void replaceNodeFunction(Node &N, Function &NewF);
+
     ///@}
   };
 
   /// A post-order depth-first RefSCC iterator over the call graph.
   ///
   /// This iterator walks the cached post-order sequence of RefSCCs. However,
   /// it trades stability for flexibility. It is restricted to a forward
   /// iterator but will survive mutations which insert new RefSCCs and continue
#endif
   /// This sets up the graph and computes all of the entry points of the graph.
   /// No function definitions are scanned until their nodes in the graph are
   /// requested during traversal.
   LazyCallGraph(Module &M);
 
   LazyCallGraph(LazyCallGraph &&G);
   LazyCallGraph &operator=(LazyCallGraph &&RHS);
 
-  edge_iterator begin() {
-    return edge_iterator(EntryEdges.begin(), EntryEdges.end());
-  }
-  edge_iterator end() {
-    return edge_iterator(EntryEdges.end(), EntryEdges.end());
-  }
+  EdgeSequence::iterator begin() { return EntryEdges.begin(); }
+  EdgeSequence::iterator end() { return EntryEdges.end(); }
 
   void buildRefSCCs();
 
   postorder_ref_scc_iterator postorder_ref_scc_begin() {
     if (!EntryEdges.empty())
       assert(!PostOrderRefSCCs.empty() &&
              "Must form RefSCCs before iterating them!");
     return postorder_ref_scc_iterator(*this);
#endif
   /// call graph. They can be used to update the core node-graph during
   /// a node-based inorder traversal that precedes any SCC-based traversal.
   ///
   /// Once you begin manipulating a call graph's SCCs, most mutation of the
   /// graph must be performed via a RefSCC method. There are some exceptions
   /// below.
 
   /// Update the call graph after inserting a new edge.
-  void insertEdge(Node &Caller, Function &Callee, Edge::Kind EK);
+  void insertEdge(Node &SourceN, Node &TargetN, Edge::Kind EK);
 
   /// Update the call graph after inserting a new edge.
-  void insertEdge(Function &Caller, Function &Callee, Edge::Kind EK) {
-    return insertEdge(get(Caller), Callee, EK);
+  void insertEdge(Function &Source, Function &Target, Edge::Kind EK) {
+    return insertEdge(get(Source), get(Target), EK);
   }
 
   /// Update the call graph after deleting an edge.
-  void removeEdge(Node &Caller, Function &Callee);
+  void removeEdge(Node &SourceN, Node &TargetN);
 
   /// Update the call graph after deleting an edge.
-  void removeEdge(Function &Caller, Function &Callee) {
-    return removeEdge(get(Caller), Callee);
+  void removeEdge(Function &Source, Function &Target) {
+    return removeEdge(get(Source), get(Target));
   }
 
   ///@}
 
   ///@{
   /// \name General Mutation API
   ///
   /// There are a very limited set of mutations allowed on the graph as a whole
private:
   typedef iterator_range<node_stack_iterator> node_stack_range;
 
   /// Allocator that holds all the call graph nodes.
   SpecificBumpPtrAllocator<Node> BPA;
 
   /// Maps function->node for fast lookup.
   DenseMap<const Function *, Node *> NodeMap;
 
-  /// The entry nodes to the graph.
+  /// The entry edges into the graph.
   ///
-  /// These nodes are reachable through "external" means. Put another way, they
+  /// These edges are from "external" sources. Put another way, they
   /// escape at the module scope.
-  EdgeVectorT EntryEdges;
-
-  /// Map of the entry nodes in the graph to their indices in \c EntryEdges.
-  DenseMap<Function *, int> EntryIndexMap;
+  EdgeSequence EntryEdges;
 
   /// Allocator that holds all the call graph SCCs.
   SpecificBumpPtrAllocator<SCC> SCCBPA;
 
   /// Maps Function -> SCC for fast lookup.
   DenseMap<Node *, SCC *> SCCMap;
 
   /// Allocator that holds all the call graph RefSCCs.
int getRefSCCIndex(RefSCC &RC) {
     assert(IndexIt != RefSCCIndices.end() && "RefSCC doesn't have an index!");
     assert(PostOrderRefSCCs[IndexIt->second] == &RC &&
            "Index does not point back at RC!");
     return IndexIt->second;
   }
 };
 
 inline LazyCallGraph::Edge::Edge() : Value() {}
-inline LazyCallGraph::Edge::Edge(Function &F, Kind K) : Value(&F, K) {}
 inline LazyCallGraph::Edge::Edge(Node &N, Kind K) : Value(&N, K) {}
 
-inline LazyCallGraph::Edge::operator bool() const {
-  return !Value.getPointer().isNull();
-}
+inline LazyCallGraph::Edge::operator bool() const { return Value.getPointer(); }
 
 inline LazyCallGraph::Edge::Kind LazyCallGraph::Edge::getKind() const {
   assert(*this && "Queried a null edge!");
   return Value.getInt();
 }
 
 inline bool LazyCallGraph::Edge::isCall() const {
   assert(*this && "Queried a null edge!");
   return getKind() == Call;
 }
 
-inline Function &LazyCallGraph::Edge::getFunction() const {
+inline LazyCallGraph::Node &LazyCallGraph::Edge::getNode() const {
   assert(*this && "Queried a null edge!");
-  auto P = Value.getPointer();
-  if (auto *F = P.dyn_cast<Function *>())
-    return *F;
-
-  return P.get<Node *>()->getFunction();
+  return *Value.getPointer();
 }
 
-inline LazyCallGraph::Node *LazyCallGraph::Edge::getNode() const {
-  assert(*this && "Queried a null edge!");
-  auto P = Value.getPointer();
-  if (auto *N = P.dyn_cast<Node *>())
-    return N;
-
-  return nullptr;
-}
-
-inline LazyCallGraph::Node &LazyCallGraph::Edge::getNode(LazyCallGraph &G) {
+inline Function &LazyCallGraph::Edge::getFunction() const {
   assert(*this && "Queried a null edge!");
-  auto P = Value.getPointer();
-  if (auto *N = P.dyn_cast<Node *>())
-    return *N;
-
-  Node &N = G.get(*P.get<Function *>());
-  Value.setPointer(&N);
-  return N;
+  return getNode().getFunction();
 }
 
 // Provide GraphTraits specializations for call graphs.
 template <> struct GraphTraits<LazyCallGraph::Node *> {
   typedef LazyCallGraph::Node *NodeRef;
-  typedef LazyCallGraph::edge_iterator ChildIteratorType;
+  typedef LazyCallGraph::EdgeSequence::iterator ChildIteratorType;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
-  static ChildIteratorType child_begin(NodeRef N) { return N->begin(); }
-  static ChildIteratorType child_end(NodeRef N) { return N->end(); }
+  static ChildIteratorType child_begin(NodeRef N) { return (*N)->begin(); }
+  static ChildIteratorType child_end(NodeRef N) { return (*N)->end(); }
 };
 template <> struct GraphTraits<LazyCallGraph *> {
   typedef LazyCallGraph::Node *NodeRef;
-  typedef LazyCallGraph::edge_iterator ChildIteratorType;
+  typedef LazyCallGraph::EdgeSequence::iterator ChildIteratorType;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
-  static ChildIteratorType child_begin(NodeRef N) { return N->begin(); }
-  static ChildIteratorType child_end(NodeRef N) { return N->end(); }
+  static ChildIteratorType child_begin(NodeRef N) { return (*N)->begin(); }
+  static ChildIteratorType child_end(NodeRef N) { return (*N)->end(); }
 };
 
 /// An analysis pass which computes the call graph for a module.
 class LazyCallGraphAnalysis : public AnalysisInfoMixin<LazyCallGraphAnalysis> {
   friend AnalysisInfoMixin<LazyCallGraphAnalysis>;
   static AnalysisKey Key;
 
 public:

llvm/trunk/lib/Analysis/CGSCCPassManager.cpp

LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
   SCC *C = &InitialC;
   RefSCC *RC = &InitialRC;
   Function &F = N.getFunction();
 
   // Walk the function body and build up the set of retained, promoted, and
   // demoted edges.
   SmallVector<Constant *, 16> Worklist;
   SmallPtrSet<Constant *, 16> Visited;
-  SmallPtrSet<Function *, 16> RetainedEdges;
-  SmallSetVector<Function *, 4> PromotedRefTargets;
-  SmallSetVector<Function *, 4> DemotedCallTargets;
+  SmallPtrSet<Node *, 16> RetainedEdges;
+  SmallSetVector<Node *, 4> PromotedRefTargets;
+  SmallSetVector<Node *, 4> DemotedCallTargets;
 
   // First walk the function and handle all called functions. We do this first
   // because if there is a single call edge, whether there are ref edges is
   // irrelevant.
   for (Instruction &I : instructions(F))
     if (auto CS = CallSite(&I))
       if (Function *Callee = CS.getCalledFunction())
         if (Visited.insert(Callee).second && !Callee->isDeclaration()) {
-          const Edge *E = N.lookup(*Callee);
+          Node &CalleeN = *G.lookup(*Callee);
+          Edge *E = N->lookup(CalleeN);
           // FIXME: We should really handle adding new calls. While it will
           // make downstream usage more complex, there is no fundamental
           // limitation and it will allow passes within the CGSCC to be a bit
           // more flexible in what transforms they can do. Until then, we
           // verify that new calls haven't been introduced.
           assert(E && "No function transformations should introduce *new* "
                       "call edges! Any new calls should be modeled as "
                       "promoted existing ref edges!");
-          RetainedEdges.insert(Callee);
+          RetainedEdges.insert(&CalleeN);
           if (!E->isCall())
-            PromotedRefTargets.insert(Callee);
+            PromotedRefTargets.insert(&CalleeN);
         }
 
   // Now walk all references.
   for (Instruction &I : instructions(F))
     for (Value *Op : I.operand_values())
       if (Constant *C = dyn_cast<Constant>(Op))
         if (Visited.insert(C).second)
           Worklist.push_back(C);
 
   LazyCallGraph::visitReferences(Worklist, Visited, [&](Function &Referee) {
-    const Edge *E = N.lookup(Referee);
+    Node &RefereeN = *G.lookup(Referee);
+    Edge *E = N->lookup(RefereeN);
     // FIXME: Similarly to new calls, we also currently preclude
     // introducing new references. See above for details.
     assert(E && "No function transformations should introduce *new* ref "
                 "edges! Any new ref edges would require IPO which "
                 "function passes aren't allowed to do!");
-    RetainedEdges.insert(&Referee);
+    RetainedEdges.insert(&RefereeN);
     if (E->isCall())
-      DemotedCallTargets.insert(&Referee);
+      DemotedCallTargets.insert(&RefereeN);
   });
 
   // First remove all of the edges that are no longer present in this function.
   // We have to build a list of dead targets first and then remove them as the
   // data structures will all be invalidated by removing them.
   SmallVector<PointerIntPair<Node *, 1, Edge::Kind>, 4> DeadTargets;
-  for (Edge &E : N)
-    if (!RetainedEdges.count(&E.getFunction()))
-      DeadTargets.push_back({E.getNode(), E.getKind()});
+  for (Edge &E : *N)
+    if (!RetainedEdges.count(&E.getNode()))
+      DeadTargets.push_back({&E.getNode(), E.getKind()});
   for (auto DeadTarget : DeadTargets) {
     Node &TargetN = *DeadTarget.getPointer();
     bool IsCall = DeadTarget.getInt() == Edge::Call;
     SCC &TargetC = *G.lookupSCC(TargetN);
     RefSCC &TargetRC = TargetC.getOuterRefSCC();
 
     if (&TargetRC != RC) {
       RC->removeOutgoingEdge(N, TargetN);
if (!NewRefSCCs.empty()) {
                  << "\n";
       }
     }
   }
 
   // Next demote all the call edges that are now ref edges. This helps make
   // the SCCs small which should minimize the work below as we don't want to
   // form cycles that this would break.
-  for (Function *RefTarget : DemotedCallTargets) {
-    Node &TargetN = *G.lookup(*RefTarget);
-    SCC &TargetC = *G.lookupSCC(TargetN);
+  for (Node *RefTarget : DemotedCallTargets) {
+    SCC &TargetC = *G.lookupSCC(*RefTarget);
     RefSCC &TargetRC = TargetC.getOuterRefSCC();
 
     // The easy case is when the target RefSCC is not this RefSCC. This is
     // only supported when the target RefSCC is a child of this RefSCC.
     if (&TargetRC != RC) {
       assert(RC->isAncestorOf(TargetRC) &&
              "Cannot potentially form RefSCC cycles here!");
-      RC->switchOutgoingEdgeToRef(N, TargetN);
+      RC->switchOutgoingEdgeToRef(N, *RefTarget);
       if (DebugLogging)
         dbgs() << "Switch outgoing call edge to a ref edge from '" << N
-               << "' to '" << TargetN << "'\n";
+               << "' to '" << *RefTarget << "'\n";
       continue;
     }
 
     // We are switching an internal call edge to a ref edge. This may split up
     // some SCCs.
     if (C != &TargetC) {
       // For separate SCCs this is trivial.
-      RC->switchTrivialInternalEdgeToRef(N, TargetN);
+      RC->switchTrivialInternalEdgeToRef(N, *RefTarget);
       continue;
     }
 
     // Otherwise we may end up re-structuring the call graph. First, invalidate
     // any SCC analyses. We have to do this before we split functions into new
     // SCCs and lose track of where their analyses are cached.
     // FIXME: We should accept a more precise preserved set here. For example,
     // it might be possible to preserve some function analyses even as the SCC
     // structure is changed.
     AM.invalidate(*C, PreservedAnalyses::none());
     // Now update the call graph.
-    C = incorporateNewSCCRange(RC->switchInternalEdgeToRef(N, TargetN), G,
-                               N, C, AM, UR, DebugLogging);
+    C = incorporateNewSCCRange(RC->switchInternalEdgeToRef(N, *RefTarget), G, N,
+                               C, AM, UR, DebugLogging);
   }
 
   // Now promote ref edges into call edges.
-  for (Function *CallTarget : PromotedRefTargets) {
-    Node &TargetN = *G.lookup(*CallTarget);
-    SCC &TargetC = *G.lookupSCC(TargetN);
+  for (Node *CallTarget : PromotedRefTargets) {
+    SCC &TargetC = *G.lookupSCC(*CallTarget);
     RefSCC &TargetRC = TargetC.getOuterRefSCC();
 
     // The easy case is when the target RefSCC is not this RefSCC. This is
     // only supported when the target RefSCC is a child of this RefSCC.
     if (&TargetRC != RC) {
       assert(RC->isAncestorOf(TargetRC) &&
              "Cannot potentially form RefSCC cycles here!");
-      RC->switchOutgoingEdgeToCall(N, TargetN);
+      RC->switchOutgoingEdgeToCall(N, *CallTarget);
       if (DebugLogging)
         dbgs() << "Switch outgoing ref edge to a call edge from '" << N
-               << "' to '" << TargetN << "'\n";
+               << "' to '" << *CallTarget << "'\n";
       continue;
     }
     if (DebugLogging)
       dbgs() << "Switch an internal ref edge to a call edge from '" << N
-             << "' to '" << TargetN << "'\n";
+             << "' to '" << *CallTarget << "'\n";
 
     // Otherwise we are switching an internal ref edge to a call edge. This
     // may merge away some SCCs, and we add those to the UpdateResult. We also
     // need to make sure to update the worklist in the event SCCs have moved
     // before the current one in the post-order sequence.
     auto InitialSCCIndex = RC->find(*C) - RC->begin();
-    auto InvalidatedSCCs = RC->switchInternalEdgeToCall(N, TargetN);
+    auto InvalidatedSCCs = RC->switchInternalEdgeToCall(N, *CallTarget);
     if (!InvalidatedSCCs.empty()) {
       C = &TargetC;
       assert(G.lookupSCC(N) == C && "Failed to update current SCC!");
 
       // Any analyses cached for this SCC are no longer precise as the shape
       // has changed by introducing this cycle.
       AM.invalidate(*C, PreservedAnalyses::none());
 

llvm/trunk/lib/Analysis/LazyCallGraph.cpp

 #include "llvm/Support/Debug.h"
 #include "llvm/Support/GraphWriter.h"
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "lcg"
 
+void LazyCallGraph::EdgeSequence::insertEdgeInternal(Node &TargetN,
+                                                     Edge::Kind EK) {
+  EdgeIndexMap.insert({&TargetN, Edges.size()});
+  Edges.emplace_back(TargetN, EK);
+}
+
+void LazyCallGraph::EdgeSequence::setEdgeKind(Node &TargetN, Edge::Kind EK) {
+  Edges[EdgeIndexMap.find(&TargetN)->second].setKind(EK);
+}
+
+bool LazyCallGraph::EdgeSequence::removeEdgeInternal(Node &TargetN) {
+  auto IndexMapI = EdgeIndexMap.find(&TargetN);
+  if (IndexMapI == EdgeIndexMap.end())
+    return false;
+
+  Edges[IndexMapI->second] = Edge();
+  EdgeIndexMap.erase(IndexMapI);
+  return true;
+}
+
 static void addEdge(SmallVectorImpl<LazyCallGraph::Edge> &Edges,
-                    DenseMap<Function *, int> &EdgeIndexMap, Function &F,
-                    LazyCallGraph::Edge::Kind EK) {
-  if (!EdgeIndexMap.insert({&F, Edges.size()}).second)
+                    DenseMap<LazyCallGraph::Node *, int> &EdgeIndexMap,
+                    LazyCallGraph::Node &N, LazyCallGraph::Edge::Kind EK) {
+  if (!EdgeIndexMap.insert({&N, Edges.size()}).second)
     return;
 
-  DEBUG(dbgs() << "    Added callable function: " << F.getName() << "\n");
-  Edges.emplace_back(LazyCallGraph::Edge(F, EK));
+  DEBUG(dbgs() << "    Added callable function: " << N.getName() << "\n");
+  Edges.emplace_back(LazyCallGraph::Edge(N, EK));
 }
 
-LazyCallGraph::Node::Node(LazyCallGraph &G, Function &F)
-    : G(&G), F(F), DFSNumber(0), LowLink(0) {
-  DEBUG(dbgs() << "  Adding functions called by '" << F.getName()
+LazyCallGraph::EdgeSequence &LazyCallGraph::Node::populateSlow() {
+  assert(!Edges && "Must not have already populated the edges for this node!");
+
+  DEBUG(dbgs() << "  Adding functions called by '" << getName()
                << "' to the graph.\n");
 
+  Edges = EdgeSequence();
+
   SmallVector<Constant *, 16> Worklist;
   SmallPtrSet<Function *, 4> Callees;
   SmallPtrSet<Constant *, 16> Visited;
 
   // Find all the potential call graph edges in this function. We track both
   // actual call edges and indirect references to functions. The direct calls
   // are trivially added, but to accumulate the latter we walk the instructions
   // and add every operand which is a constant to the worklist to process
   // afterward.
   //
   // Note that we consider *any* function with a definition to be a viable
   // edge. Even if the function's definition is subject to replacement by
   // some other module (say, a weak definition) there may still be
   // optimizations which essentially speculate based on the definition and
   // a way to check that the specific definition is in fact the one being
   // used. For example, this could be done by moving the weak definition to
   // a strong (internal) definition and making the weak definition be an
   // alias. Then a test of the address of the weak function against the new
   // strong definition's address would be an effective way to determine the
   // safety of optimizing a direct call edge.
-  for (BasicBlock &BB : F)
+  for (BasicBlock &BB : *F)
     for (Instruction &I : BB) {
       if (auto CS = CallSite(&I))
         if (Function *Callee = CS.getCalledFunction())
           if (!Callee->isDeclaration())
             if (Callees.insert(Callee).second) {
               Visited.insert(Callee);
-              addEdge(Edges, EdgeIndexMap, *Callee, LazyCallGraph::Edge::Call);
+              addEdge(Edges->Edges, Edges->EdgeIndexMap, G->get(*Callee),
+                      LazyCallGraph::Edge::Call);
             }
 
       for (Value *Op : I.operand_values())
         if (Constant *C = dyn_cast<Constant>(Op))
           if (Visited.insert(C).second)
             Worklist.push_back(C);
     }
 
   // We've collected all the constant (and thus potentially function or
   // function containing) operands to all of the instructions in the function.
   // Process them (recursively) collecting every function found.
   visitReferences(Worklist, Visited, [&](Function &F) {
-    addEdge(Edges, EdgeIndexMap, F, LazyCallGraph::Edge::Ref);
+    addEdge(Edges->Edges, Edges->EdgeIndexMap, G->get(F),
+            LazyCallGraph::Edge::Ref);
   });
-}
-
-void LazyCallGraph::Node::insertEdgeInternal(Function &Target, Edge::Kind EK) {
-  if (Node *N = G->lookup(Target))
-    return insertEdgeInternal(*N, EK);
 
-  EdgeIndexMap.insert({&Target, Edges.size()});
-  Edges.emplace_back(Target, EK);
-}
-
-void LazyCallGraph::Node::insertEdgeInternal(Node &TargetN, Edge::Kind EK) {
-  EdgeIndexMap.insert({&TargetN.getFunction(), Edges.size()});
-  Edges.emplace_back(TargetN, EK);
+  return *Edges;
 }
 
-void LazyCallGraph::Node::setEdgeKind(Function &TargetF, Edge::Kind EK) {
-  Edges[EdgeIndexMap.find(&TargetF)->second].setKind(EK);
-}
-
-void LazyCallGraph::Node::removeEdgeInternal(Function &Target) {
-  auto IndexMapI = EdgeIndexMap.find(&Target);
-  assert(IndexMapI != EdgeIndexMap.end() &&
-         "Target not in the edge set for this caller?");
-
-  Edges[IndexMapI->second] = Edge();
-  EdgeIndexMap.erase(IndexMapI);
+void LazyCallGraph::Node::replaceFunction(Function &NewF) {
+  assert(F != &NewF && "Must not replace a function with itself!");
+  F = &NewF;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void LazyCallGraph::Node::dump() const {
   dbgs() << *this << '\n';
 }
 #endif
 
 LazyCallGraph::LazyCallGraph(Module &M) {
   DEBUG(dbgs() << "Building CG for module: " << M.getModuleIdentifier()
                << "\n");
   for (Function &F : M)
-    if (!F.isDeclaration() && !F.hasLocalLinkage())
-      if (EntryIndexMap.insert({&F, EntryEdges.size()}).second) {
-        DEBUG(dbgs() << "  Adding '" << F.getName()
-                     << "' to entry set of the graph.\n");
-        EntryEdges.emplace_back(F, Edge::Ref);
-      }
+    if (!F.isDeclaration() && !F.hasLocalLinkage()) {
+      addEdge(EntryEdges.Edges, EntryEdges.EdgeIndexMap, get(F), Edge::Ref);
 
   // Now add entry nodes for functions reachable via initializers to globals.
   SmallVector<Constant *, 16> Worklist;
   SmallPtrSet<Constant *, 16> Visited;
   for (GlobalVariable &GV : M.globals())
     if (GV.hasInitializer())
       if (Visited.insert(GV.getInitializer()).second)
         Worklist.push_back(GV.getInitializer());
 
   DEBUG(dbgs() << "  Adding functions referenced by global initializers to the "
                   "entry set.\n");
   visitReferences(Worklist, Visited, [&](Function &F) {
-    addEdge(EntryEdges, EntryIndexMap, F, LazyCallGraph::Edge::Ref);
+    addEdge(EntryEdges.Edges, EntryEdges.EdgeIndexMap, get(F),
+            LazyCallGraph::Edge::Ref);
   });
 }
 
 LazyCallGraph::LazyCallGraph(LazyCallGraph &&G)
     : BPA(std::move(G.BPA)), NodeMap(std::move(G.NodeMap)),
-      EntryEdges(std::move(G.EntryEdges)),
-      EntryIndexMap(std::move(G.EntryIndexMap)), SCCBPA(std::move(G.SCCBPA)),
+      EntryEdges(std::move(G.EntryEdges)), SCCBPA(std::move(G.SCCBPA)),
       SCCMap(std::move(G.SCCMap)), LeafRefSCCs(std::move(G.LeafRefSCCs)) {
   updateGraphPtrs();
 }
 
 LazyCallGraph &LazyCallGraph::operator=(LazyCallGraph &&G) {
   BPA = std::move(G.BPA);
   NodeMap = std::move(G.NodeMap);
   EntryEdges = std::move(G.EntryEdges);
-  EntryIndexMap = std::move(G.EntryIndexMap);
   SCCBPA = std::move(G.SCCBPA);
   SCCMap = std::move(G.SCCMap);
   LeafRefSCCs = std::move(G.LeafRefSCCs);
   updateGraphPtrs();
   return *this;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
void LazyCallGraph::SCC::verify() {
   for (Node *N : Nodes) {
     assert(N && "Can't have a null node!");
     assert(OuterRefSCC->G->lookupSCC(*N) == this &&
            "Node does not map to this SCC!");
     assert(N->DFSNumber == -1 &&
            "Must set DFS numbers to -1 when adding a node to an SCC!");
     assert(N->LowLink == -1 &&
            "Must set low link to -1 when adding a node to an SCC!");
-    for (Edge &E : *N)
-      assert(E.getNode() && "Can't have an edge to a raw function!");
+    for (Edge &E : **N)
+      assert(E.getNode() && "Can't have an unpopulated node!");
   }
 }
 #endif
 
 bool LazyCallGraph::SCC::isParentOf(const SCC &C) const {
   if (this == &C)
     return false;
 
   for (Node &N : *this)
-    for (Edge &E : N.calls())
-      if (Node *CalleeN = E.getNode())
-        if (OuterRefSCC->G->lookupSCC(*CalleeN) == &C)
-          return true;
+    for (Edge &E : N->calls())
+      if (OuterRefSCC->G->lookupSCC(E.getNode()) == &C)
 
   // No edges found.
   return false;
 }
 
 bool LazyCallGraph::SCC::isAncestorOf(const SCC &TargetC) const {
   if (this == &TargetC)
     return false;
 
   LazyCallGraph &G = *OuterRefSCC->G;
 
   // Start with this SCC.
   SmallPtrSet<const SCC *, 16> Visited = {this};
   SmallVector<const SCC *, 16> Worklist = {this};
 
   // Walk down the graph until we run out of edges or find a path to TargetC.
   do {
     const SCC &C = *Worklist.pop_back_val();
     for (Node &N : C)
-      for (Edge &E : N.calls()) {
-        Node *CalleeN = E.getNode();
-        if (!CalleeN)
-          continue;
-        SCC *CalleeC = G.lookupSCC(*CalleeN);
+      for (Edge &E : N->calls()) {
+        SCC *CalleeC = G.lookupSCC(E.getNode());
         if (!CalleeC)
           continue;
 
         // If the callee's SCC is the TargetC, we're done.
         if (CalleeC == &TargetC)
           return true;
 
         // If this is the first time we've reached this SCC, put it on the
for (auto &SCCIndexPair : SCCIndices) {
     assert(SCCSet.count(C) && "Found an index for an SCC not in the RefSCC!");
     assert(SCCs[i] == C && "Index doesn't point to SCC!");
   }
 
   // Check that the SCCs are in fact in post-order.
   for (int i = 0, Size = SCCs.size(); i < Size; ++i) {
     SCC &SourceSCC = *SCCs[i];
     for (Node &N : SourceSCC)
-      for (Edge &E : N) {
+      for (Edge &E : *N) {
         if (!E.isCall())
           continue;
-        SCC &TargetSCC = *G->lookupSCC(*E.getNode());
+        SCC &TargetSCC = *G->lookupSCC(E.getNode());
         if (&TargetSCC.getOuterRefSCC() == this) {
           assert(SCCIndices.find(&TargetSCC)->second <= i &&
                  "Edge between SCCs violates post-order relationship.");
           continue;
         }
         assert(TargetSCC.getOuterRefSCC().Parents.count(this) &&
                "Edge to a RefSCC missing us in its parent set.");
       }
   }
 
   // Check that our parents are actually parents.
   for (RefSCC *ParentRC : Parents) {
     assert(ParentRC != this && "Cannot be our own parent!");
     auto HasConnectingEdge = [&] {
       for (SCC &C : *ParentRC)
         for (Node &N : C)
-          for (Edge &E : N)
-            if (G->lookupRefSCC(*E.getNode()) == this)
+          for (Edge &E : *N)
+            if (G->lookupRefSCC(E.getNode()) == this)
               return true;
       return false;
     };
     assert(HasConnectingEdge() && "No edge connects the parent to us!");
   }
 }
 #endif
 
updatePostorderSequenceForEdgeInsertion(
   // because target connects back to source, and we know that all of the SCCs
   // between the source and target in the postorder sequence participate in that
   // cycle.
   return make_range(SCCs.begin() + SourceIdx, SCCs.begin() + TargetIdx);
 }
 
 SmallVector<LazyCallGraph::SCC *, 1>
 LazyCallGraph::RefSCC::switchInternalEdgeToCall(Node &SourceN, Node &TargetN) {
-  assert(!SourceN[TargetN].isCall() && "Must start with a ref edge!");
+  assert(!(*SourceN)[TargetN].isCall() && "Must start with a ref edge!");
   SmallVector<SCC *, 1> DeletedSCCs;
 
 #ifndef NDEBUG
   // In a debug build, verify the RefSCC is valid to start with and when this
   // routine finishes.
   verify();
   auto VerifyOnExit = make_scope_exit([&]() { verify(); });
 #endif
 
   SCC &SourceSCC = *G->lookupSCC(SourceN);
   SCC &TargetSCC = *G->lookupSCC(TargetN);
 
   // If the two nodes are already part of the same SCC, we're also done as
   // we've just added more connectivity.
   if (&SourceSCC == &TargetSCC) {
-    SourceN.setEdgeKind(TargetN.getFunction(), Edge::Call);
+    SourceN->setEdgeKind(TargetN, Edge::Call);
     return DeletedSCCs;
   }
 
   // At this point we leverage the postorder list of SCCs to detect when the
   // insertion of an edge changes the SCC structure in any way.
   //
   // First and foremost, we can eliminate the need for any changes when the
   // edge is toward the beginning of the postorder sequence because all edges
   // flow in that direction already. Thus adding a new one cannot form a cycle.
   int SourceIdx = SCCIndices[&SourceSCC];
   int TargetIdx = SCCIndices[&TargetSCC];
   if (TargetIdx < SourceIdx) {
-    SourceN.setEdgeKind(TargetN.getFunction(), Edge::Call);
+    SourceN->setEdgeKind(TargetN, Edge::Call);
     return DeletedSCCs;
   }
 
   // Compute the SCCs which (transitively) reach the source.
   auto ComputeSourceConnectedSet = [&](SmallPtrSetImpl<SCC *> &ConnectedSet) {
 #ifndef NDEBUG
     // Check that the RefSCC is still valid before computing this as the
     // results will be nonsensical of we've broken its invariants.
     verify();
 #endif
     ConnectedSet.insert(&SourceSCC);
     auto IsConnected = [&](SCC &C) {
       for (Node &N : C)
-        for (Edge &E : N.calls()) {
-          assert(E.getNode() && "Must have formed a node within an SCC!");
-          if (ConnectedSet.count(G->lookupSCC(*E.getNode())))
+        for (Edge &E : N->calls())
+          if (ConnectedSet.count(G->lookupSCC(E.getNode())))
             return true;
-        }
 
       return false;
     };
 
     for (SCC *C :
          make_range(SCCs.begin() + SourceIdx + 1, SCCs.begin() + TargetIdx + 1))
       if (IsConnected(*C))
         ConnectedSet.insert(C);
#ifndef NDEBUG
     verify();
 #endif
     ConnectedSet.insert(&TargetSCC);
     SmallVector<SCC *, 4> Worklist;
     Worklist.push_back(&TargetSCC);
     do {
       SCC &C = *Worklist.pop_back_val();
       for (Node &N : C)
-        for (Edge &E : N) {
-          assert(E.getNode() && "Must have formed a node within an SCC!");
+        for (Edge &E : *N) {
           if (!E.isCall())
             continue;
-          SCC &EdgeC = *G->lookupSCC(*E.getNode());
+          SCC &EdgeC = *G->lookupSCC(E.getNode());
           if (&EdgeC.getOuterRefSCC() != this)
             // Not in this RefSCC...
             continue;
           if (SCCIndices.find(&EdgeC)->second <= SourceIdx)
             // Not in the postorder sequence between source and target.
             continue;
 
           if (ConnectedSet.insert(&EdgeC).second)
#endif
   auto MergeRange = updatePostorderSequenceForEdgeInsertion(
       SourceSCC, TargetSCC, SCCs, SCCIndices, ComputeSourceConnectedSet,
       ComputeTargetConnectedSet);
 
   // If the merge range is empty, then adding the edge didn't actually form any
   // new cycles. We're done.
   if (MergeRange.begin() == MergeRange.end()) {
     // Now that the SCC structure is finalized, flip the kind to call.
-    SourceN.setEdgeKind(TargetN.getFunction(), Edge::Call);
+    SourceN->setEdgeKind(TargetN, Edge::Call);
     return DeletedSCCs;
   }
 
 #ifndef NDEBUG
   // Before merging, check that the RefSCC remains valid after all the
   // postorder updates.
   verify();
 #endif
#endif
   // Erase the merged SCCs from the list and update the indices of the
   // remaining SCCs.
   int IndexOffset = MergeRange.end() - MergeRange.begin();
   auto EraseEnd = SCCs.erase(MergeRange.begin(), MergeRange.end());
   for (SCC *C : make_range(EraseEnd, SCCs.end()))
     SCCIndices[C] -= IndexOffset;
 
   // Now that the SCC structure is finalized, flip the kind to call.
-  SourceN.setEdgeKind(TargetN.getFunction(), Edge::Call);
+  SourceN->setEdgeKind(TargetN, Edge::Call);
 
   // And we're done!
   return DeletedSCCs;
 }
 
 void LazyCallGraph::RefSCC::switchTrivialInternalEdgeToRef(Node &SourceN,
                                                            Node &TargetN) {
-  assert(SourceN[TargetN].isCall() && "Must start with a call edge!");
+  assert((*SourceN)[TargetN].isCall() && "Must start with a call edge!");
 
 #ifndef NDEBUG
   // In a debug build, verify the RefSCC is valid to start with and when this
   // routine finishes.
   verify();
   auto VerifyOnExit = make_scope_exit([&]() { verify(); });
 #endif
 
   assert(G->lookupRefSCC(SourceN) == this &&
          "Source must be in this RefSCC.");
   assert(G->lookupRefSCC(TargetN) == this &&
          "Target must be in this RefSCC.");
   assert(G->lookupSCC(SourceN) != G->lookupSCC(TargetN) &&
          "Source and Target must be in separate SCCs for this to be trivial!");
 
   // Set the edge kind.
-  SourceN.setEdgeKind(TargetN.getFunction(), Edge::Ref);
+  SourceN->setEdgeKind(TargetN, Edge::Ref);
 }
 
 iterator_range<LazyCallGraph::RefSCC::iterator>
 LazyCallGraph::RefSCC::switchInternalEdgeToRef(Node &SourceN, Node &TargetN) {
-  assert(SourceN[TargetN].isCall() && "Must start with a call edge!");
+  assert((*SourceN)[TargetN].isCall() && "Must start with a call edge!");
 
 #ifndef NDEBUG
   // In a debug build, verify the RefSCC is valid to start with and when this
   // routine finishes.
   verify();
   auto VerifyOnExit = make_scope_exit([&]() { verify(); });
 #endif
 
   assert(G->lookupRefSCC(SourceN) == this &&
          "Source must be in this RefSCC.");
   assert(G->lookupRefSCC(TargetN) == this &&
          "Target must be in this RefSCC.");
 
   SCC &TargetSCC = *G->lookupSCC(TargetN);
   assert(G->lookupSCC(SourceN) == &TargetSCC && "Source and Target must be in "
                                                 "the same SCC to require the "
                                                 "full CG update.");
 
   // Set the edge kind.
-  SourceN.setEdgeKind(TargetN.getFunction(), Edge::Ref);
+  SourceN->setEdgeKind(TargetN, Edge::Ref);
 
   // Otherwise we are removing a call edge from a single SCC. This may break
   // the cycle. In order to compute the new set of SCCs, we need to do a small
   // DFS over the nodes within the SCC to form any sub-cycles that remain as
   // distinct SCCs and compute a postorder over the resulting SCCs.
   //
   // However, we specially handle the target node. The target node is known to
   // reach all other nodes in the original SCC by definition. This means that
   // we want the old SCC to be replaced with an SCC contaning that node as it
   // will be the root of whatever SCC DAG results from the DFS. Assumptions
   // about an SCC such as the set of functions called will continue to hold,
   // etc.
 
   SCC &OldSCC = TargetSCC;
-  SmallVector<std::pair<Node *, call_edge_iterator>, 16> DFSStack;
+  SmallVector<std::pair<Node *, EdgeSequence::call_iterator>, 16> DFSStack;
   SmallVector<Node *, 16> PendingSCCStack;
   SmallVector<SCC *, 4> NewSCCs;
 
   // Prepare the nodes for a fresh DFS.
   SmallVector<Node *, 16> Worklist;
   Worklist.swap(OldSCC.Nodes);
   for (Node *N : Worklist) {
     N->DFSNumber = N->LowLink = 0;
if (RootN->DFSNumber != 0) {
       assert(RootN->DFSNumber == -1 &&
              "Shouldn't have any mid-DFS root nodes!");
       continue;
     }
 
     RootN->DFSNumber = RootN->LowLink = 1;
     int NextDFSNumber = 2;
 
-    DFSStack.push_back({RootN, RootN->call_begin()});
+    DFSStack.push_back({RootN, (*RootN)->call_begin()});
     do {
       Node *N;
-      call_edge_iterator I;
+      EdgeSequence::call_iterator I;
       std::tie(N, I) = DFSStack.pop_back_val();
-      auto E = N->call_end();
+      auto E = (*N)->call_end();
       while (I != E) {
-        Node &ChildN = *I->getNode();
+        Node &ChildN = I->getNode();
         if (ChildN.DFSNumber == 0) {
           // We haven't yet visited this child, so descend, pushing the current
           // node onto the stack.
           DFSStack.push_back({N, I});
 
           assert(!G->SCCMap.count(&ChildN) &&
                  "Found a node with 0 DFS number but already in an SCC!");
           ChildN.DFSNumber = ChildN.LowLink = NextDFSNumber++;
           N = &ChildN;
-          I = N->call_begin();
-          E = N->call_end();
+          I = (*N)->call_begin();
+          E = (*N)->call_end();
           continue;
         }
 
         // Check for the child already being part of some component.
         if (ChildN.DFSNumber == -1) {
           if (G->lookupSCC(ChildN) == &OldSCC) {
             // If the child is part of the old SCC, we know that it can reach
             // every other node, so we have formed a cycle. Pull the entire DFS
for (int Idx = OldIdx, Size = SCCs.size(); Idx < Size; ++Idx)
     SCCIndices[SCCs[Idx]] = Idx;
 
   return make_range(SCCs.begin() + OldIdx,
                     SCCs.begin() + OldIdx + NewSCCs.size());
 }
 
 void LazyCallGraph::RefSCC::switchOutgoingEdgeToCall(Node &SourceN,
                                                      Node &TargetN) {
-  assert(!SourceN[TargetN].isCall() && "Must start with a ref edge!");
+  assert(!(*SourceN)[TargetN].isCall() && "Must start with a ref edge!");
 
   assert(G->lookupRefSCC(SourceN) == this && "Source must be in this RefSCC.");
   assert(G->lookupRefSCC(TargetN) != this &&
          "Target must not be in this RefSCC.");
 #ifdef EXPENSIVE_CEHCKS
   assert(G->lookupRefSCC(TargetN)->isDescendantOf(*this) &&
          "Target must be a descendant of the Source.");
 #endif
 
   // Edges between RefSCCs are the same regardless of call or ref, so we can
   // just flip the edge here.
-  SourceN.setEdgeKind(TargetN.getFunction(), Edge::Call);
+  SourceN->setEdgeKind(TargetN, Edge::Call);
 
 #ifndef NDEBUG
   // Check that the RefSCC is still valid.
   verify();
 #endif
 }
 
 void LazyCallGraph::RefSCC::switchOutgoingEdgeToRef(Node &SourceN,
                                                     Node &TargetN) {
-  assert(SourceN[TargetN].isCall() && "Must start with a call edge!");
+  assert((*SourceN)[TargetN].isCall() && "Must start with a call edge!");
 
   assert(G->lookupRefSCC(SourceN) == this && "Source must be in this RefSCC.");
   assert(G->lookupRefSCC(TargetN) != this &&
          "Target must not be in this RefSCC.");
 #ifdef EXPENSIVE_CEHCKS
   assert(G->lookupRefSCC(TargetN)->isDescendantOf(*this) &&
          "Target must be a descendant of the Source.");
 #endif
 
   // Edges between RefSCCs are the same regardless of call or ref, so we can
   // just flip the edge here.
-  SourceN.setEdgeKind(TargetN.getFunction(), Edge::Ref);
+  SourceN->setEdgeKind(TargetN, Edge::Ref);
 
 #ifndef NDEBUG
   // Check that the RefSCC is still valid.
   verify();
 #endif
 }
 
 void LazyCallGraph::RefSCC::insertInternalRefEdge(Node &SourceN,
                                                   Node &TargetN) {
   assert(G->lookupRefSCC(SourceN) == this && "Source must be in this RefSCC.");
   assert(G->lookupRefSCC(TargetN) == this && "Target must be in this RefSCC.");
 
-  SourceN.insertEdgeInternal(TargetN, Edge::Ref);
+  SourceN->insertEdgeInternal(TargetN, Edge::Ref);
 
 #ifndef NDEBUG
   // Check that the RefSCC is still valid.
   verify();
 #endif
 }
 
 void LazyCallGraph::RefSCC::insertOutgoingEdge(Node &SourceN, Node &TargetN,
                                                Edge::Kind EK) {
   // First insert it into the caller.
-  SourceN.insertEdgeInternal(TargetN, EK);
+  SourceN->insertEdgeInternal(TargetN, EK);
 
   assert(G->lookupRefSCC(SourceN) == this && "Source must be in this RefSCC.");
 
   RefSCC &TargetC = *G->lookupRefSCC(TargetN);
   assert(&TargetC != this && "Target must not be in this RefSCC.");
 #ifdef EXPENSIVE_CEHCKS
   assert(TargetC.isDescendantOf(*this) &&
          "Target must be a descendant of the Source.");
#endif
   auto ComputeTargetConnectedSet = [&](SmallPtrSetImpl<RefSCC *> &Set) {
     Set.insert(this);
     SmallVector<RefSCC *, 4> Worklist;
     Worklist.push_back(this);
     do {
       RefSCC &RC = *Worklist.pop_back_val();
       for (SCC &C : RC)
         for (Node &N : C)
-          for (Edge &E : N) {
-            assert(E.getNode() && "Must have formed a node!");
-            RefSCC &EdgeRC = *G->lookupRefSCC(*E.getNode());
+          for (Edge &E : *N) {
+            RefSCC &EdgeRC = *G->lookupRefSCC(E.getNode());
             if (G->getRefSCCIndex(EdgeRC) <= SourceIdx)
               // Not in the postorder sequence between source and target.
               continue;
 
             if (Set.insert(&EdgeRC).second)
               Worklist.push_back(&EdgeRC);
           }
     } while (!Worklist.empty());
for (RefSCC *RC : MergeRange) {
     // update any parent sets.
     // FIXME: We should try to find a way to avoid this (rather expensive) edge
     // walk by updating the parent sets in some other manner.
     for (SCC &InnerC : *RC) {
       InnerC.OuterRefSCC = this;
       SCCIndices[&InnerC] = SCCIndex++;
       for (Node &N : InnerC) {
         G->SCCMap[&N] = &InnerC;
-        for (Edge &E : N) {
-          assert(E.getNode() &&
-                 "Cannot have a null node within a visited SCC!");
-          RefSCC &ChildRC = *G->lookupRefSCC(*E.getNode());
+        for (Edge &E : *N) {
+          RefSCC &ChildRC = *G->lookupRefSCC(E.getNode());
           if (MergeSet.count(&ChildRC))
             continue;
           ChildRC.Parents.erase(RC);
           ChildRC.Parents.insert(this);
         }
       }
     }
 
#endif
   int IndexOffset = MergeRange.end() - MergeRange.begin();
   auto EraseEnd =
       G->PostOrderRefSCCs.erase(MergeRange.begin(), MergeRange.end());
   for (RefSCC *RC : make_range(EraseEnd, G->PostOrderRefSCCs.end()))
     G->RefSCCIndices[RC] -= IndexOffset;
 
   // At this point we have a merged RefSCC with a post-order SCCs list, just
   // connect the nodes to form the new edge.
-  SourceN.insertEdgeInternal(TargetN, Edge::Ref);
+  SourceN->insertEdgeInternal(TargetN, Edge::Ref);
 
   // We return the list of SCCs which were merged so that callers can
   // invalidate any data they have associated with those SCCs. Note that these
   // SCCs are no longer in an interesting state (they are totally empty) but
   // the pointers will remain stable for the life of the graph itself.
   return DeletedRefSCCs;
 }
 

 #ifndef NDEBUG
   // In a debug build, verify the RefSCC is valid to start with and when this
   // routine finishes.
   verify();
   auto VerifyOnExit = make_scope_exit([&]() { verify(); });
 #endif
 
   // First remove it from the node.
-  SourceN.removeEdgeInternal(TargetN.getFunction());
+  bool Removed = SourceN->removeEdgeInternal(TargetN);
+  (void)Removed;
+  assert(Removed && "Target not in the edge set for this caller?");
 
   bool HasOtherEdgeToChildRC = false;
   bool HasOtherChildRC = false;
   for (SCC *InnerC : SCCs) {
     for (Node &N : *InnerC) {
-      for (Edge &E : N) {
-        assert(E.getNode() && "Cannot have a missing node in a visited SCC!");
-        RefSCC &OtherChildRC = *G->lookupRefSCC(*E.getNode());
+      for (Edge &E : *N) {
+        RefSCC &OtherChildRC = *G->lookupRefSCC(E.getNode());
         if (&OtherChildRC == &TargetRC) {
           HasOtherEdgeToChildRC = true;
           break;
         }
         if (&OtherChildRC != this)
           HasOtherChildRC = true;
       }
       if (HasOtherEdgeToChildRC)
if (!HasOtherEdgeToChildRC) {
     // It may make the Source SCC a leaf SCC.
     if (!HasOtherChildRC)
       G->LeafRefSCCs.push_back(this);
   }
 }
 
 SmallVector<LazyCallGraph::RefSCC *, 1>
 LazyCallGraph::RefSCC::removeInternalRefEdge(Node &SourceN, Node &TargetN) {
-  assert(!SourceN[TargetN].isCall() &&
+  assert(!(*SourceN)[TargetN].isCall() &&
          "Cannot remove a call edge, it must first be made a ref edge");
 
 #ifndef NDEBUG
   // In a debug build, verify the RefSCC is valid to start with and when this
   // routine finishes.
   verify();
   auto VerifyOnExit = make_scope_exit([&]() { verify(); });
 #endif
 
   // First remove the actual edge.
-  SourceN.removeEdgeInternal(TargetN.getFunction());
+  bool Removed = SourceN->removeEdgeInternal(TargetN);
+  (void)Removed;
+  assert(Removed && "Target not in the edge set for this caller?");
 
   // We return a list of the resulting *new* RefSCCs in post-order.
   SmallVector<RefSCC *, 1> Result;
 
   // Direct recursion doesn't impact the SCC graph at all.
   if (&SourceN == &TargetN)
     return Result;
 
for (SCC *C : SCCs) {
     Worklist.append(C->Nodes.begin(), C->Nodes.end());
   }
 
   auto MarkNodeForSCCNumber = [&PostOrderMapping](Node &N, int Number) {
     N.DFSNumber = N.LowLink = -1;
     PostOrderMapping[&N] = Number;
   };
 
-  SmallVector<std::pair<Node *, edge_iterator>, 4> DFSStack;
+  SmallVector<std::pair<Node *, EdgeSequence::iterator>, 4> DFSStack;
   SmallVector<Node *, 4> PendingRefSCCStack;
   do {
     assert(DFSStack.empty() &&
            "Cannot begin a new root with a non-empty DFS stack!");
     assert(PendingRefSCCStack.empty() &&
            "Cannot begin a new root with pending nodes for an SCC!");
 
     Node *RootN = Worklist.pop_back_val();
     // Skip any nodes we've already reached in the DFS.
     if (RootN->DFSNumber != 0) {
       assert(RootN->DFSNumber == -1 &&
              "Shouldn't have any mid-DFS root nodes!");
       continue;
     }
 
     RootN->DFSNumber = RootN->LowLink = 1;
     int NextDFSNumber = 2;
 
-    DFSStack.push_back({RootN, RootN->begin()});
+    DFSStack.push_back({RootN, (*RootN)->begin()});
     do {
       Node *N;
-      edge_iterator I;
+      EdgeSequence::iterator I;
       std::tie(N, I) = DFSStack.pop_back_val();
-      auto E = N->end();
+      auto E = (*N)->end();
 
       assert(N->DFSNumber != 0 && "We should always assign a DFS number "
                                   "before processing a node.");
 
       while (I != E) {
-        Node &ChildN = I->getNode(*G);
+        Node &ChildN = I->getNode();
         if (ChildN.DFSNumber == 0) {
           // Mark that we should start at this child when next this node is the
           // top of the stack. We don't start at the next child to ensure this
           // child's lowlink is reflected.
           DFSStack.push_back({N, I});
 
           // Continue, resetting to the child node.
           ChildN.LowLink = ChildN.DFSNumber = NextDFSNumber++;
           N = &ChildN;
-          I = ChildN.begin();
-          E = ChildN.end();
+          I = ChildN->begin();
+          E = ChildN->end();
           continue;
         }
         if (ChildN.DFSNumber == -1) {
           // Check if this edge's target node connects to the deleted edge's
           // target node. If so, we know that every node connected will end up
           // in this RefSCC, so collapse the entire current stack into the root
           // slot in our SCC numbering. See above for the motivation of
           // optimizing the target connected nodes in this way.

 
 #ifndef NDEBUG
   // Now we need to reconnect the current (root) SCC to the graph. We do this
   // manually because we can special case our leaf handling and detect errors.
   bool IsLeaf = true;
 #endif
   for (SCC *C : SCCs)
     for (Node &N : *C) {
-      for (Edge &E : N) {
-        assert(E.getNode() && "Cannot have a missing node in a visited SCC!");
-        RefSCC &ChildRC = *G->lookupRefSCC(*E.getNode());
+      for (Edge &E : *N) {
+        RefSCC &ChildRC = *G->lookupRefSCC(E.getNode());
         if (&ChildRC == this)
           continue;
         ChildRC.Parents.insert(this);
 #ifndef NDEBUG
         IsLeaf = false;
 #endif
       }
     }
 #ifndef NDEBUG
   if (!Result.empty())
     assert(!IsLeaf && "This SCC cannot be a leaf as we have split out new "
                       "SCCs by removing this edge.");
   if (none_of(G->LeafRefSCCs, [&](RefSCC *C) { return C == this; }))
     assert(!IsLeaf && "This SCC cannot be a leaf as it already had child "
                       "SCCs before we removed this edge.");
 #endif
   // And connect both this RefSCC and all the new ones to the correct parents.
   // The easiest way to do this is just to re-analyze the old parent set.
   SmallVector<RefSCC *, 4> OldParents(Parents.begin(), Parents.end());
   Parents.clear();
   for (RefSCC *ParentRC : OldParents)
     for (SCC &ParentC : *ParentRC)
       for (Node &ParentN : ParentC)
-        for (Edge &E : ParentN) {
-          assert(E.getNode() && "Cannot have a missing node in a visited SCC!");
-          RefSCC &RC = *G->lookupRefSCC(*E.getNode());
+        for (Edge &E : *ParentN) {
+          RefSCC &RC = *G->lookupRefSCC(E.getNode());
           if (&RC != ParentRC)
             RC.Parents.insert(ParentRC);
         }
 
   // If this SCC stopped being a leaf through this edge removal, remove it from
   // the leaf SCC list. Note that this DTRT in the case where this was never
   // a leaf.
   // FIXME: As LeafRefSCCs could be very large, we might want to not walk the
#ifdef EXPENSIVE_CHECKS
   SCC &TargetC = *G->lookupSCC(TargetN);
   if (&SourceC != &TargetC)
     assert(SourceC.isAncestorOf(TargetC) &&
            "Call edge is not trivial in the SCC graph!");
 #endif // EXPENSIVE_CHECKS
 #endif // NDEBUG
 
   // First insert it into the source or find the existing edge.
-  auto InsertResult = SourceN.EdgeIndexMap.insert(
-      {&TargetN.getFunction(), SourceN.Edges.size()});
+  auto InsertResult =
+      SourceN->EdgeIndexMap.insert({&TargetN, SourceN->Edges.size()});
   if (!InsertResult.second) {
     // Already an edge, just update it.
-    Edge &E = SourceN.Edges[InsertResult.first->second];
+    Edge &E = SourceN->Edges[InsertResult.first->second];
     if (E.isCall())
       return; // Nothing to do!
     E.setKind(Edge::Call);
   } else {
     // Create the new edge.
-    SourceN.Edges.emplace_back(TargetN, Edge::Call);
+    SourceN->Edges.emplace_back(TargetN, Edge::Call);
   }
 
   // Now that we have the edge, handle the graph fallout.
   handleTrivialEdgeInsertion(SourceN, TargetN);
 }
 
 void LazyCallGraph::RefSCC::insertTrivialRefEdge(Node &SourceN, Node &TargetN) {
 #ifndef NDEBUG
   // Check that the RefSCC is still valid when we finish.
   auto ExitVerifier = make_scope_exit([this] { verify(); });
 
 #ifdef EXPENSIVE_CHECKS
   // Check that we aren't breaking some invariants of the RefSCC graph.
   RefSCC &SourceRC = *G->lookupRefSCC(SourceN);
   RefSCC &TargetRC = *G->lookupRefSCC(TargetN);
   if (&SourceRC != &TargetRC)
     assert(SourceRC.isAncestorOf(TargetRC) &&
            "Ref edge is not trivial in the RefSCC graph!");
 #endif // EXPENSIVE_CHECKS
 #endif // NDEBUG
 
   // First insert it into the source or find the existing edge.
-  auto InsertResult = SourceN.EdgeIndexMap.insert(
-      {&TargetN.getFunction(), SourceN.Edges.size()});
+  auto InsertResult =
+      SourceN->EdgeIndexMap.insert({&TargetN, SourceN->Edges.size()});
   if (!InsertResult.second)
     // Already an edge, we're done.
     return;
 
   // Create the new edge.
-  SourceN.Edges.emplace_back(TargetN, Edge::Ref);
+  SourceN->Edges.emplace_back(TargetN, Edge::Ref);
 
   // Now that we have the edge, handle the graph fallout.
   handleTrivialEdgeInsertion(SourceN, TargetN);
 }
 
-void LazyCallGraph::insertEdge(Node &SourceN, Function &Target, Edge::Kind EK) {
+void LazyCallGraph::RefSCC::replaceNodeFunction(Node &N, Function &NewF) {
+  Function &OldF = N.getFunction();
+
+#ifndef NDEBUG
+  // Check that the RefSCC is still valid when we finish.
+  auto ExitVerifier = make_scope_exit([this] { verify(); });
+
+  assert(G->lookupRefSCC(N) == this &&
+         "Cannot replace the function of a node outside this RefSCC.");
+
+  assert(G->NodeMap.find(&NewF) == G->NodeMap.end() &&
+         "Must not have already walked the new function!'");
+
+  // It is important that this replacement not introduce graph changes so we
+  // insist that the caller has already removed every use of the original
+  // function and that all uses of the new function correspond to existing
+  // edges in the graph. The common and expected way to use this is when
+  // replacing the function itself in the IR without changing the call graph
+  // shape and just updating the analysis based on that.
+  assert(&OldF != &NewF && "Cannot replace a function with itself!");
+  assert(OldF.use_empty() &&
+         "Must have moved all uses from the old function to the new!");
+#endif
+
+  N.replaceFunction(NewF);
+
+  // Update various call graph maps.
+  G->NodeMap.erase(&OldF);
+  G->NodeMap[&NewF] = &N;
+}
+
+void LazyCallGraph::insertEdge(Node &SourceN, Node &TargetN, Edge::Kind EK) {
   assert(SCCMap.empty() &&
          "This method cannot be called after SCCs have been formed!");
 
-  return SourceN.insertEdgeInternal(Target, EK);
+  return SourceN->insertEdgeInternal(TargetN, EK);
 }
 
-void LazyCallGraph::removeEdge(Node &SourceN, Function &Target) {
+void LazyCallGraph::removeEdge(Node &SourceN, Node &TargetN) {
   assert(SCCMap.empty() &&
          "This method cannot be called after SCCs have been formed!");
 
-  return SourceN.removeEdgeInternal(Target);
+  bool Removed = SourceN->removeEdgeInternal(TargetN);
+  (void)Removed;
+  assert(Removed && "Target not in the edge set for this caller?");
 }
 
 void LazyCallGraph::removeDeadFunction(Function &F) {
   // FIXME: This is unnecessarily restrictive. We should be able to remove
   // functions which recursively call themselves.
   assert(F.use_empty() &&
          "This routine should only be called on trivially dead functions!");
 
-  auto EII = EntryIndexMap.find(&F);
-  if (EII != EntryIndexMap.end()) {
-    EntryEdges[EII->second] = Edge();
-    EntryIndexMap.erase(EII);
-  }
-
   auto NI = NodeMap.find(&F);
   if (NI == NodeMap.end())
     // Not in the graph at all!
     return;
 
   Node &N = *NI->second;
   NodeMap.erase(NI);
 
+  // Remove this from the entry edges if present.
+  EntryEdges.removeEdgeInternal(N);
+
   if (SCCMap.empty()) {
     // No SCCs have been formed, so removing this is fine and there is nothing
     // else necessary at this point but clearing out the node.
     N.clear();
     return;
   }
 
   // Cannot remove a function which has yet to be visited in the DFS walk, so
   // if we have a node at all then we must have an SCC and RefSCC.
   auto CI = SCCMap.find(&N);
   assert(CI != SCCMap.end() &&
          "Tried to remove a node without an SCC after DFS walk started!");
   SCC &C = *CI->second;
   SCCMap.erase(CI);
   RefSCC &RC = C.getOuterRefSCC();
 
   // This node must be the only member of its SCC as it has no callers, and
   // that SCC must be the only member of a RefSCC as it has no references.
   // Validate these properties first.
   assert(C.size() == 1 && "Dead functions must be in a singular SCC");
   assert(RC.size() == 1 && "Dead functions must be in a singular RefSCC");
   assert(RC.Parents.empty() && "Cannot have parents of a dead RefSCC!");
 
   // Now remove this RefSCC from any parents sets and the leaf list.
-  for (Edge &E : N)
-    if (Node *TargetN = E.getNode())
-      if (RefSCC *TargetRC = lookupRefSCC(*TargetN))
-        TargetRC->Parents.erase(&RC);
+  for (Edge &E : *N)
+    if (RefSCC *TargetRC = lookupRefSCC(E.getNode()))
   // FIXME: This is a linear operation which could become hot and benefit from
   // an index map.
   auto LRI = find(LeafRefSCCs, &RC);
   if (LRI != LeafRefSCCs.end())
     LeafRefSCCs.erase(LRI);
 
   auto RCIndexI = RefSCCIndices.find(&RC);
   int RCIndex = RCIndexI->second;
LazyCallGraph::Node &LazyCallGraph::insertInto(Function &F, Node *&MappedN) {
   return *new (MappedN = BPA.Allocate()) Node(*this, F);
 }
 
 void LazyCallGraph::updateGraphPtrs() {
   // Process all nodes updating the graph pointers.
   {
     SmallVector<Node *, 16> Worklist;
     for (Edge &E : EntryEdges)
-      if (Node *EntryN = E.getNode())
-        Worklist.push_back(EntryN);
+      Worklist.push_back(&E.getNode());
 
     while (!Worklist.empty()) {
-      Node *N = Worklist.pop_back_val();
-      N->G = this;
-      for (Edge &E : N->Edges)
-        if (Node *TargetN = E.getNode())
-          Worklist.push_back(TargetN);
+      Node &N = *Worklist.pop_back_val();
+      N.G = this;
+      if (N)
+        for (Edge &E : *N)
+          Worklist.push_back(&E.getNode());
     }
   }
 
   // Process all SCCs updating the graph pointers.
   {
     SmallVector<RefSCC *, 16> Worklist(LeafRefSCCs.begin(), LeafRefSCCs.end());
 
     while (!Worklist.empty()) {
for (Node *N : Nodes) {
     // This node will go into the next RefSCC, clear out its DFS and low link
     // as we scan.
     N->DFSNumber = N->LowLink = 0;
   }
 
   // Each RefSCC contains a DAG of the call SCCs. To build these, we do
   // a direct walk of the call edges using Tarjan's algorithm. We reuse the
   // internal storage as we won't need it for the outer graph's DFS any longer.
-  buildGenericSCCs(Nodes, [](Node &N) { return N.call_begin(); },
-                   [](Node &N) { return N.call_end(); },
-                   [](call_edge_iterator I) -> Node & {
-                     // For SCCs, all the nodes should already be formed.
-                     return *I->getNode();
-                   },
-                   [this, &RC](node_stack_range Nodes) {
-                     RC.SCCs.push_back(createSCC(RC, Nodes));
-                     for (Node &N : *RC.SCCs.back()) {
-                       N.DFSNumber = N.LowLink = -1;
-                       SCCMap[&N] = RC.SCCs.back();
-                     }
-                   });
+  buildGenericSCCs(
+      Nodes, [](Node &N) { return N->call_begin(); },
+      [](Node &N) { return N->call_end(); },
+      [](EdgeSequence::call_iterator I) -> Node & { return I->getNode(); },
 
   // Wire up the SCC indices.
   for (int i = 0, Size = RC.SCCs.size(); i < Size; ++i)
     RC.SCCIndices[RC.SCCs[i]] = i;
 }
 
 void LazyCallGraph::buildRefSCCs() {
   if (EntryEdges.empty() || !PostOrderRefSCCs.empty())
     // RefSCCs are either non-existent or already built!
     return;
 
   assert(RefSCCIndices.empty() && "Already mapped RefSCC indices!");
 
   SmallVector<Node *, 16> Roots;
   for (Edge &E : *this)
-    Roots.push_back(&E.getNode(*this));
+    Roots.push_back(&E.getNode());
 
   // The roots will be popped of a stack, so use reverse to get a less
   // surprising order. This doesn't change any of the semantics anywhere.
   std::reverse(Roots.begin(), Roots.end());
 
   buildGenericSCCs(
-      Roots, [](Node &N) { return N.begin(); }, [](Node &N) { return N.end(); },
-      [this](edge_iterator I) -> Node & {
-        // Form the node if we haven't yet.
-        return I->getNode(*this);
+      Roots,
+      [](Node &N) {
+        // We need to populate each node as we begin to walk its edges.
+        N.populate();
+        return N->begin();
       },
+      [](Node &N) { return N->end(); },
+      [](EdgeSequence::iterator I) -> Node & { return I->getNode(); },
       [this](node_stack_range Nodes) {
         RefSCC *NewRC = createRefSCC(*this);
         buildSCCs(*NewRC, Nodes);
         connectRefSCC(*NewRC);
 
         // Push the new node into the postorder list and remember its position
         // in the index map.
         bool Inserted =

 // tracking into their DFS in order to remove a full walk of all edges.
 void LazyCallGraph::connectRefSCC(RefSCC &RC) {
   // Walk all edges in the RefSCC (this remains linear as we only do this once
   // when we build the RefSCC) to connect it to the parent sets of its
   // children.
   bool IsLeaf = true;
   for (SCC &C : RC)
     for (Node &N : C)
-      for (Edge &E : N) {
-        assert(E.getNode() &&
-               "Cannot have a missing node in a visited part of the graph!");
-        RefSCC &ChildRC = *lookupRefSCC(*E.getNode());
+      for (Edge &E : *N) {
+        RefSCC &ChildRC = *lookupRefSCC(E.getNode());
         if (&ChildRC == &RC)
           continue;
         ChildRC.Parents.insert(&RC);
         IsLeaf = false;
       }
 
   // For the SCCs where we find no child SCCs, add them to the leaf list.
   if (IsLeaf)
     LeafRefSCCs.push_back(&RC);
 }
 
 AnalysisKey LazyCallGraphAnalysis::Key;
 
 LazyCallGraphPrinterPass::LazyCallGraphPrinterPass(raw_ostream &OS) : OS(OS) {}
 
 static void printNode(raw_ostream &OS, LazyCallGraph::Node &N) {
   OS << "  Edges in function: " << N.getFunction().getName() << "\n";
-  for (const LazyCallGraph::Edge &E : N)
+  for (LazyCallGraph::Edge &E : N.populate())
     OS << "    " << (E.isCall() ? "call" : "ref ") << " -> "
        << E.getFunction().getName() << "\n";
 
   OS << "\n";
 }
 
 static void printSCC(raw_ostream &OS, LazyCallGraph::SCC &C) {
   ptrdiff_t Size = std::distance(C.begin(), C.end());

 }
 
 LazyCallGraphDOTPrinterPass::LazyCallGraphDOTPrinterPass(raw_ostream &OS)
     : OS(OS) {}
 
 static void printNodeDOT(raw_ostream &OS, LazyCallGraph::Node &N) {
   std::string Name = "\"" + DOT::EscapeString(N.getFunction().getName()) + "\"";
 
-  for (const LazyCallGraph::Edge &E : N) {
+  for (LazyCallGraph::Edge &E : N.populate()) {
     OS << "  " << Name << " -> \""
        << DOT::EscapeString(E.getFunction().getName()) << "\"";
     if (!E.isCall()) // It is a ref edge.
       OS << " [style=dashed,label=\"ref\"]";
     OS << ";\n";
   }
 
   OS << "\n";

llvm/trunk/lib/Transforms/IPO/Inliner.cpp

do {
     // chain. While in some cases these edges are direct calls inside the
     // callee, they have to be modeled in the inliner as reference edges as
     // there may be a reference edge anywhere along the chain from the current
     // caller to the callee that causes the whole thing to appear like
     // a (transitive) reference edge that will require promotion to a call edge
     // below.
     for (Function *InlinedCallee : InlinedCallees) {
       LazyCallGraph::Node &CalleeN = *CG.lookup(*InlinedCallee);
-      for (LazyCallGraph::Edge &E : CalleeN)
-        RC->insertTrivialRefEdge(N, *E.getNode());
+      for (LazyCallGraph::Edge &E : *CalleeN)
+        RC->insertTrivialRefEdge(N, E.getNode());
     }
     InlinedCallees.clear();
 
     // At this point, since we have made changes we have at least removed
     // a call instruction. However, in the process we do some incremental
     // simplification of the surrounding code. This simplification can
     // essentially do all of the same things as a function pass and we can
     // re-use the exact same logic for updating the call graph to reflect the

llvm/trunk/unittests/Analysis/LazyCallGraphTest.cpp

TEST(LazyCallGraphTest, BasicGraphFormation) {
   LLVMContext Context;
   std::unique_ptr<Module> M = parseAssembly(Context, DiamondOfTriangles);
   LazyCallGraph CG(*M);
 
   // The order of the entry nodes should be stable w.r.t. the source order of
   // the IR, and everything in our module is an entry node, so just directly
   // build variables for each node.
   auto I = CG.begin();
-  LazyCallGraph::Node &A1 = (I++)->getNode(CG);
+  LazyCallGraph::Node &A1 = (I++)->getNode();
   EXPECT_EQ("a1", A1.getFunction().getName());
-  LazyCallGraph::Node &A2 = (I++)->getNode(CG);
+  LazyCallGraph::Node &A2 = (I++)->getNode();
   EXPECT_EQ("a2", A2.getFunction().getName());
-  LazyCallGraph::Node &A3 = (I++)->getNode(CG);
+  LazyCallGraph::Node &A3 = (I++)->getNode();
   EXPECT_EQ("a3", A3.getFunction().getName());
-  LazyCallGraph::Node &B1 = (I++)->getNode(CG);
+  LazyCallGraph::Node &B1 = (I++)->getNode();
   EXPECT_EQ("b1", B1.getFunction().getName());
-  LazyCallGraph::Node &B2 = (I++)->getNode(CG);
+  LazyCallGraph::Node &B2 = (I++)->getNode();
   EXPECT_EQ("b2", B2.getFunction().getName());
-  LazyCallGraph::Node &B3 = (I++)->getNode(CG);
+  LazyCallGraph::Node &B3 = (I++)->getNode();
   EXPECT_EQ("b3", B3.getFunction().getName());
-  LazyCallGraph::Node &C1 = (I++)->getNode(CG);
+  LazyCallGraph::Node &C1 = (I++)->getNode();
   EXPECT_EQ("c1", C1.getFunction().getName());
-  LazyCallGraph::Node &C2 = (I++)->getNode(CG);
+  LazyCallGraph::Node &C2 = (I++)->getNode();
   EXPECT_EQ("c2", C2.getFunction().getName());
-  LazyCallGraph::Node &C3 = (I++)->getNode(CG);
+  LazyCallGraph::Node &C3 = (I++)->getNode();
   EXPECT_EQ("c3", C3.getFunction().getName());
-  LazyCallGraph::Node &D1 = (I++)->getNode(CG);
+  LazyCallGraph::Node &D1 = (I++)->getNode();
   EXPECT_EQ("d1", D1.getFunction().getName());
-  LazyCallGraph::Node &D2 = (I++)->getNode(CG);
+  LazyCallGraph::Node &D2 = (I++)->getNode();
   EXPECT_EQ("d2", D2.getFunction().getName());
-  LazyCallGraph::Node &D3 = (I++)->getNode(CG);
+  LazyCallGraph::Node &D3 = (I++)->getNode();
   EXPECT_EQ("d3", D3.getFunction().getName());
   EXPECT_EQ(CG.end(), I);
 
   // Build vectors and sort them for the rest of the assertions to make them
   // independent of order.
   std::vector<std::string> Nodes;
 
-  for (LazyCallGraph::Edge &E : A1)
+  for (LazyCallGraph::Edge &E : A1.populate())
     Nodes.push_back(E.getFunction().getName());
   std::sort(Nodes.begin(), Nodes.end());
   EXPECT_EQ("a2", Nodes[0]);
   EXPECT_EQ("b2", Nodes[1]);
   EXPECT_EQ("c3", Nodes[2]);
   Nodes.clear();
 
-  EXPECT_EQ(A2.end(), std::next(A2.begin()));
-  EXPECT_EQ("a3", A2.begin()->getFunction().getName());
-  EXPECT_EQ(A3.end(), std::next(A3.begin()));
-  EXPECT_EQ("a1", A3.begin()->getFunction().getName());
+  A2.populate();
+  EXPECT_EQ(A2->end(), std::next(A2->begin()));
+  EXPECT_EQ("a3", A2->begin()->getFunction().getName());
+  A3.populate();
+  EXPECT_EQ(A3->end(), std::next(A3->begin()));
+  EXPECT_EQ("a1", A3->begin()->getFunction().getName());
 
-  for (LazyCallGraph::Edge &E : B1)
+  for (LazyCallGraph::Edge &E : B1.populate())
     Nodes.push_back(E.getFunction().getName());
   std::sort(Nodes.begin(), Nodes.end());
   EXPECT_EQ("b2", Nodes[0]);
   EXPECT_EQ("d3", Nodes[1]);
   Nodes.clear();
 
-  EXPECT_EQ(B2.end(), std::next(B2.begin()));
-  EXPECT_EQ("b3", B2.begin()->getFunction().getName());
-  EXPECT_EQ(B3.end(), std::next(B3.begin()));
-  EXPECT_EQ("b1", B3.begin()->getFunction().getName());
+  B2.populate();
+  EXPECT_EQ(B2->end(), std::next(B2->begin()));
+  EXPECT_EQ("b3", B2->begin()->getFunction().getName());
+  B3.populate();
+  EXPECT_EQ(B3->end(), std::next(B3->begin()));
+  EXPECT_EQ("b1", B3->begin()->getFunction().getName());
 
-  for (LazyCallGraph::Edge &E : C1)
+  for (LazyCallGraph::Edge &E : C1.populate())
     Nodes.push_back(E.getFunction().getName());
   std::sort(Nodes.begin(), Nodes.end());
   EXPECT_EQ("c2", Nodes[0]);
   EXPECT_EQ("d2", Nodes[1]);
   Nodes.clear();
 
-  EXPECT_EQ(C2.end(), std::next(C2.begin()));
-  EXPECT_EQ("c3", C2.begin()->getFunction().getName());
-  EXPECT_EQ(C3.end(), std::next(C3.begin()));
-  EXPECT_EQ("c1", C3.begin()->getFunction().getName());
-
-  EXPECT_EQ(D1.end(), std::next(D1.begin()));
-  EXPECT_EQ("d2", D1.begin()->getFunction().getName());
-  EXPECT_EQ(D2.end(), std::next(D2.begin()));
-  EXPECT_EQ("d3", D2.begin()->getFunction().getName());
-  EXPECT_EQ(D3.end(), std::next(D3.begin()));
-  EXPECT_EQ("d1", D3.begin()->getFunction().getName());
+  C2.populate();
+  EXPECT_EQ(C2->end(), std::next(C2->begin()));
+  EXPECT_EQ("c3", C2->begin()->getFunction().getName());
+  C3.populate();
+  EXPECT_EQ(C3->end(), std::next(C3->begin()));
+  EXPECT_EQ("c1", C3->begin()->getFunction().getName());
+
+  D1.populate();
+  EXPECT_EQ(D1->end(), std::next(D1->begin()));
+  EXPECT_EQ("d2", D1->begin()->getFunction().getName());
+  D2.populate();
+  EXPECT_EQ(D2->end(), std::next(D2->begin()));
+  EXPECT_EQ("d3", D2->begin()->getFunction().getName());
+  D3.populate();
+  EXPECT_EQ(D3->end(), std::next(D3->begin()));
+  EXPECT_EQ("d1", D3->begin()->getFunction().getName());
 
   // Now lets look at the RefSCCs and SCCs.
   CG.buildRefSCCs();
   auto J = CG.postorder_ref_scc_begin();
 
   LazyCallGraph::RefSCC &D = *J++;
   ASSERT_EQ(1, D.size());
   for (LazyCallGraph::Node &N : *D.begin())
std::unique_ptr<Module> M = parseAssembly(Context, "define void @a() {\n"
                                                      "define void @c() {\n"
                                                      "entry:\n"
                                                      "  ret void\n"
                                                      "}\n");
   LazyCallGraph CG(*M);
 
   LazyCallGraph::Node &A = CG.get(lookupFunction(*M, "a"));
   LazyCallGraph::Node &B = CG.get(lookupFunction(*M, "b"));
-  EXPECT_EQ(2, std::distance(A.begin(), A.end()));
-  EXPECT_EQ(0, std::distance(B.begin(), B.end()));
+  A.populate();
+  EXPECT_EQ(2, std::distance(A->begin(), A->end()));
+  B.populate();
+  EXPECT_EQ(0, std::distance(B->begin(), B->end()));
+
+  LazyCallGraph::Node &C = CG.get(lookupFunction(*M, "c"));
+  C.populate();
+  CG.insertEdge(B, C, LazyCallGraph::Edge::Call);
+  EXPECT_EQ(1, std::distance(B->begin(), B->end()));
+  EXPECT_EQ(0, std::distance(C->begin(), C->end()));
+
+  CG.insertEdge(C, B, LazyCallGraph::Edge::Call);
+  EXPECT_EQ(1, std::distance(C->begin(), C->end()));
+  EXPECT_EQ(&B, &C->begin()->getNode());
+
+  CG.insertEdge(C, C, LazyCallGraph::Edge::Call);
+  EXPECT_EQ(2, std::distance(C->begin(), C->end()));
+  EXPECT_EQ(&B, &C->begin()->getNode());
+  EXPECT_EQ(&C, &std::next(C->begin())->getNode());
+
+  CG.removeEdge(C, B);
+  EXPECT_EQ(1, std::distance(C->begin(), C->end()));
+  EXPECT_EQ(&C, &C->begin()->getNode());
 
-  CG.insertEdge(B, lookupFunction(*M, "c"), LazyCallGraph::Edge::Call);
-  EXPECT_EQ(1, std::distance(B.begin(), B.end()));
-  LazyCallGraph::Node &C = B.begin()->getNode(CG);
-  EXPECT_EQ(0, std::distance(C.begin(), C.end()));
-
-  CG.insertEdge(C, B.getFunction(), LazyCallGraph::Edge::Call);
-  EXPECT_EQ(1, std::distance(C.begin(), C.end()));
-  EXPECT_EQ(&B, C.begin()->getNode());
-
-  CG.insertEdge(C, C.getFunction(), LazyCallGraph::Edge::Call);
-  EXPECT_EQ(2, std::distance(C.begin(), C.end()));
-  EXPECT_EQ(&B, C.begin()->getNode());
-  EXPECT_EQ(&C, std::next(C.begin())->getNode());
-
-  CG.removeEdge(C, B.getFunction());
-  EXPECT_EQ(1, std::distance(C.begin(), C.end()));
-  EXPECT_EQ(&C, C.begin()->getNode());
+  CG.removeEdge(C, C);
+  EXPECT_EQ(0, std::distance(C->begin(), C->end()));
 
-  CG.removeEdge(C, C.getFunction());
-  EXPECT_EQ(0, std::distance(C.begin(), C.end()));
-
-  CG.removeEdge(B, C.getFunction());
-  EXPECT_EQ(0, std::distance(B.begin(), B.end()));
+  CG.removeEdge(B, C);
+  EXPECT_EQ(0, std::distance(B->begin(), B->end()));
 }
 
 TEST(LazyCallGraphTest, InnerSCCFormation) {
   LLVMContext Context;
   std::unique_ptr<Module> M = parseAssembly(Context, DiamondOfTriangles);
   LazyCallGraph CG(*M);
 
   // Now mutate the graph to connect every node into a single RefSCC to ensure
   // that our inner SCC formation handles the rest.
-  CG.insertEdge(lookupFunction(*M, "d1"), lookupFunction(*M, "a1"),
-                LazyCallGraph::Edge::Ref);
+  LazyCallGraph::Node &D1 = CG.get(lookupFunction(*M, "d1"));
+  LazyCallGraph::Node &A1 = CG.get(lookupFunction(*M, "a1"));
+  A1.populate();
+  D1.populate();
+  CG.insertEdge(D1, A1, LazyCallGraph::Edge::Ref);
 
   // Build vectors and sort them for the rest of the assertions to make them
   // independent of order.
   std::vector<std::string> Nodes;
 
   // We should build a single RefSCC for the entire graph.
   CG.buildRefSCCs();
   auto I = CG.postorder_ref_scc_begin();
TEST(LazyCallGraphTest, OutgoingEdgeMutation) {
   EXPECT_FALSE(DC.isChildOf(AC));
   EXPECT_TRUE(DRC.isDescendantOf(ARC));
   EXPECT_TRUE(DC.isDescendantOf(AC));
   EXPECT_TRUE(DRC.isChildOf(BRC));
   EXPECT_TRUE(DC.isChildOf(BC));
   EXPECT_TRUE(DRC.isChildOf(CRC));
   EXPECT_TRUE(DC.isChildOf(CC));
 
-  EXPECT_EQ(2, std::distance(A.begin(), A.end()));
+  EXPECT_EQ(2, std::distance(A->begin(), A->end()));
   ARC.insertOutgoingEdge(A, D, LazyCallGraph::Edge::Call);
-  EXPECT_EQ(3, std::distance(A.begin(), A.end()));
-  const LazyCallGraph::Edge &NewE = A[D];
+  EXPECT_EQ(3, std::distance(A->begin(), A->end()));
+  const LazyCallGraph::Edge &NewE = (*A)[D];
   EXPECT_TRUE(NewE);
   EXPECT_TRUE(NewE.isCall());
-  EXPECT_EQ(&D, NewE.getNode());
+  EXPECT_EQ(&D, &NewE.getNode());
 
   // Only the parent and child tests sholud have changed. The rest of the graph
   // remains the same.
   EXPECT_TRUE(ARC.isParentOf(DRC));
   EXPECT_TRUE(AC.isParentOf(DC));
   EXPECT_TRUE(ARC.isAncestorOf(DRC));
   EXPECT_TRUE(AC.isAncestorOf(DC));
   EXPECT_TRUE(DRC.isChildOf(ARC));
TEST(LazyCallGraphTest, OutgoingEdgeMutation) {
   EXPECT_EQ(&CC, CG.lookupSCC(C));
   EXPECT_EQ(&DC, CG.lookupSCC(D));
   EXPECT_EQ(&ARC, CG.lookupRefSCC(A));
   EXPECT_EQ(&BRC, CG.lookupRefSCC(B));
   EXPECT_EQ(&CRC, CG.lookupRefSCC(C));
   EXPECT_EQ(&DRC, CG.lookupRefSCC(D));
 
   ARC.removeOutgoingEdge(A, D);
-  EXPECT_EQ(2, std::distance(A.begin(), A.end()));
+  EXPECT_EQ(2, std::distance(A->begin(), A->end()));
 
   // Now the parent and child tests fail again but the rest remains the same.
   EXPECT_FALSE(ARC.isParentOf(DRC));
   EXPECT_FALSE(AC.isParentOf(DC));
   EXPECT_TRUE(ARC.isAncestorOf(DRC));
   EXPECT_TRUE(AC.isAncestorOf(DC));
   EXPECT_FALSE(DRC.isChildOf(ARC));
   EXPECT_FALSE(DC.isChildOf(AC));
TEST(LazyCallGraphTest, IncomingEdgeInsertion) {
   ASSERT_EQ(&ARC, CG.lookupRefSCC(A2));
   ASSERT_EQ(&ARC, CG.lookupRefSCC(A3));
   ASSERT_EQ(&BRC, CG.lookupRefSCC(B2));
   ASSERT_EQ(&BRC, CG.lookupRefSCC(B3));
   ASSERT_EQ(&CRC, CG.lookupRefSCC(C2));
   ASSERT_EQ(&CRC, CG.lookupRefSCC(C3));
   ASSERT_EQ(&DRC, CG.lookupRefSCC(D2));
   ASSERT_EQ(&DRC, CG.lookupRefSCC(D3));
-  ASSERT_EQ(1, std::distance(D2.begin(), D2.end()));
+  ASSERT_EQ(1, std::distance(D2->begin(), D2->end()));
 
   // Add an edge to make the graph:
   //
   //         d1         |
   //        /  \        |
   //       d3--d2---.   |
   //      /     \    |  |
   //     b1     c1   |  |
   //   /  \    /  \ /   |
   //  b3--b2  c3--c2    |
   //       \  /         |
   //        a1          |
   //       /  \         |
   //      a3--a2        |
   auto MergedRCs = CRC.insertIncomingRefEdge(D2, C2);
   // Make sure we connected the nodes.
-  for (LazyCallGraph::Edge E : D2) {
-    if (E.getNode() == &D3)
+  for (LazyCallGraph::Edge E : *D2) {
+    if (&E.getNode() == &D3)
       continue;
-    EXPECT_EQ(&C2, E.getNode());
+    EXPECT_EQ(&C2, &E.getNode());
   }
   // And marked the D ref-SCC as no longer valid.
   EXPECT_EQ(1u, MergedRCs.size());
   EXPECT_EQ(&DRC, MergedRCs[0]);
 
   // Make sure we have the correct nodes in the SCC sets.
   EXPECT_EQ(&ARC, CG.lookupRefSCC(A1));
   EXPECT_EQ(&ARC, CG.lookupRefSCC(A2));
TEST(LazyCallGraphTest, IncomingEdgeInsertionRefGraph) {
   ASSERT_EQ(&ARC, CG.lookupRefSCC(A2));
   ASSERT_EQ(&ARC, CG.lookupRefSCC(A3));
   ASSERT_EQ(&BRC, CG.lookupRefSCC(B2));
   ASSERT_EQ(&BRC, CG.lookupRefSCC(B3));
   ASSERT_EQ(&CRC, CG.lookupRefSCC(C2));
   ASSERT_EQ(&CRC, CG.lookupRefSCC(C3));
   ASSERT_EQ(&DRC, CG.lookupRefSCC(D2));
   ASSERT_EQ(&DRC, CG.lookupRefSCC(D3));
-  ASSERT_EQ(1, std::distance(D2.begin(), D2.end()));
+  ASSERT_EQ(1, std::distance(D2->begin(), D2->end()));
 
   // Add an edge to make the graph:
   //
   //         d1         |
   //        /  \        |
   //       d3--d2---.   |
   //      /     \    |  |
   //     b1     c1   |  |
   //   /  \    /  \ /   |
   //  b3--b2  c3--c2    |
   //       \  /         |
   //        a1          |
   //       /  \         |
   //      a3--a2        |
   auto MergedRCs = CRC.insertIncomingRefEdge(D2, C2);
   // Make sure we connected the nodes.
-  for (LazyCallGraph::Edge E : D2) {
-    if (E.getNode() == &D3)
+  for (LazyCallGraph::Edge E : *D2) {
+    if (&E.getNode() == &D3)
       continue;
-    EXPECT_EQ(&C2, E.getNode());
+    EXPECT_EQ(&C2, &E.getNode());
   }
   // And marked the D ref-SCC as no longer valid.
   EXPECT_EQ(1u, MergedRCs.size());
   EXPECT_EQ(&DRC, MergedRCs[0]);
 
   // Make sure we have the correct nodes in the SCC sets.
   EXPECT_EQ(&ARC, CG.lookupRefSCC(A1));
   EXPECT_EQ(&ARC, CG.lookupRefSCC(A2));
TEST(LazyCallGraphTest, IncomingEdgeInsertionLargeCallCycle) {
   LazyCallGraph::RefSCC &ARC = *CG.lookupRefSCC(A);
   LazyCallGraph::RefSCC &BRC = *CG.lookupRefSCC(B);
   LazyCallGraph::RefSCC &CRC = *CG.lookupRefSCC(C);
   LazyCallGraph::RefSCC &DRC = *CG.lookupRefSCC(D);
 
   // Connect the top to the bottom forming a large RefSCC made up mostly of calls.
   auto MergedRCs = ARC.insertIncomingRefEdge(D, A);
   // Make sure we connected the nodes.
-  EXPECT_NE(D.begin(), D.end());
-  EXPECT_EQ(&A, D.begin()->getNode());
+  EXPECT_NE(D->begin(), D->end());
+  EXPECT_EQ(&A, &D->begin()->getNode());
 
   // Check that we have the dead RCs, but ignore the order.
   EXPECT_EQ(3u, MergedRCs.size());
   EXPECT_NE(find(MergedRCs, &BRC), MergedRCs.end());
   EXPECT_NE(find(MergedRCs, &CRC), MergedRCs.end());
   EXPECT_NE(find(MergedRCs, &DRC), MergedRCs.end());
 
   // Make sure the nodes point to the right place now.
TEST(LazyCallGraphTest, IncomingEdgeInsertionLargeRefCycle) {
   LazyCallGraph::RefSCC &BRC = *CG.lookupRefSCC(B);
   LazyCallGraph::RefSCC &CRC = *CG.lookupRefSCC(C);
   LazyCallGraph::RefSCC &DRC = *CG.lookupRefSCC(D);
 
   // Connect the top to the bottom forming a large RefSCC made up just of
   // references.
   auto MergedRCs = ARC.insertIncomingRefEdge(D, A);
   // Make sure we connected the nodes.
-  EXPECT_NE(D.begin(), D.end());
-  EXPECT_EQ(&A, D.begin()->getNode());
+  EXPECT_NE(D->begin(), D->end());
+  EXPECT_EQ(&A, &D->begin()->getNode());
 
   // Check that we have the dead RCs, but ignore the order.
   EXPECT_EQ(3u, MergedRCs.size());
   EXPECT_NE(find(MergedRCs, &BRC), MergedRCs.end());
   EXPECT_NE(find(MergedRCs, &CRC), MergedRCs.end());
   EXPECT_NE(find(MergedRCs, &DRC), MergedRCs.end());
 
   // Make sure the nodes point to the right place now.
TEST(LazyCallGraphTest, InlineAndDeleteFunction) {
   ASSERT_EQ(&ARC, CG.lookupRefSCC(A2));
   ASSERT_EQ(&ARC, CG.lookupRefSCC(A3));
   ASSERT_EQ(&BRC, CG.lookupRefSCC(B2));
   ASSERT_EQ(&BRC, CG.lookupRefSCC(B3));
   ASSERT_EQ(&CRC, CG.lookupRefSCC(C2));
   ASSERT_EQ(&CRC, CG.lookupRefSCC(C3));
   ASSERT_EQ(&DRC, CG.lookupRefSCC(D2));
   ASSERT_EQ(&DRC, CG.lookupRefSCC(D3));
-  ASSERT_EQ(1, std::distance(D2.begin(), D2.end()));
+  ASSERT_EQ(1, std::distance(D2->begin(), D2->end()));
 
   // Delete d2 from the graph, as if it had been inlined.
   //
   //         d1         |
   //        / /         |
   //       d3--.        |
   //      /     \       |
   //     b1     c1      |
TEST(LazyCallGraphTest, InternalEdgeMutation) {
   EXPECT_EQ(&RC, CG.lookupRefSCC(C));
   EXPECT_EQ(1, RC.size());
   EXPECT_EQ(&*RC.begin(), CG.lookupSCC(A));
   EXPECT_EQ(&*RC.begin(), CG.lookupSCC(B));
   EXPECT_EQ(&*RC.begin(), CG.lookupSCC(C));
 
   // Insert an edge from 'a' to 'c'. Nothing changes about the graph.
   RC.insertInternalRefEdge(A, C);
-  EXPECT_EQ(2, std::distance(A.begin(), A.end()));
+  EXPECT_EQ(2, std::distance(A->begin(), A->end()));
   EXPECT_EQ(&RC, CG.lookupRefSCC(A));
   EXPECT_EQ(&RC, CG.lookupRefSCC(B));
   EXPECT_EQ(&RC, CG.lookupRefSCC(C));
   EXPECT_EQ(1, RC.size());
   EXPECT_EQ(&*RC.begin(), CG.lookupSCC(A));
   EXPECT_EQ(&*RC.begin(), CG.lookupSCC(B));
   EXPECT_EQ(&*RC.begin(), CG.lookupSCC(C));
 
TEST(LazyCallGraphTest, HandleBlockAddress) {
 
   LazyCallGraph::Node &F = *CG.lookup(lookupFunction(*M, "f"));
   LazyCallGraph::Node &G = *CG.lookup(lookupFunction(*M, "g"));
   EXPECT_EQ(&FRC, CG.lookupRefSCC(F));
   EXPECT_EQ(&GRC, CG.lookupRefSCC(G));
   EXPECT_TRUE(GRC.isParentOf(FRC));
 }
 
+TEST(LazyCallGraphTest, ReplaceNodeFunction) {
+  LLVMContext Context;
+  // A graph with several different kinds of edges pointing at a particular
+  // function.
+  std::unique_ptr<Module> M =
+      parseAssembly(Context,
+                    "define void @a(i8** %ptr) {\n"
+                    "entry:\n"
+                    "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+                    "  ret void\n"
+                    "}\n"
+                    "define void @b(i8** %ptr) {\n"
+                    "entry:\n"
+                    "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+                    "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+                    "  call void @d(i8** %ptr)"
+                    "  ret void\n"
+                    "}\n"
+                    "define void @c(i8** %ptr) {\n"
+                    "entry:\n"
+                    "  call void @d(i8** %ptr)"
+                    "  call void @d(i8** %ptr)"
+                    "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+                    "  ret void\n"
+                    "}\n"
+                    "define void @d(i8** %ptr) {\n"
+                    "entry:\n"
+                    "  store i8* bitcast (void(i8**)* @b to i8*), i8** %ptr\n"
+                    "  call void @c(i8** %ptr)"
+                    "  call void @d(i8** %ptr)"
+                    "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+                    "  ret void\n"
+                    "}\n");
+  LazyCallGraph CG(*M);
+
+  // Force the graph to be fully expanded.
+  CG.buildRefSCCs();
+  auto I = CG.postorder_ref_scc_begin();
+  LazyCallGraph::RefSCC &RC1 = *I++;
+  LazyCallGraph::RefSCC &RC2 = *I++;
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+
+  ASSERT_EQ(2u, RC1.size());
+  LazyCallGraph::SCC &C1 = RC1[0];
+  LazyCallGraph::SCC &C2 = RC1[1];
+
+  LazyCallGraph::Node &AN = *CG.lookup(lookupFunction(*M, "a"));
+  LazyCallGraph::Node &BN = *CG.lookup(lookupFunction(*M, "b"));
+  LazyCallGraph::Node &CN = *CG.lookup(lookupFunction(*M, "c"));
+  LazyCallGraph::Node &DN = *CG.lookup(lookupFunction(*M, "d"));
+  EXPECT_EQ(&C1, CG.lookupSCC(DN));
+  EXPECT_EQ(&C1, CG.lookupSCC(CN));
+  EXPECT_EQ(&C2, CG.lookupSCC(BN));
+  EXPECT_EQ(&RC1, CG.lookupRefSCC(DN));
+  EXPECT_EQ(&RC1, CG.lookupRefSCC(CN));
+  EXPECT_EQ(&RC1, CG.lookupRefSCC(BN));
+  EXPECT_EQ(&RC2, CG.lookupRefSCC(AN));
+
+  // Now we need to build a new function 'e' with the same signature as 'd'.
+  Function &D = DN.getFunction();
+  Function &E = *Function::Create(D.getFunctionType(), D.getLinkage(), "e");
+  D.getParent()->getFunctionList().insert(D.getIterator(), &E);
+
+  // Change each use of 'd' to use 'e'. This is particularly easy as they have
+  // the same type.
+  D.replaceAllUsesWith(&E);
+
+  // Splice the body of the old function into the new one.
+  E.getBasicBlockList().splice(E.begin(), D.getBasicBlockList());
+  // And fix up the one argument.
+  D.arg_begin()->replaceAllUsesWith(&*E.arg_begin());
+  E.arg_begin()->takeName(&*D.arg_begin());
+
+  // Now replace the function in the graph.
+  RC1.replaceNodeFunction(DN, E);
+
+  EXPECT_EQ(&E, &DN.getFunction());
+  EXPECT_EQ(&DN, &(*CN)[DN].getNode());
+  EXPECT_EQ(&DN, &(*BN)[DN].getNode());
+}
 }