MIR SRoA: Update escaping_locals to only be about escaping

Author: scottmcm

compiler/rustc_mir_transform/src/sroa.rs

@@ -31,8 +31,9 @@ impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {
         let mut excluded = excluded_locals(body);
         let typing_env = body.typing_env(tcx);
         loop {
+            add_type_based_exclusions(tcx, &mut excluded, body);
             debug!(?excluded);
-            let escaping = escaping_locals(tcx, &excluded, body);
+            let escaping = escaping_locals(&excluded, body);
             debug!(?escaping);
             let replacements = compute_flattening(tcx, typing_env, body, escaping);
             debug!(?replacements);
@@ -55,18 +56,13 @@ impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {
     }
 }
 
-/// Identify all locals that are not eligible for SROA.
-///
-/// There are 3 cases:
-/// - the aggregated local is used or passed to other code (function parameters and arguments);
-/// - the locals is a union or an enum;
-/// - the local's address is taken, and thus the relative addresses of the fields are observable to
-///   client code.
-fn escaping_locals<'tcx>(
+/// Also exclude locals whose types are ineligible for SRoA because of
+/// implementation restrictions (whether here or in codegen later).
+fn add_type_based_exclusions<'tcx>(
     tcx: TyCtxt<'tcx>,
-    excluded: &DenseBitSet<Local>,
+    excluded: &mut DenseBitSet<Local>,
     body: &Body<'tcx>,
-) -> DenseBitSet<Local> {
+) {
     let is_excluded_ty = |ty: Ty<'tcx>| {
         if ty.is_union() || ty.is_enum() {
             return true;
@@ -86,10 +82,28 @@ fn escaping_locals<'tcx>(
         false
     };
 
+    for (local, decl) in body.local_decls().iter_enumerated() {
+        if is_excluded_ty(decl.ty) {
+            excluded.insert(local);
+        }
+    }
+}
+
+/// Identify all locals that are not eligible for SROA.
+///
+/// There are 3 cases:
+/// - the aggregated local is used or passed to other code (function parameters and arguments);
+/// - the locals is a union or an enum;
+/// - the local's address is taken, and thus the relative addresses of the fields are observable to
+///   client code.
+fn escaping_locals<'tcx>(
+    excluded: &DenseBitSet<Local>,
+    body: &Body<'tcx>,
+) -> DenseBitSet<Local> {
     let mut set = DenseBitSet::new_empty(body.local_decls.len());
     set.insert_range(RETURN_PLACE..=Local::from_usize(body.arg_count));
-    for (local, decl) in body.local_decls().iter_enumerated() {
-        if excluded.contains(local) || is_excluded_ty(decl.ty) {
+    for local in body.local_decls().indices() {
+        if excluded.contains(local) {
             set.insert(local);
         }
     }