From 9ba0ca9397e65b706a75408e9062e74b03f0bc77 Mon Sep 17 00:00:00 2001 From: Ben Kimock Date: Thu, 19 Sep 2024 20:03:15 -0400 Subject: [PATCH] Deduplicate Postorder logic --- compiler/rustc_middle/src/mir/basic_blocks.rs | 2 +- compiler/rustc_middle/src/mir/traversal.rs | 101 +++++++----------- 2 files changed, 39 insertions(+), 64 deletions(-) diff --git a/compiler/rustc_middle/src/mir/basic_blocks.rs b/compiler/rustc_middle/src/mir/basic_blocks.rs index c08bfc1fa95f4..b8ac21cdfcc2b 100644 --- a/compiler/rustc_middle/src/mir/basic_blocks.rs +++ b/compiler/rustc_middle/src/mir/basic_blocks.rs @@ -71,7 +71,7 @@ impl<'tcx> BasicBlocks<'tcx> { #[inline] pub fn reverse_postorder(&self) -> &[BasicBlock] { self.cache.reverse_postorder.get_or_init(|| { - let mut rpo: Vec<_> = Postorder::new(&self.basic_blocks, START_BLOCK).collect(); + let mut rpo: Vec<_> = Postorder::new(&self.basic_blocks, START_BLOCK, ()).collect(); rpo.reverse(); rpo }) diff --git a/compiler/rustc_middle/src/mir/traversal.rs b/compiler/rustc_middle/src/mir/traversal.rs index eebe8a9ac1df2..b8b74da401c9f 100644 --- a/compiler/rustc_middle/src/mir/traversal.rs +++ b/compiler/rustc_middle/src/mir/traversal.rs @@ -104,36 +104,46 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> { /// ``` /// /// A Postorder traversal of this graph is `D B C A` or `D C B A` -pub struct Postorder<'a, 'tcx> { +pub struct Postorder<'a, 'tcx, C> { basic_blocks: &'a IndexSlice>, visited: BitSet, visit_stack: Vec<(BasicBlock, Successors<'a>)>, root_is_start_block: bool, + extra: C, } -impl<'a, 'tcx> Postorder<'a, 'tcx> { +impl<'a, 'tcx, C> Postorder<'a, 'tcx, C> +where + C: Customization<'tcx>, +{ pub fn new( basic_blocks: &'a IndexSlice>, root: BasicBlock, - ) -> Postorder<'a, 'tcx> { + extra: C, + ) -> Postorder<'a, 'tcx, C> { let mut po = Postorder { basic_blocks, visited: BitSet::new_empty(basic_blocks.len()), visit_stack: Vec::new(), root_is_start_block: root == START_BLOCK, + extra, }; - let data = &po.basic_blocks[root]; - - if let Some(ref term) = data.terminator { - po.visited.insert(root); - po.visit_stack.push((root, term.successors())); - po.traverse_successor(); - } + po.visit(root); + po.traverse_successor(); po } + fn visit(&mut self, bb: BasicBlock) { + if !self.visited.insert(bb) { + return; + } + let data = &self.basic_blocks[bb]; + let successors = C::successors(data, self.extra); + self.visit_stack.push((bb, successors)); + } + fn traverse_successor(&mut self) { // This is quite a complex loop due to 1. the borrow checker not liking it much // and 2. what exactly is going on is not clear @@ -183,16 +193,15 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> { // since we've already visited `E`, that child isn't added to the stack. The last // two iterations yield `B` and finally `A` for a final traversal of [E, D, C, B, A] while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) { - if self.visited.insert(bb) { - if let Some(term) = &self.basic_blocks[bb].terminator { - self.visit_stack.push((bb, term.successors())); - } - } + self.visit(bb); } } } -impl<'tcx> Iterator for Postorder<'_, 'tcx> { +impl<'tcx, C> Iterator for Postorder<'_, 'tcx, C> +where + C: Customization<'tcx>, +{ type Item = BasicBlock; fn next(&mut self) -> Option { @@ -232,57 +241,23 @@ pub fn postorder<'a, 'tcx>( reverse_postorder(body).rev() } -struct MonoReachablePostorder<'a, 'tcx> { - basic_blocks: &'a IndexSlice>, - visited: BitSet, - visit_stack: Vec<(BasicBlock, Successors<'a>)>, - tcx: TyCtxt<'tcx>, - instance: Instance<'tcx>, +/// Lets us plug in some additional logic and data into a Postorder traversal. Or not. +pub trait Customization<'tcx>: Copy { + fn successors<'a>(_: &'a BasicBlockData<'tcx>, _: Self) -> Successors<'a>; } -impl<'a, 'tcx> MonoReachablePostorder<'a, 'tcx> { - fn new( - body: &'a Body<'tcx>, - tcx: TyCtxt<'tcx>, - instance: Instance<'tcx>, - ) -> MonoReachablePostorder<'a, 'tcx> { - let basic_blocks = &body.basic_blocks; - let mut po = MonoReachablePostorder { - basic_blocks, - visited: BitSet::new_empty(basic_blocks.len()), - visit_stack: Vec::new(), - tcx, - instance, - }; - - po.visit(START_BLOCK); - po.traverse_successor(); - po - } - - fn visit(&mut self, bb: BasicBlock) { - if !self.visited.insert(bb) { - return; - } - let data = &self.basic_blocks[bb]; - let successors = data.mono_successors(self.tcx, self.instance); - self.visit_stack.push((bb, successors)); - } - - fn traverse_successor(&mut self) { - while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) { - self.visit(bb); - } +impl<'tcx> Customization<'tcx> for () { + fn successors<'a>(data: &'a BasicBlockData<'tcx>, _: ()) -> Successors<'a> { + data.terminator().successors() } } -impl<'tcx> Iterator for MonoReachablePostorder<'_, 'tcx> { - type Item = BasicBlock; - - fn next(&mut self) -> Option { - let (bb, _) = self.visit_stack.pop()?; - self.traverse_successor(); - Some(bb) +impl<'tcx> Customization<'tcx> for (TyCtxt<'tcx>, Instance<'tcx>) { + fn successors<'a>( + data: &'a BasicBlockData<'tcx>, + (tcx, instance): (TyCtxt<'tcx>, Instance<'tcx>), + ) -> Successors<'a> { + data.mono_successors(tcx, instance) } } @@ -291,7 +266,7 @@ pub fn mono_reachable_reverse_postorder<'a, 'tcx>( tcx: TyCtxt<'tcx>, instance: Instance<'tcx>, ) -> Vec { - let mut iter = MonoReachablePostorder::new(body, tcx, instance); + let mut iter = Postorder::new(&body.basic_blocks, START_BLOCK, (tcx, instance)); let mut items = Vec::with_capacity(body.basic_blocks.len()); while let Some(block) = iter.next() { items.push(block);