1 files changed, 34 insertions, 3 deletions
diff --git a/src/cmd/compile/internal/ssa/layout.go b/src/cmd/compile/internal/ssa/layout.go
index 338cd91c47..30b7b97d04 100644
--- a/src/cmd/compile/internal/ssa/layout.go
+++ b/src/cmd/compile/internal/ssa/layout.go
@@ -46,13 +46,44 @@ func layoutOrder(f *Func) []*Block {
 	exit := f.newSparseSet(f.NumBlocks()) // exit blocks
 	defer f.retSparseSet(exit)
 
-	// Initialize indegree of each block
+	// Populate idToBlock and find exit blocks.
 	for _, b := range f.Blocks {
 		idToBlock[b.ID] = b
 		if b.Kind == BlockExit {
-			// exit blocks are always scheduled last
-			// TODO: also add blocks post-dominated by exit blocks
 			exit.add(b.ID)
+		}
+	}
+
+	// Expand exit to include blocks post-dominated by exit blocks.
+	for {
+		changed := false
+		for _, id := range exit.contents() {
+			b := idToBlock[id]
+		NextPred:
+			for _, pe := range b.Preds {
+				p := pe.b
+				if exit.contains(p.ID) {
+					continue
+				}
+				for _, s := range p.Succs {
+					if !exit.contains(s.b.ID) {
+						continue NextPred
+					}
+				}
+				// All Succs are in exit; add p.
+				exit.add(p.ID)
+				changed = true
+			}
+		}
+		if !changed {
+			break
+		}
+	}
+
+	// Initialize indegree of each block
+	for _, b := range f.Blocks {
+		if exit.contains(b.ID) {
+			// exit blocks are always scheduled last
 			continue
 		}
 		indegree[b.ID] = len(b.Preds)