make sure that retagging body arguments stays the first instruction

Author: dingxiangfei2009

compiler/rustc_mir_transform/src/coroutine/relocate_upvars.rs

@@ -102,6 +102,7 @@
 
 use std::borrow::Cow;
 
+use itertools::Itertools;
 use rustc_abi::FieldIdx;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_index::{IndexSlice, IndexVec};
@@ -261,6 +262,28 @@ impl RelocateUpvars {
             | ty::InstanceKind::FnPtrAddrShim(..) => unreachable!(),
         };
 
+        // HACK: in case `AddRetag` is already run, we have one `Retag` at the body entrance
+        // so we need to make sure that first `Retag` is run.
+        let last_retag = body.basic_blocks[START_BLOCK]
+            .statements
+            .iter()
+            .find_position(|stmt| !matches!(stmt.kind, StatementKind::Retag(_, _)));
+        let retagged_start_block = if let Some((index, _)) = last_retag {
+            let span = body.span;
+            let bbs = body.basic_blocks_mut();
+            let stmts = bbs[START_BLOCK].statements.split_off(index);
+            let terminator = bbs[START_BLOCK].terminator.take();
+            let split_start_block = bbs.push(BasicBlockData::new_stmts(stmts, terminator, false));
+            bbs[START_BLOCK].statements.shrink_to_fit();
+            bbs[START_BLOCK].terminator = Some(Terminator {
+                kind: TerminatorKind::Goto { target: split_start_block },
+                source_info: SourceInfo::outermost(span),
+            });
+            Some(split_start_block)
+        } else {
+            None
+        };
+
         if let Some(mir_dumper) = MirDumper::new(tcx, "RelocateUpvars", body) {
             mir_dumper.set_disambiguator(&"before").dump_mir(body);
         }
@@ -308,7 +331,7 @@ impl RelocateUpvars {
         SubstituteUpvarVisitor { tcx, mappings: &substitution_mapping }.visit_body(body);
 
         rewrite_drop_coroutine_struct(tcx, body, &substitution_mapping);
-        insert_substitution_prologue(body, &substitution_mapping);
+        insert_substitution_prologue(body, retagged_start_block, &substitution_mapping);
         patch_missing_storage_deads(tcx, body, &substitution_mapping);
         hydrate_var_debug_info(body, &substitution_mapping);
         if let Some(mir_dumper) = MirDumper::new(tcx, "RelocateUpvars", body) {
@@ -445,10 +468,11 @@ fn rewrite_drop_coroutine_struct<'tcx>(
 
 fn insert_substitution_prologue<'tcx>(
     body: &mut Body<'tcx>,
+    retagged_start_block: Option<BasicBlock>,
     substitution_mapping: &IndexSlice<FieldIdx, UpvarSubstitution<'tcx>>,
 ) {
     let mut patch = MirPatch::new(body);
-    let mut stmts = Vec::with_capacity(2 * substitution_mapping.len());
+    let mut stmts = Vec::with_capacity(5 * substitution_mapping.len());
     for &UpvarSubstitution { local, reloc, upvar_place, span, name: _ } in substitution_mapping {
         // For each upvar-local _$i
         let source_info = SourceInfo::outermost(span);
@@ -481,22 +505,30 @@ fn insert_substitution_prologue<'tcx>(
     let source_info = SourceInfo::outermost(body.span);
     let prologue = patch.new_block(BasicBlockData::new_stmts(
         stmts,
-        Some(Terminator { source_info, kind: TerminatorKind::Goto { target: START_BLOCK } }),
+        Some(Terminator {
+            source_info,
+            kind: TerminatorKind::Goto { target: retagged_start_block.unwrap_or(START_BLOCK) },
+        }),
         false,
     ));
     patch.apply(body);
 
     // Manually patch so that prologue is the new entry-point
-    let preds = body.basic_blocks.predecessors()[START_BLOCK].clone();
-    let basic_blocks = body.basic_blocks.as_mut();
-    for pred in preds {
-        basic_blocks[pred].terminator_mut().successors_mut(|target| {
-            if *target == START_BLOCK {
-                *target = prologue;
-            }
-        });
+    if retagged_start_block.is_some() {
+        let basic_blocks = body.basic_blocks.as_mut();
+        basic_blocks[START_BLOCK].terminator_mut().successors_mut(|target| *target = prologue);
+    } else {
+        let preds = body.basic_blocks.predecessors()[START_BLOCK].clone();
+        let basic_blocks = body.basic_blocks.as_mut();
+        for pred in preds {
+            basic_blocks[pred].terminator_mut().successors_mut(|target| {
+                if *target == START_BLOCK {
+                    *target = prologue;
+                }
+            });
+        }
+        basic_blocks.swap(START_BLOCK, prologue);
     }
-    basic_blocks.swap(START_BLOCK, prologue);
 }
 
 /// Occasionally there are upvar locals left without `StorageDead` because

tests/mir-opt/coroutine_relocate_upvars.main-{closure#0}.RelocateUpvars.after.panic-abort.mir

@@ -16,12 +16,7 @@ yields ()
     let _10: &mut std::string::String;
 
     bb0: {
-        StorageLive(_9);
-        StorageLive(_10);
-        _10 = move (_1.0: &mut std::string::String);
-        _9 = move _10;
-        StorageDead(_10);
-        goto -> bb18;
+        goto -> bb19;
     }
 
     bb1: {
@@ -47,7 +42,7 @@ yields ()
         StorageDead(_5);
         StorageDead(_4);
         _0 = const ();
-        goto -> bb13;
+        goto -> bb14;
     }
 
     bb5: {
@@ -61,55 +56,64 @@ yields ()
     }
 
     bb7: {
-        goto -> bb17;
+        goto -> bb18;
     }
 
     bb8: {
         coroutine_drop;
     }
 
-    bb9 (cleanup): {
-        unreachable;
+    bb9: {
+        StorageLive(_3);
+        _3 = String::new() -> [return: bb1, unwind unreachable];
     }
 
     bb10 (cleanup): {
-        StorageDead(_9);
-        goto -> bb9;
+        unreachable;
     }
 
     bb11 (cleanup): {
+        StorageDead(_9);
         goto -> bb10;
     }
 
-    bb12: {
+    bb12 (cleanup): {
+        goto -> bb11;
+    }
+
+    bb13: {
         StorageDead(_9);
         goto -> bb5;
     }
 
-    bb13: {
-        goto -> bb12;
+    bb14: {
+        goto -> bb13;
     }
 
-    bb14 (cleanup): {
+    bb15 (cleanup): {
         StorageDead(_9);
-        goto -> bb9;
+        goto -> bb10;
     }
 
-    bb15 (cleanup): {
-        goto -> bb14;
+    bb16 (cleanup): {
+        goto -> bb15;
     }
 
-    bb16: {
+    bb17: {
         StorageDead(_9);
         goto -> bb8;
     }
 
-    bb17: {
-        goto -> bb16;
+    bb18: {
+        goto -> bb17;
     }
 
-    bb18: {
-        StorageLive(_3);
-        _3 = String::new() -> [return: bb1, unwind unreachable];
+    bb19: {
+        StorageLive(_9);
+        StorageLive(_10);
+        _10 = move (_1.0: &mut std::string::String);
+        _9 = move _10;
+        StorageDead(_10);
+        goto -> bb9;
     }
 }

tests/mir-opt/coroutine_relocate_upvars.main-{closure#0}.RelocateUpvars.after.panic-unwind.mir

@@ -16,12 +16,7 @@ yields ()
     let _10: &mut std::string::String;
 
     bb0: {
-        StorageLive(_9);
-        StorageLive(_10);
-        _10 = move (_1.0: &mut std::string::String);
-        _9 = move _10;
-        StorageDead(_10);
-        goto -> bb23;
+        goto -> bb24;
     }
 
     bb1: {
@@ -53,7 +48,7 @@ yields ()
         StorageDead(_5);
         StorageDead(_4);
         _0 = const ();
-        goto -> bb15;
+        goto -> bb16;
     }
 
     bb6: {
@@ -67,7 +62,7 @@ yields ()
     }
 
     bb8: {
-        goto -> bb19;
+        goto -> bb20;
     }
 
     bb9: {
@@ -76,64 +71,73 @@ yields ()
 
     bb10 (cleanup): {
         StorageDead(_3);
-        goto -> bb22;
+        goto -> bb23;
     }
 
     bb11 (cleanup): {
         resume;
     }
 
-    bb12 (cleanup): {
+    bb12: {
+        StorageLive(_3);
+        _3 = String::new() -> [return: bb1, unwind: bb10];
+    }
+
+    bb13 (cleanup): {
         StorageDead(_9);
         goto -> bb11;
     }
 
-    bb13 (cleanup): {
-        goto -> bb12;
+    bb14 (cleanup): {
+        goto -> bb13;
     }
 
-    bb14: {
+    bb15: {
         StorageDead(_9);
         goto -> bb6;
     }
 
-    bb15: {
-        goto -> bb14;
+    bb16: {
+        goto -> bb15;
     }
 
-    bb16 (cleanup): {
+    bb17 (cleanup): {
         StorageDead(_9);
         goto -> bb11;
     }
 
-    bb17 (cleanup): {
-        goto -> bb16;
+    bb18 (cleanup): {
+        goto -> bb17;
     }
 
-    bb18: {
+    bb19: {
         StorageDead(_9);
         goto -> bb9;
     }
 
-    bb19: {
-        goto -> bb18;
+    bb20: {
+        goto -> bb19;
     }
 
-    bb20 (cleanup): {
+    bb21 (cleanup): {
         terminate(cleanup);
     }
 
-    bb21 (cleanup): {
+    bb22 (cleanup): {
         StorageDead(_9);
         goto -> bb11;
     }
 
-    bb22 (cleanup): {
-        goto -> bb21;
+    bb23 (cleanup): {
+        goto -> bb22;
     }
 
-    bb23: {
-        StorageLive(_3);
-        _3 = String::new() -> [return: bb1, unwind: bb10];
+    bb24: {
+        StorageLive(_9);
+        StorageLive(_10);
+        _10 = move (_1.0: &mut std::string::String);
+        _9 = move _10;
+        StorageDead(_10);
+        goto -> bb12;
     }
 }

tests/mir-opt/coroutine_relocate_upvars.main-{closure#0}.RelocateUpvars.before.panic-abort.mir

@@ -13,8 +13,7 @@ yields ()
     let mut _8: &mut std::string::String;
 
     bb0: {
-        StorageLive(_3);
-        _3 = String::new() -> [return: bb1, unwind unreachable];
+        goto -> bb9;
     }
 
     bb1: {
@@ -60,4 +59,9 @@ yields ()
     bb8: {
         coroutine_drop;
     }
+
+    bb9: {
+        StorageLive(_3);
+        _3 = String::new() -> [return: bb1, unwind unreachable];
+    }
 }

tests/mir-opt/coroutine_relocate_upvars.main-{closure#0}.RelocateUpvars.before.panic-unwind.mir

@@ -13,8 +13,7 @@ yields ()
     let mut _8: &mut std::string::String;
 
     bb0: {
-        StorageLive(_3);
-        _3 = String::new() -> [return: bb1, unwind: bb10];
+        goto -> bb12;
     }
 
     bb1: {
@@ -75,4 +74,9 @@ yields ()
     bb11 (cleanup): {
         resume;
     }
+
+    bb12: {
+        StorageLive(_3);
+        _3 = String::new() -> [return: bb1, unwind: bb10];
+    }
 }