3 files changed, 32 insertions, 11 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
index 406e253fe4c..fcc1d8f4ed4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
@@ -989,14 +989,29 @@ void SelectionDAGISel::EnforceNodeIdInvariant(SDNode *Node) {
     for (auto *U : N->uses()) {
       auto UId = U->getNodeId();
       if (UId > 0) {
-        int InvalidatedUId = -UId + 1;
-        U->setNodeId(InvalidatedUId);
+        InvalidateNodeId(U);
         Nodes.push_back(U);
       }
     }
   }
 }
 
+// InvalidateNodeId - As discusses in EnforceNodeIdInvariant, mark a
+// NodeId with the equivalent node id which is invalid for topological
+// pruning.
+void SelectionDAGISel::InvalidateNodeId(SDNode *N) {
+  int InvalidId = -N->getNodeId() + 1;
+  N->setNodeId(InvalidId);
+}
+
+// getUninvalidatedNodeId - get original uninvalidated node id.
+int SelectionDAGISel::getUninvalidatedNodeId(SDNode *N) {
+  int Id = N->getNodeId();
+  if (Id < 0)
+    return -Id + 1;
+  return Id;
+}
+
 void SelectionDAGISel::DoInstructionSelection() {
   DEBUG(dbgs() << "===== Instruction selection begins: "
                << printMBBReference(*FuncInfo->MBB) << " '"
diff --git a/llvm/lib/Target/SystemZ/SystemZISelDAGToDAG.cpp b/llvm/lib/Target/SystemZ/SystemZISelDAGToDAG.cpp
index 6e2130828bb..59dd46c0697 100644
--- a/llvm/lib/Target/SystemZ/SystemZISelDAGToDAG.cpp
+++ b/llvm/lib/Target/SystemZ/SystemZISelDAGToDAG.cpp
@@ -593,16 +593,16 @@ bool SystemZDAGToDAGISel::selectAddress(SDValue Addr,
 // The selection DAG must no longer depend on their uniqueness when this
 // function is used.
 static void insertDAGNode(SelectionDAG *DAG, SDNode *Pos, SDValue N) {
-  if (N.getNode()->getNodeId() == -1 ||
-      N.getNode()->getNodeId() > Pos->getNodeId()) {
+  if (N->getNodeId() == -1 ||
+      (SelectionDAGISel::getUninvalidatedNodeId(N.getNode()) >
+       SelectionDAGISel::getUninvalidatedNodeId(Pos))) {
     DAG->RepositionNode(Pos->getIterator(), N.getNode());
     // Mark Node as invalid for pruning as after this it may be a successor to a
     // selected node but otherwise be in the same position of Pos.
     // Conservatively mark it with the same -abs(Id) to assure node id
     // invariant is preserved.
-    int PId = Pos->getNodeId();
-    int InvalidatedPId = -(PId + 1);
-    N->setNodeId((PId > 0) ? InvalidatedPId : PId);
+    N->setNodeId(Pos->getNodeId());
+    SelectionDAGISel::InvalidateNodeId(N.getNode());
   }
 }
 
diff --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
index 375f965b7b2..1edeb85ddb0 100644
--- a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -1070,10 +1070,16 @@ bool X86DAGToDAGISel::matchAdd(SDValue N, X86ISelAddressMode &AM,
 // IDs! The selection DAG must no longer depend on their uniqueness when this
 // is used.
 static void insertDAGNode(SelectionDAG &DAG, SDValue Pos, SDValue N) {
-  if (N.getNode()->getNodeId() == -1 ||
-      N.getNode()->getNodeId() > Pos.getNode()->getNodeId()) {
-    DAG.RepositionNode(Pos.getNode()->getIterator(), N.getNode());
-    N.getNode()->setNodeId(Pos.getNode()->getNodeId());
+  if (N->getNodeId() == -1 ||
+      (SelectionDAGISel::getUninvalidatedNodeId(N.getNode()) >
+       SelectionDAGISel::getUninvalidatedNodeId(Pos.getNode()))) {
+    DAG.RepositionNode(Pos->getIterator(), N.getNode());
+    // Mark Node as invalid for pruning as after this it may be a successor to a
+    // selected node but otherwise be in the same position of Pos.
+    // Conservatively mark it with the same -abs(Id) to assure node id
+    // invariant is preserved.
+    N->setNodeId(Pos->getNodeId());
+    SelectionDAGISel::InvalidateNodeId(N.getNode());
   }
 }