Misc improvements to coroutine transform code

compiler/rustc_middle/src/mir/pretty.rs

@@ -15,6 +15,7 @@ use crate::mir::interpret::{
 };
 use crate::mir::visit::Visitor;
 use crate::mir::*;
+use crate::ty::CoroutineArgsExt;
 
 const INDENT: &str = "    ";
 /// Alignment for lining up comments following MIR statements
@@ -185,9 +186,6 @@ impl<'a, 'tcx> MirDumper<'a, 'tcx> {
             Some(promoted) => write!(w, "::{promoted:?}`")?,
         }
         writeln!(w, " {} {}", self.disambiguator, self.pass_name)?;
-        if let Some(ref layout) = body.coroutine_layout_raw() {
-            writeln!(w, "/* coroutine_layout = {layout:#?} */")?;
-        }
         writeln!(w)?;
         (self.writer.extra_data)(PassWhere::BeforeCFG, w)?;
         write_user_type_annotations(self.tcx(), body, w)?;
@@ -429,6 +427,31 @@ fn write_scope_tree(
         }
     }
 
+    // Coroutine debuginfo.
+    if let Some(layout) = body.coroutine_layout_raw() {
+        for (field, name) in layout.field_names.iter_enumerated() {
+            let source_info = layout.field_tys[field].source_info;
+            if let Some(name) = name
+                && source_info.scope == parent
+            {
+                let indented_debug_info =
+                    format!("{0:1$}coroutine debug {2} => {3:?};", INDENT, indent, name, field);
+
+                if options.include_extra_comments {
+                    writeln!(
+                        w,
+                        "{0:1$} // in {2}",
+                        indented_debug_info,
+                        ALIGN,
+                        comment(tcx, source_info),
+                    )?;
+                } else {
+                    writeln!(w, "{indented_debug_info}")?;
+                }
+            }
+        }
+    }
+
     // Local variable types.
     for (local, local_decl) in body.local_decls.iter_enumerated() {
         if (1..body.arg_count + 1).contains(&local.index()) {
@@ -530,6 +553,50 @@ impl Debug for VarDebugInfo<'_> {
     }
 }
 
+fn write_coroutine_layout<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    layout: &CoroutineLayout<'_>,
+    w: &mut dyn io::Write,
+    options: PrettyPrintMirOptions,
+) -> io::Result<()> {
+    let CoroutineLayout {
+        field_tys,
+        field_names: _, // Dumped in scope tree with debug info.
+        variant_fields,
+        variant_source_info,
+        storage_conflicts,
+    } = layout;
+
+    writeln!(w, "{INDENT}coroutine layout {{")?;
+
+    for (field, CoroutineSavedTy { ty, source_info, ignore_for_traits }) in
+        field_tys.iter_enumerated()
+    {
+        let ignore_for_traits = if *ignore_for_traits { " (ignored for traits)" } else { "" };
+        let indented_body = format!("{INDENT}{INDENT}field {field:?}: {ty}{ignore_for_traits};",);
+        if options.include_extra_comments {
+            writeln!(w, "{0:ALIGN$} // in {1}", indented_body, comment(tcx, *source_info))?;
+        } else {
+            writeln!(w, "{}", indented_body)?;
+        }
+    }
+
+    writeln!(w, "{INDENT}{INDENT}variant_fields = {{")?;
+    for (variant, fields) in variant_fields.iter_enumerated() {
+        let variant_name = ty::CoroutineArgs::variant_name(variant);
+        let header = format!("{INDENT}{INDENT}{INDENT}{variant_name:9}({variant:?}): {fields:?},");
+        if options.include_extra_comments {
+            let source_info = variant_source_info[variant];
+            writeln!(w, "{0:ALIGN$} // in {1}", header, comment(tcx, source_info))?;
+        } else {
+            writeln!(w, "{}", header)?;
+        }
+    }
+    writeln!(w, "{INDENT}{INDENT}}}")?;
+    writeln!(w, "{INDENT}{INDENT}storage_conflicts = {storage_conflicts:?}")?;
+    writeln!(w, "{INDENT}}}")
+}
+
 /// Write out a human-readable textual representation of the MIR's `fn` type and the types of its
 /// local variables (both user-defined bindings and compiler temporaries).
 fn write_mir_intro<'tcx>(
@@ -541,6 +608,10 @@ fn write_mir_intro<'tcx>(
     write_mir_sig(tcx, body, w)?;
     writeln!(w, "{{")?;
 
+    if let Some(ref layout) = body.coroutine_layout_raw() {
+        write_coroutine_layout(tcx, layout, w, options)?;
+    }
+
     // construct a scope tree and write it out
     let mut scope_tree: FxHashMap<SourceScope, Vec<SourceScope>> = Default::default();
     for (index, scope_data) in body.source_scopes.iter_enumerated() {

compiler/rustc_mir_transform/src/coroutine.rs

@@ -1077,31 +1077,42 @@ fn compute_layout<'tcx>(
 /// Replaces the entry point of `body` with a block that switches on the coroutine discriminant and
 /// dispatches to blocks according to `cases`.
 ///
-/// After this function, the former entry point of the function will be bb1.
+/// After this function, the former entry point of the function will be the last block.
 fn insert_switch<'tcx>(
     body: &mut Body<'tcx>,
     cases: Vec<(usize, BasicBlock)>,
     transform: &TransformVisitor<'tcx>,
     default_block: BasicBlock,
 ) {
     let (assign, discr) = transform.get_discr(body);
-    let switch_targets =
-        SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
-    let switch = TerminatorKind::SwitchInt { discr: Operand::Move(discr), targets: switch_targets };
 
-    let source_info = SourceInfo::outermost(body.span);
-    body.basic_blocks_mut().raw.insert(
-        0,
-        BasicBlockData::new_stmts(
-            vec![assign],
-            Some(Terminator { source_info, kind: switch }),
-            false,
-        ),
+    // MIR validation ensures that no block targets `ENTRY_BLOCK`.
+    #[cfg(debug_assertions)]
+    for bb in body.basic_blocks.iter() {
+        for target in bb.terminator().successors() {
+            assert_ne!(target, START_BLOCK);
+        }
+    }
+
+    // Add the switch as entry block, and put the former entry block at the end.
+    let former_entry = std::mem::replace(
+        &mut body.basic_blocks_mut()[START_BLOCK],
+        BasicBlockData::new_stmts(vec![assign], None, false),
     );
+    let former_entry = body.basic_blocks_mut().push(former_entry);
 
-    for b in body.basic_blocks_mut().iter_mut() {
-        b.terminator_mut().successors_mut(|target| *target += 1);
+    // We may point to `START_BLOCK` in our `cases`, replace it with `former_entry`.
+    let mut switch_targets =
+        SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
+    for bb in switch_targets.all_targets_mut() {
+        if *bb == START_BLOCK {
+            *bb = former_entry;
+        }
     }
+
+    let switch = TerminatorKind::SwitchInt { discr: Operand::Move(discr), targets: switch_targets };
+    body.basic_blocks_mut()[START_BLOCK].terminator =
+        Some(Terminator { source_info: SourceInfo::outermost(body.span), kind: switch });
 }
 
 fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock {
@@ -1172,41 +1183,8 @@ fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
         return false;
     }
 
-    // Unwinds can only start at certain terminators.
-    for block in body.basic_blocks.iter() {
-        match block.terminator().kind {
-            // These never unwind.
-            TerminatorKind::Goto { .. }
-            | TerminatorKind::SwitchInt { .. }
-            | TerminatorKind::UnwindTerminate(_)
-            | TerminatorKind::Return
-            | TerminatorKind::Unreachable
-            | TerminatorKind::CoroutineDrop
-            | TerminatorKind::FalseEdge { .. }
-            | TerminatorKind::FalseUnwind { .. } => {}
-
-            // Resume will *continue* unwinding, but if there's no other unwinding terminator it
-            // will never be reached.
-            TerminatorKind::UnwindResume => {}
-
-            TerminatorKind::Yield { .. } => {
-                unreachable!("`can_unwind` called before coroutine transform")
-            }
-
-            // These may unwind.
-            TerminatorKind::Drop { .. }
-            | TerminatorKind::Call { .. }
-            | TerminatorKind::InlineAsm { .. }
-            | TerminatorKind::Assert { .. } => return true,
-
-            TerminatorKind::TailCall { .. } => {
-                unreachable!("tail calls can't be present in generators")
-            }
-        }
-    }
-
-    // If we didn't find an unwinding terminator, the function cannot unwind.
-    false
+    // If we don't find an unwinding terminator, the function cannot unwind.
+    body.basic_blocks.iter().any(|block| block.terminator().unwind().is_some())
 }
 
 // Poison the coroutine when it unwinds