@@ -27,12 +27,17 @@ func loopRotate(f *Func) {
27
27
return
28
28
}
29
29
30
+ idToIdx := make ([]int , f .NumBlocks ())
31
+ for i , b := range f .Blocks {
32
+ idToIdx [b .ID ] = i
33
+ }
34
+
30
35
// Set of blocks we're moving, by ID.
31
36
move := map [ID ]struct {}{}
32
37
33
- // Map from block ID to the moving block that should
38
+ // Map from block ID to the moving blocks that should
34
39
// come right after it.
35
- after := map [ID ]* Block {}
40
+ after := map [ID ][] * Block {}
36
41
37
42
// Check each loop header and decide if we want to move it.
38
43
for _ , loop := range loopnest .loops {
@@ -50,10 +55,27 @@ func loopRotate(f *Func) {
50
55
if p == nil || p == b {
51
56
continue
52
57
}
58
+ after [p .ID ] = []* Block {b }
59
+ for {
60
+ nextIdx := idToIdx [b .ID ] + 1
61
+ if nextIdx >= len (f .Blocks ) { // reached end of function (maybe impossible?)
62
+ break
63
+ }
64
+ nextb := f .Blocks [nextIdx ]
65
+ if nextb == p { // original loop precedessor is next
66
+ break
67
+ }
68
+ if loopnest .b2l [nextb .ID ] != loop { // about to leave loop
69
+ break
70
+ }
71
+ after [p .ID ] = append (after [p .ID ], nextb )
72
+ b = nextb
73
+ }
53
74
54
75
// Place b after p.
55
- move [b .ID ] = struct {}{}
56
- after [p .ID ] = b
76
+ for _ , b := range after [p .ID ] {
77
+ move [b .ID ] = struct {}{}
78
+ }
57
79
}
58
80
59
81
// Move blocks to their destinations in a single pass.
@@ -67,7 +89,7 @@ func loopRotate(f *Func) {
67
89
}
68
90
f .Blocks [j ] = b
69
91
j ++
70
- if a := after [b .ID ]; a != nil {
92
+ for _ , a := range after [b .ID ] {
71
93
if j > i {
72
94
f .Fatalf ("head before tail in loop %s" , b )
73
95
}
0 commit comments