@@ -25,6 +25,20 @@ pub fn collect_roots(module: &Module) -> FxHashSet<Word> {
25
25
rooted
26
26
}
27
27
28
+ // Exactly the same as Function::all_inst_iter, except return type is `impl DoubleEndedIterator`
29
+ // instead of `impl Iterator`
30
+ fn all_inst_iter ( func : & Function ) -> impl DoubleEndedIterator < Item = & Instruction > {
31
+ func. def
32
+ . iter ( )
33
+ . chain ( func. parameters . iter ( ) )
34
+ . chain (
35
+ func. blocks
36
+ . iter ( )
37
+ . flat_map ( |b| b. label . iter ( ) . chain ( b. instructions . iter ( ) ) ) ,
38
+ )
39
+ . chain ( func. end . iter ( ) )
40
+ }
41
+
28
42
fn spread_roots ( module : & Module , rooted : & mut FxHashSet < Word > ) -> bool {
29
43
let mut any = false ;
30
44
for inst in module. global_inst_iter ( ) {
@@ -40,18 +54,7 @@ fn spread_roots(module: &Module, rooted: &mut FxHashSet<Word>) -> bool {
40
54
// earlier insts, by reversing the iteration order, we're more likely to root the
41
55
// entire relevant function at once.
42
56
// See https://github.com/EmbarkStudios/rust-gpu/pull/691#discussion_r681477091
43
- for inst in func
44
- . end
45
- . iter ( )
46
- . chain (
47
- func. blocks
48
- . iter ( )
49
- . rev ( )
50
- . flat_map ( |b| b. instructions . iter ( ) . rev ( ) . chain ( b. label . iter ( ) ) ) ,
51
- )
52
- . chain ( func. parameters . iter ( ) . rev ( ) )
53
- . chain ( func. def . iter ( ) )
54
- {
57
+ for inst in all_inst_iter ( func) . rev ( ) {
55
58
if !instruction_is_pure ( inst) {
56
59
any |= root ( inst, rooted) ;
57
60
} else if let Some ( id) = inst. result_id {
@@ -111,13 +114,13 @@ fn kill_unrooted(module: &mut Module, rooted: &FxHashSet<Word>) {
111
114
module
112
115
. functions
113
116
. retain ( |f| is_rooted ( f. def . as_ref ( ) . unwrap ( ) , rooted) ) ;
114
- module. functions . iter_mut ( ) . for_each ( |fun| {
115
- fun. blocks . iter_mut ( ) . for_each ( |block| {
117
+ for fun in & mut module. functions {
118
+ for block in & mut fun. blocks {
116
119
block
117
120
. instructions
118
121
. retain ( |inst| !instruction_is_pure ( inst) || is_rooted ( inst, rooted) ) ;
119
- } ) ;
120
- } ) ;
122
+ }
123
+ }
121
124
}
122
125
123
126
pub fn dce_phi ( func : & mut Function ) {
0 commit comments