@@ -232,6 +232,155 @@ pub fn postorder<'a, 'tcx>(
232232 reverse_postorder ( body) . rev ( )
233233}
234234
235+ pub struct MonoReachablePostorder < ' a , ' tcx > {
236+ basic_blocks : & ' a IndexSlice < BasicBlock , BasicBlockData < ' tcx > > ,
237+ visited : BitSet < BasicBlock > ,
238+ visit_stack : Vec < ( BasicBlock , Successors < ' a > ) > ,
239+ locals : BitSet < Local > ,
240+ tcx : TyCtxt < ' tcx > ,
241+ instance : Instance < ' tcx > ,
242+ }
243+
244+ impl < ' a , ' tcx > MonoReachablePostorder < ' a , ' tcx > {
245+ pub fn new (
246+ body : & ' a Body < ' tcx > ,
247+ tcx : TyCtxt < ' tcx > ,
248+ instance : Instance < ' tcx > ,
249+ ) -> MonoReachablePostorder < ' a , ' tcx > {
250+ let basic_blocks = & body. basic_blocks ;
251+ let mut po = MonoReachablePostorder {
252+ basic_blocks,
253+ visited : BitSet :: new_empty ( basic_blocks. len ( ) ) ,
254+ visit_stack : Vec :: new ( ) ,
255+ locals : BitSet :: new_empty ( body. local_decls . len ( ) ) ,
256+ tcx,
257+ instance,
258+ } ;
259+
260+ let root = START_BLOCK ;
261+ let data = & po. basic_blocks [ root] ;
262+
263+ UsedLocals { locals : & mut po. locals } . visit_basic_block_data ( root, data) ;
264+ if let Some ( ref term) = data. terminator {
265+ po. visited . insert ( root) ;
266+
267+ let successors = if let Some ( ( bits, targets) ) =
268+ Body :: try_const_mono_switchint ( tcx, instance, data)
269+ {
270+ targets. successors_for_value ( bits)
271+ } else {
272+ term. successors ( )
273+ } ;
274+
275+ po. visit_stack . push ( ( root, successors) ) ;
276+ po. traverse_successor ( ) ;
277+ }
278+
279+ po
280+ }
281+
282+ fn traverse_successor ( & mut self ) {
283+ // This is quite a complex loop due to 1. the borrow checker not liking it much
284+ // and 2. what exactly is going on is not clear
285+ //
286+ // It does the actual traversal of the graph, while the `next` method on the iterator
287+ // just pops off of the stack. `visit_stack` is a stack containing pairs of nodes and
288+ // iterators over the successors of those nodes. Each iteration attempts to get the next
289+ // node from the top of the stack, then pushes that node and an iterator over the
290+ // successors to the top of the stack. This loop only grows `visit_stack`, stopping when
291+ // we reach a child that has no children that we haven't already visited.
292+ //
293+ // For a graph that looks like this:
294+ //
295+ // A
296+ // / \
297+ // / \
298+ // B C
299+ // | |
300+ // | |
301+ // | D
302+ // \ /
303+ // \ /
304+ // E
305+ //
306+ // The state of the stack starts out with just the root node (`A` in this case);
307+ // [(A, [B, C])]
308+ //
309+ // When the first call to `traverse_successor` happens, the following happens:
310+ //
311+ // [(C, [D]), // `C` taken from the successors of `A`, pushed to the
312+ // // top of the stack along with the successors of `C`
313+ // (A, [B])]
314+ //
315+ // [(D, [E]), // `D` taken from successors of `C`, pushed to stack
316+ // (C, []),
317+ // (A, [B])]
318+ //
319+ // [(E, []), // `E` taken from successors of `D`, pushed to stack
320+ // (D, []),
321+ // (C, []),
322+ // (A, [B])]
323+ //
324+ // Now that the top of the stack has no successors we can traverse, each item will
325+ // be popped off during iteration until we get back to `A`. This yields [E, D, C].
326+ //
327+ // When we yield `C` and call `traverse_successor`, we push `B` to the stack, but
328+ // since we've already visited `E`, that child isn't added to the stack. The last
329+ // two iterations yield `B` and finally `A` for a final traversal of [E, D, C, B, A]
330+ while let Some ( bb) = self . visit_stack . last_mut ( ) . and_then ( |( _, iter) | iter. next_back ( ) ) {
331+ if self . visited . insert ( bb) {
332+ let data = & self . basic_blocks [ bb] ;
333+ UsedLocals { locals : & mut self . locals } . visit_basic_block_data ( bb, data) ;
334+
335+ let Some ( term) = & data. terminator else {
336+ continue ;
337+ } ;
338+
339+ let successors = if let Some ( ( bits, targets) ) =
340+ Body :: try_const_mono_switchint ( self . tcx , self . instance , data)
341+ {
342+ targets. successors_for_value ( bits)
343+ } else {
344+ term. successors ( )
345+ } ;
346+
347+ self . visit_stack . push ( ( bb, successors) ) ;
348+ }
349+ }
350+ }
351+ }
352+
353+ impl < ' tcx > Iterator for MonoReachablePostorder < ' _ , ' tcx > {
354+ type Item = BasicBlock ;
355+
356+ fn next ( & mut self ) -> Option < BasicBlock > {
357+ let ( bb, _) = self . visit_stack . pop ( ) ?;
358+ self . traverse_successor ( ) ;
359+
360+ Some ( bb)
361+ }
362+
363+ fn size_hint ( & self ) -> ( usize , Option < usize > ) {
364+ // All the blocks, minus the number of blocks we've visited.
365+ let remaining = self . basic_blocks . len ( ) - self . visited . count ( ) ;
366+ ( remaining, Some ( remaining) )
367+ }
368+ }
369+
370+ pub fn mono_reachable_reverse_postorder < ' a , ' tcx > (
371+ body : & ' a Body < ' tcx > ,
372+ tcx : TyCtxt < ' tcx > ,
373+ instance : Instance < ' tcx > ,
374+ ) -> ( Vec < BasicBlock > , BitSet < Local > ) {
375+ let mut iter = MonoReachablePostorder :: new ( body, tcx, instance) ;
376+ let mut items = Vec :: with_capacity ( body. basic_blocks . len ( ) ) ;
377+ while let Some ( block) = iter. next ( ) {
378+ items. push ( block) ;
379+ }
380+ items. reverse ( ) ;
381+ ( items, iter. locals )
382+ }
383+
235384/// Returns an iterator over all basic blocks reachable from the `START_BLOCK` in no particular
236385/// order.
237386///
@@ -320,6 +469,22 @@ pub struct MonoReachable<'a, 'tcx> {
320469 worklist : BitSet < BasicBlock > ,
321470}
322471
472+ struct UsedLocals < ' a > {
473+ locals : & ' a mut BitSet < Local > ,
474+ }
475+
476+ use crate :: mir:: visit:: Visitor ;
477+ impl < ' a , ' tcx > Visitor < ' tcx > for UsedLocals < ' a > {
478+ fn visit_local (
479+ & mut self ,
480+ local : Local ,
481+ _ctx : crate :: mir:: visit:: PlaceContext ,
482+ _location : Location ,
483+ ) {
484+ self . locals . insert ( local) ;
485+ }
486+ }
487+
323488impl < ' a , ' tcx > MonoReachable < ' a , ' tcx > {
324489 pub fn new (
325490 body : & ' a Body < ' tcx > ,
0 commit comments