[SCCIterator] Fix another potential use-after-free

[loladiro]

Summary

This is a follow up to https://reviews.llvm.org/D72529. Since this
code is a bit old, I decided to see if there were any similar issues
hiding and believe I have found one. The Inliner transform deletes
CallGraphNodes, but doesn't remove the deleted node from the scc
iterator. At first glance, one might think this is fine, because
we don't access the CallGraphNodes through that pointer and it
simply serves as an indication whether or not we have visited
a given CallGraphNode. However, there is a subtle problem:
The pointer of the CallGraphNode we have freed may be re-used
by a later allocation of a new CallGraphNode. This could now
present a problem, because the scc iterator will associate the
old state with the new CallGraphNode. Fix that by adding the
ability to inform the scc iterator of node deletion.

I suspect this situation never actually causes a problem in
practice because the sequence of events (node deleted -> node
reallocated with the same pointer -> node accessed in CFG
walk) is quite long and may not actually lead to a crash or
even incorrect code (I believe the symptom would be a missing
node in the SCC). I'd like to get it fixed anyway. As a bonus,
I believe we can remove an outstanding todo from the Inliner.

[loladiro]

Turns out this doesn't quite yet address the TODO in Inliner.
Reverted that part of this change for a complete solution in a future
revision.

[wristow]

This is a follow up to https://reviews.llvm.org/D72529.

For reference, we independently encountered a similar issue to the one described in the above review, and we proposed a fix (at D72469) that has now been committed.

[lebedev.ri]

What's the status here?

[lebedev.ri]

Reverse-ping.

I'm suspecting this might or might not be relevant for the
http://lab.llvm.org:8011/builders/llvm-clang-x86_64-expensive-checks-win/builds/23032/steps/test-check-all/logs/FAIL%3A%20LLVM%3A%3Afp80.ll
failures from before rG0c22cb0.
Is anything holding this patch up? Can we land this?

[jdoerfert]

llvm/lib/Analysis/CallGraphSCCPass.cpp

557

This part actually landed by now.

[lebedev.ri]

@loladiro rebase or abandon this?

llvm/include/llvm/ADT/SCCIterator.h

public:
   /// Test if the current SCC has a loop.
   ///
   /// If the SCC has more than one node, this is trivially true.  If not, it may
   /// still contain a loop if the node has an edge back to itself.
   bool hasLoop() const;
 
   /// This informs the \c scc_iterator that the specified \c Old node
   /// has been deleted, and \c New is to be used in its place.
+  /// If \c New is null, the node has been deleted without replacement.
   void ReplaceNode(NodeRef Old, NodeRef New) {
     auto it = nodeVisitNumbers.find(Old);
     assert(it != nodeVisitNumbers.end() && "Old not in scc_iterator?");
     unsigned OldVisitNumber = it->second;
+    // Since we're not updating any of our internal book keeping data structures,
+    // we need to make sure that the node we're replacing belongs to an SCC
+    // that's already been visited.
+    assert(OldVisitNumber == ~0U && "Old not in an scc that was already visited");
     nodeVisitNumbers.erase(it);
-    nodeVisitNumbers[New] = OldVisitNumber;
+    if (New) {
+        bool Added = nodeVisitNumbers.try_emplace(New, OldVisitNumber).second;
+        (void)Added;
+        assert(Added && "New already in scc_iterator?");
+    }
   }
 };
 
 template <class GraphT, class GT>
 void scc_iterator<GraphT, GT>::DFSVisitOne(NodeRef N) {
   ++visitNum;
   nodeVisitNumbers[N] = visitNum;
   SCCNodeStack.push_back(N);

llvm/include/llvm/Analysis/CallGraphSCCPass.h

 // call graph up-to-date if they do anything which could modify it.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_CALLGRAPHSCCPASS_H
 #define LLVM_ANALYSIS_CALLGRAPHSCCPASS_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SCCIterator.h"
 #include "llvm/Pass.h"
 #include <vector>
 
 namespace llvm {
 
 class CallGraph;
 class CallGraphNode;
 class CallGraphSCC;
protected:
   /// Optional passes call this function to check whether the pass should be
   /// skipped. This is the case when optimization bisect is over the limit.
   bool skipSCC(CallGraphSCC &SCC) const;
 };
 
 /// CallGraphSCC - This is a single SCC that a CallGraphSCCPass is run on.
 class CallGraphSCC {
   const CallGraph &CG; // The call graph for this SCC.
-  void *Context; // The CGPassManager object that is vending this.
+  scc_iterator<CallGraph *> &Context; // The CGPassManager object that is vending this.
   std::vector<CallGraphNode *> Nodes;
 
 public:
-  CallGraphSCC(CallGraph &cg, void *context) : CG(cg), Context(context) {}
+  CallGraphSCC(CallGraph &cg, scc_iterator<CallGraph *> &context) : CG(cg), Context(context) {}
 
   void initialize(ArrayRef<CallGraphNode *> NewNodes) {
     Nodes.assign(NewNodes.begin(), NewNodes.end());
   }
 
   bool isSingular() const { return Nodes.size() == 1; }
   unsigned size() const { return Nodes.size(); }
 
   /// ReplaceNode - This informs the SCC and the pass manager that the specified
   /// Old node has been deleted, and New is to be used in its place.
+  /// If New is null, the node has been deleted without replacement.
   void ReplaceNode(CallGraphNode *Old, CallGraphNode *New);
 
   using iterator = std::vector<CallGraphNode *>::const_iterator;
 
   iterator begin() const { return Nodes.begin(); }
   iterator end() const { return Nodes.end(); }
 
   const CallGraph &getCallGraph() { return CG; }
+  scc_iterator<CallGraph *> &getContext() { return Context; }
 };
 
 void initializeDummyCGSCCPassPass(PassRegistry &);
 
 /// This pass is required by interprocedural register allocation. It forces
 /// codegen to follow bottom up order on call graph.
 class DummyCGSCCPass : public CallGraphSCCPass {
 public:

llvm/lib/Analysis/CallGraphSCCPass.cpp

 /// whether any of the passes modifies the module, and if so, return true.
 bool CGPassManager::runOnModule(Module &M) {
   CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
   bool Changed = doInitialization(CG);
 
   // Walk the callgraph in bottom-up SCC order.
   scc_iterator<CallGraph*> CGI = scc_begin(&CG);
 
-  CallGraphSCC CurSCC(CG, &CGI);
+  CallGraphSCC CurSCC(CG, CGI);
   while (!CGI.isAtEnd()) {
     // Copy the current SCC and increment past it so that the pass can hack
     // on the SCC if it wants to without invalidating our iterator.
     const std::vector<CallGraphNode *> &NodeVec = *CGI;
     CurSCC.initialize(NodeVec);
     ++CGI;
 
     // At the top level, we run all the passes in this pass manager on the

 //===----------------------------------------------------------------------===//
 // CallGraphSCC Implementation
 //===----------------------------------------------------------------------===//
 
 /// This informs the SCC and the pass manager that the specified
 /// Old node has been deleted, and New is to be used in its place.
 void CallGraphSCC::ReplaceNode(CallGraphNode *Old, CallGraphNode *New) {
   assert(Old != New && "Should not replace node with self");
-  for (unsigned i = 0; ; ++i) {
-    assert(i != Nodes.size() && "Node not in SCC");
-    if (Nodes[i] != Old) continue;
-    Nodes[i] = New;
+  for (auto it = Nodes.begin(); ; ++it) {
+    assert(it != Nodes.end() && "Node not in SCC");
+    if (*it != Old)
+      continue;
+    if (New) {
+      *it = New;
+    } else {
+      Nodes.erase(it);
+    }
     break;
   }
 
   // Update the active scc_iterator so that it doesn't contain dangling
   // pointers to the old CallGraphNode.
-  scc_iterator<CallGraph*> *CGI = (scc_iterator<CallGraph*>*)Context;
-  CGI->ReplaceNode(Old, New);
+  Context.ReplaceNode(Old, New);
 }
 
 //===----------------------------------------------------------------------===//
 // CallGraphSCCPass Implementation
 //===----------------------------------------------------------------------===//
 
 /// Assign pass manager to manage this pass.
 void CallGraphSCCPass::assignPassManager(PMStack &PMS,

llvm/lib/Transforms/IPO/Inliner.cpp

for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
           for (Value *Ptr : InlineInfo.InlinedCalls)
             CallSites.push_back(std::make_pair(CallSite(Ptr), NewHistoryID));
         }
       }
 
       // If we inlined or deleted the last possible call site to the function,
       // delete the function body now.
       if (Callee && Callee->use_empty() && Callee->hasLocalLinkage() &&
-          // TODO: Can remove if in SCC now.
-          !SCCFunctions.count(Callee) &&
           // The function may be apparently dead, but if there are indirect
           // callgraph references to the node, we cannot delete it yet, this
           // could invalidate the CGSCC iterator.
           CG[Callee]->getNumReferences() == 0) {
         LLVM_DEBUG(dbgs() << "    -> Deleting dead function: "
                           << Callee->getName() << "\n");
         CallGraphNode *CalleeNode = CG[Callee];
 
         // Remove any call graph edges from the callee to its callees.
         CalleeNode->removeAllCalledFunctions();
 
         // Removing the node for callee from the call graph and delete it.
         delete CG.removeFunctionFromModule(CalleeNode);
+        if (SCCFunctions.count(Callee))
+          SCC.ReplaceNode(CalleeNode, nullptr);
+        else
+          SCC.getContext().ReplaceNode(CalleeNode, nullptr);
         ++NumDeleted;
       }
 
       // Remove this call site from the list.  If possible, use
       // swap/pop_back for efficiency, but do not use it if doing so would
       // move a call site to a function in this SCC before the
       // 'FirstCallInSCC' barrier.
       if (SCC.isSingular()) {