[llvm] [MachineOutliner] Leaf Descendants (PR #90275)

[Xuan Zhang via llvm-commits]

================
@@ -105,6 +112,70 @@ void SuffixTree::setSuffixIndices() {
   }
 }
 
+void SuffixTree::setLeafNodes() {
+  // A stack that keeps track of nodes to visit for post-order DFS traversal.
+  SmallVector<SuffixTreeNode *> ToVisit;
+  ToVisit.push_back(Root);
+
+  // This keeps track of the index of the next leaf node to be added to
+  // the LeafNodes vector of the suffix tree.
+  unsigned LeafCounter = 0;
+
+  // This keeps track of nodes whose children have been added to the stack.
+  // The value is a pair, representing a node's first and last children.
+  DenseMap<SuffixTreeInternalNode *,
+           std::pair<SuffixTreeNode *, SuffixTreeNode *>>
+      ChildrenMap;
+
+  // Traverse the tree in post-order.
+  while (!ToVisit.empty()) {
+    SuffixTreeNode *CurrNode = ToVisit.pop_back_val();
+    if (auto *CurrInternalNode = dyn_cast<SuffixTreeInternalNode>(CurrNode)) {
+      // The current node is an internal node.
+      auto I = ChildrenMap.find(CurrInternalNode);
+      if (I == ChildrenMap.end()) {
+        // This is the first time we visit this node.
+        // Its children have not been added to the stack yet.
+        // We add current node back, and add its children to the stack.
+        // We keep track of the first and last children of the current node.
+        auto it = CurrInternalNode->Children.begin();
+        if (it != CurrInternalNode->Children.end()) {
+          ToVisit.push_back(CurrNode);
+          SuffixTreeNode *FirstChild = it->second;
+          SuffixTreeNode *LastChild = nullptr;
+          for (; it != CurrInternalNode->Children.end(); ++it) {
+            LastChild = it->second;
+            ToVisit.push_back(LastChild);
+          }
+          ChildrenMap[CurrInternalNode] = {FirstChild, LastChild};
+        }
+      } else {
+        // This is the second time we visit this node.
+        // All of its children have already been processed.
+        // Now, we can set its LeftLeafIdx and RightLeafIdx;
+        auto [FirstChild, LastChild] = I->second;
+        // Get the first child to use its RightLeafIdx.
+        // The first child is the first one added to the stack, so it is
+        // the last one to be processed. Hence, the leaf descendants
+        // of the first child are assigned the largest index numbers.
+        CurrNode->setRightLeafIdx(FirstChild->getRightLeafIdx());
+        // Get the last child to use its LeftLeafIdx.
+        CurrNode->setLeftLeafIdx(LastChild->getLeftLeafIdx());
+        assert(CurrNode->getLeftLeafIdx() <= CurrNode->getRightLeafIdx() &&
+               "LeftLeafIdx should not be larger than RightLeafIdx");
+      }
+    } else {
+      // The current node is a leaf node.
+      // We can simply set its LeftLeafIdx and RightLeafIdx.
+      CurrNode->setLeftLeafIdx(LeafCounter);
+      CurrNode->setRightLeafIdx(LeafCounter);
+      LeafCounter++;
----------------
xuanzh-meta wrote:

Interesting, will fix!

https://github.com/llvm/llvm-project/pull/90275