From 0b4b0869a7968a2612c9bd93019a021ce168c557 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9my=20Rakic?= <remy.rakic+github@gmail.com>
Date: Wed, 14 Jun 2023 20:00:23 +0000
Subject: [PATCH] make `traversal::postorder` traverse RPO cache backwards

---
 compiler/rustc_middle/src/mir/traversal.rs | 40 +++++++++++++++++++---
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_middle/src/mir/traversal.rs b/compiler/rustc_middle/src/mir/traversal.rs
index 780003f5fa05..6d44234d11f2 100644
--- a/compiler/rustc_middle/src/mir/traversal.rs
+++ b/compiler/rustc_middle/src/mir/traversal.rs
@@ -188,10 +188,6 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
     }
 }
 
-pub fn postorder<'a, 'tcx>(body: &'a Body<'tcx>) -> Postorder<'a, 'tcx> {
-    Postorder::new(&body.basic_blocks, START_BLOCK)
-}
-
 impl<'a, 'tcx> Iterator for Postorder<'a, 'tcx> {
     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
 
@@ -219,6 +215,42 @@ impl<'a, 'tcx> Iterator for Postorder<'a, 'tcx> {
     }
 }
 
+/// Creates an iterator over the `Body`'s basic blocks, that:
+/// - returns basic blocks in a postorder,
+/// - traverses the `BasicBlocks` CFG cache's reverse postorder backwards, and does not cache the
+///   postorder itself.
+pub fn postorder<'a, 'tcx>(body: &'a Body<'tcx>) -> PostorderIter<'a, 'tcx> {
+    let blocks = body.basic_blocks.reverse_postorder();
+    let len = blocks.len();
+    PostorderIter { body, blocks, idx: len }
+}
+
+#[derive(Clone)]
+pub struct PostorderIter<'a, 'tcx> {
+    body: &'a Body<'tcx>,
+    blocks: &'a [BasicBlock],
+    idx: usize,
+}
+
+impl<'a, 'tcx> Iterator for PostorderIter<'a, 'tcx> {
+    type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
+
+    fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
+        if self.idx == 0 {
+            return None;
+        }
+        self.idx -= 1;
+
+        self.blocks.get(self.idx).map(|&bb| (bb, &self.body[bb]))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.idx, Some(self.idx))
+    }
+}
+
+impl<'a, 'tcx> ExactSizeIterator for PostorderIter<'a, 'tcx> {}
+
 /// Reverse postorder traversal of a graph
 ///
 /// Reverse postorder is the reverse order of a postorder traversal.