Automerge: [ORC] Clear stale ElemToPendingSN entries in WaitingOnGraph. (#169747)

WaitingOnGraph::processReadyOrFailed was not clearing stale entries from
the ElemToPendingSN map. If symbols were removed from the
ExecutionSession and then re-added this could lead to dependencies on
the stale entries, triggering a use-after-free bug.

llvm/llvm-project#169135

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h

@@ -338,9 +338,9 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
     // incorporate NewSNs.
     std::vector<std::unique_ptr<SuperNode>> ReadyNodes, FailedNodes;
     processReadyOrFailed(ModifiedPendingSNs, ReadyNodes, FailedNodes,
-                         SuperNodeDeps, ElemToPendingSN, FailedSNs);
+                         SuperNodeDeps, FailedSNs, &ElemToPendingSN);
     processReadyOrFailed(NewSNs, ReadyNodes, FailedNodes, SuperNodeDeps,
-                         ElemToNewSN, FailedSNs);
+                         FailedSNs, nullptr);
 
     CoalesceToPendingSNs.coalesce(ModifiedPendingSNs, ElemToPendingSN);
     CoalesceToPendingSNs.coalesce(NewSNs, ElemToPendingSN);
@@ -591,8 +591,11 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
                             std::vector<std::unique_ptr<SuperNode>> &Ready,
                             std::vector<std::unique_ptr<SuperNode>> &Failed,
                             SuperNodeDepsMap &SuperNodeDeps,
-                            ElemToSuperNodeMap &ElemToSNs,
-                            const std::vector<SuperNode *> &FailedSNs) {
+                            const std::vector<SuperNode *> &FailedSNs,
+                            ElemToSuperNodeMap *ElemToSNs) {
+
+    SmallVector<SuperNode *> ToRemoveFromElemToSNs;
+
     for (size_t I = 0; I != SNs.size();) {
       auto &SN = SNs[I];
 
@@ -609,13 +612,24 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
       bool SNReady = SN->Deps.empty();
 
       if (SNReady || SNFailed) {
+        if (ElemToSNs)
+          ToRemoveFromElemToSNs.push_back(SN.get());
         auto &NodeList = SNFailed ? Failed : Ready;
         NodeList.push_back(std::move(SN));
         std::swap(SN, SNs.back());
         SNs.pop_back();
       } else
         ++I;
     }
+
+    // Update ElemToSNs (if passed) to remove elements pointing at SN.
+    for (auto *SN : ToRemoveFromElemToSNs) {
+      for (auto &[Container, Elems] : SN->defs()) {
+        auto &Row = (*ElemToSNs)[Container];
+        for (auto &Elem : Elems)
+          Row.erase(Elem);
+      }
+    }
   }
 
   std::vector<std::unique_ptr<SuperNode>> PendingSNs;

llvm/unittests/ExecutionEngine/Orc/WaitingOnGraphTest.cpp

@@ -532,6 +532,52 @@ TEST_F(WaitingOnGraphTest, Emit_ZigZag) {
   EXPECT_TRUE(PendingSNs.empty());
 }
 
+TEST_F(WaitingOnGraphTest, Emit_ReEmit) {
+  // Test for the bug in https://github.com/llvm/llvm-project/issues/169135,
+  // which was caused by stale entries in the ElemsToPendingSNs map.
+  //
+  // To trigger the bug we need to:
+  // 1. Create a SuperNode with an unmet dependence, causing it to be added to
+  //    ElemsToPendingSNs.
+  // 2. Cause that SuperNode to become ready (bug left stale entries in map)
+  // 3. Remove the node from the Ready map (this is equivalent to removal of a
+  //    symbol in an ORC session, and allows new SuperNodes to depend on the
+  //    stale entry).
+  // 4. Add a new node that references the previously emitted/removed SuperNode
+  //    This triggers access of the stale entry, and should error out in
+  //    sanitizer builds.
+
+  SuperNodeBuilder B;
+
+  // 1. Create SuperNode with unmet dependence.
+  ContainerElementsMap Defs0({{0, {0}}});
+  ContainerElementsMap Deps0({{0, {1}}});
+  B.add(Defs0, Deps0);
+  emit(TestGraph::simplify(B.takeSuperNodes()));
+
+  EXPECT_TRUE(Ready.empty());
+
+  // 2. Cause previous SuperNode to become ready.
+  ContainerElementsMap Defs1({{0, {1}}});
+  B.add(Defs1, ContainerElementsMap());
+  emit(TestGraph::simplify(B.takeSuperNodes()));
+
+  // Check that both nodes have become ready.
+  EXPECT_EQ(Ready, merge(Defs0, Defs1));
+
+  // 3. Erase Ready nodes to simulate removal from the graph.
+  Ready.clear();
+
+  // 4. Emit a new dependence on the original def.
+  ContainerElementsMap Defs2({{0, {2}}});
+  ContainerElementsMap Deps2({{0, {0}}});
+  B.add(Defs2, Deps2);
+  auto ER = emit(TestGraph::simplify(B.takeSuperNodes()));
+
+  // We expect the new dependence to remain pending.
+  EXPECT_TRUE(ER.Ready.empty());
+}
+
 TEST_F(WaitingOnGraphTest, Fail_Empty) {
   // Check that failing an empty set is a no-op.
   auto FR = G.fail(ContainerElementsMap());