Small review fixes

src/librustc/ty/layout.rs

@@ -649,13 +649,11 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                 use SavedLocalEligibility::*;
 
                 let mut assignments: IndexVec<GeneratorSavedLocal, SavedLocalEligibility> =
-                    iter::repeat(Unassigned)
-                    .take(info.field_tys.len())
-                    .collect();
+                    IndexVec::from_elem_n(Unassigned, info.field_tys.len());
 
                 // The saved locals not eligible for overlap. These will get
                 // "promoted" to the prefix of our generator.
-                let mut eligible_locals = BitSet::new_filled(info.field_tys.len());
+                let mut ineligible_locals = BitSet::new_empty(info.field_tys.len());
 
                 // Figure out which of our saved locals are fields in only
                 // one variant. The rest are deemed ineligible for overlap.
@@ -670,7 +668,7 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                                 // point, so it is no longer a candidate.
                                 trace!("removing local {:?} in >1 variant ({:?}, {:?})",
                                        local, variant_index, idx);
-                                eligible_locals.remove(*local);
+                                ineligible_locals.insert(*local);
                                 assignments[*local] = Ineligible(None);
                             }
                             Ineligible(_) => {},
@@ -681,46 +679,50 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                 // Next, check every pair of eligible locals to see if they
                 // conflict.
                 for (local_a, conflicts_a) in info.storage_conflicts.iter_enumerated() {
-                    if !eligible_locals.contains(local_a) {
+                    if ineligible_locals.contains(local_a) {
                         continue;
                     }
 
                     for local_b in conflicts_a.iter() {
-                        // local_a and local_b have overlapping storage, therefore they
+                        // local_a and local_b are storage live at the same time, therefore they
                         // cannot overlap in the generator layout. The only way to guarantee
                         // this is if they are in the same variant, or one is ineligible
                         // (which means it is stored in every variant).
-                        if !eligible_locals.contains(local_b) ||
+                        if ineligible_locals.contains(local_b) ||
                             assignments[local_a] == assignments[local_b]
                         {
                             continue;
                         }
 
                         // If they conflict, we will choose one to make ineligible.
+                        // This is not always optimal; it's just a greedy heuristic
+                        // that seems to produce good results most of the time.
                         let conflicts_b = &info.storage_conflicts[local_b];
                         let (remove, other) = if conflicts_a.count() > conflicts_b.count() {
                             (local_a, local_b)
                         } else {
                             (local_b, local_a)
                         };
-                        eligible_locals.remove(remove);
+                        ineligible_locals.insert(remove);
                         assignments[remove] = Ineligible(None);
                         trace!("removing local {:?} due to conflict with {:?}", remove, other);
                     }
                 }
 
-                let mut ineligible_locals = BitSet::new_filled(info.field_tys.len());
-                ineligible_locals.subtract(&eligible_locals);
-
                 // Write down the order of our locals that will be promoted to
                 // the prefix.
-                for (idx, local) in ineligible_locals.iter().enumerate() {
-                    assignments[local] = Ineligible(Some(idx as u32));
+                {
+                    let mut idx = 0u32;
+                    for local in ineligible_locals.iter() {
+                        assignments[local] = Ineligible(Some(idx));
+                        idx += 1;
+                    }
                 }
                 debug!("generator saved local assignments: {:?}", assignments);
 
                 // Build a prefix layout, including "promoting" all ineligible
-                // locals as part of the prefix.
+                // locals as part of the prefix. We compute the layout of all of
+                // these fields at once to get optimal packing.
                 let discr_index = substs.prefix_tys(def_id, tcx).count();
                 let promoted_tys =
                     ineligible_locals.iter().map(|local| subst_field(info.field_tys[local]));
@@ -733,20 +735,23 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                     StructKind::AlwaysSized)?;
                 let (prefix_size, prefix_align) = (prefix.size, prefix.align);
 
-                // Split the prefix layout into the "outer" fields (upvars and
-                // discriminant) and the "promoted" fields. Promoted fields will
-                // get included in each variant that requested them in
-                // GeneratorLayout.
-                let renumber_indices = |mut index: Vec<u32>| -> Vec<u32> {
-                    debug!("renumber_indices({:?})", index);
-                    let mut inverse_index = (0..index.len() as u32).collect::<Vec<_>>();
-                    inverse_index.sort_unstable_by_key(|i| index[*i as usize]);
+                let recompute_memory_index = |offsets: &Vec<u32>| -> Vec<u32> {
+                    debug!("recompute_memory_index({:?})", offsets);
+                    let mut inverse_index = (0..offsets.len() as u32).collect::<Vec<_>>();
+                    inverse_index.sort_unstable_by_key(|i| offsets[*i as usize]);
+
+                    let mut index = vec![0; offsets.len()];
                     for i in 0..index.len() {
                         index[inverse_index[i] as usize] = i as u32;
                     }
-                    debug!("renumber_indices() => {:?}", index);
+                    debug!("recompute_memory_index() => {:?}", index);
                     index
                 };
+
+                // Split the prefix layout into the "outer" fields (upvars and
+                // discriminant) and the "promoted" fields. Promoted fields will
+                // get included in each variant that requested them in
+                // GeneratorLayout.
                 debug!("prefix = {:#?}", prefix);
                 let (outer_fields, promoted_offsets, promoted_memory_index) = match prefix.fields {
                     FieldPlacement::Arbitrary { offsets, memory_index } => {
@@ -756,11 +761,11 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                             memory_index.split_at(discr_index + 1);
                         let outer_fields = FieldPlacement::Arbitrary {
                             offsets: offsets_a.to_vec(),
-                            memory_index: renumber_indices(memory_index_a.to_vec())
+                            memory_index: recompute_memory_index(&memory_index_a.to_vec())
                         };
                         (outer_fields,
                          offsets_b.to_vec(),
-                         renumber_indices(memory_index_b.to_vec()))
+                         recompute_memory_index(&memory_index_b.to_vec()))
                     }
                     _ => bug!(),
                 };
@@ -769,15 +774,17 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                 let mut align = prefix.align;
                 let variants = info.variant_fields.iter_enumerated().map(|(index, variant_fields)| {
                     // Only include overlap-eligible fields when we compute our variant layout.
-                    let variant_only_tys = variant_fields.iter().flat_map(|local| {
-                        let ty = info.field_tys[*local];
-                        match assignments[*local] {
-                            Unassigned => bug!(),
-                            Assigned(v) if v == index => Some(subst_field(ty)),
-                            Assigned(_) => bug!("assignment does not match variant"),
-                            Ineligible(_) => None,
-                        }
-                    });
+                    let variant_only_tys = variant_fields
+                        .iter()
+                        .filter(|local| {
+                            match assignments[**local] {
+                                Unassigned => bug!(),
+                                Assigned(v) if v == index => true,
+                                Assigned(_) => bug!("assignment does not match variant"),
+                                Ineligible(_) => false,
+                            }
+                        })
+                        .map(|local| subst_field(info.field_tys[*local]));
 
                     let mut variant = univariant_uninterned(
                         &variant_only_tys
@@ -823,7 +830,7 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                     }
                     variant.fields = FieldPlacement::Arbitrary {
                         offsets: combined_offsets,
-                        memory_index: renumber_indices(combined_memory_index),
+                        memory_index: recompute_memory_index(&combined_memory_index),
                     };
 
                     size = size.max(variant.size);

src/librustc_mir/dataflow/mod.rs

@@ -431,7 +431,7 @@ fn for_each_block_location<'tcx, T: BitDenotation<'tcx>>(
         let term_loc = Location { block, statement_index: mir[block].statements.len() };
         analysis.before_terminator_effect(&mut sets, term_loc);
         f(sets.gen_set, term_loc);
-        analysis.before_statement_effect(&mut sets, term_loc);
+        analysis.terminator_effect(&mut sets, term_loc);
         f(sets.gen_set, term_loc);
     }
 }

src/librustc_mir/transform/generator.rs

@@ -394,17 +394,33 @@ impl<'tcx> Visitor<'tcx> for StorageIgnored {
     }
 }
 
+struct LivenessInfo {
+    /// Which locals are live across any suspension point.
+    ///
+    /// GeneratorSavedLocal is indexed in terms of the elements in this set;
+    /// i.e. GeneratorSavedLocal::new(1) corresponds to the second local
+    /// included in this set.
+    live_locals: liveness::LiveVarSet,
+
+    /// The set of saved locals live at each suspension point.
+    live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>,
+
+    /// For every saved local, the set of other saved locals that are
+    /// storage-live at the same time as this local. We cannot overlap locals in
+    /// the layout which have conflicting storage.
+    storage_conflicts: IndexVec<GeneratorSavedLocal, BitSet<GeneratorSavedLocal>>,
+
+    /// For every suspending block, the locals which are storage-live across
+    /// that suspension point.
+    storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>,
+}
+
 fn locals_live_across_suspend_points(
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
     body: &Body<'tcx>,
     source: MirSource<'tcx>,
     movable: bool,
-) -> (
-    liveness::LiveVarSet,
-    Vec<BitSet<GeneratorSavedLocal>>,
-    IndexVec<GeneratorSavedLocal, BitSet<GeneratorSavedLocal>>,
-    FxHashMap<BasicBlock, liveness::LiveVarSet>,
-) {
+) -> LivenessInfo {
     let dead_unwinds = BitSet::new_empty(body.basic_blocks().len());
     let def_id = source.def_id();
 
@@ -434,7 +450,7 @@ fn locals_live_across_suspend_points(
     };
 
     // Calculate the liveness of MIR locals ignoring borrows.
-    let mut set = liveness::LiveVarSet::new_empty(body.local_decls.len());
+    let mut live_locals = liveness::LiveVarSet::new_empty(body.local_decls.len());
     let mut liveness = liveness::liveness_of_locals(
         body,
     );
@@ -489,36 +505,40 @@ fn locals_live_across_suspend_points(
             // Locals live are live at this point only if they are used across
             // suspension points (the `liveness` variable)
             // and their storage is live (the `storage_liveness` variable)
-            storage_liveness.intersect(&liveness.outs[block]);
+            let mut live_locals_here = storage_liveness;
+            live_locals_here.intersect(&liveness.outs[block]);
 
             // The generator argument is ignored
-            storage_liveness.remove(self_arg());
-
-            let live_locals = storage_liveness;
+            live_locals_here.remove(self_arg());
 
             // Add the locals live at this suspension point to the set of locals which live across
             // any suspension points
-            set.union(&live_locals);
+            live_locals.union(&live_locals_here);
 
-            live_locals_at_suspension_points.push(live_locals);
+            live_locals_at_suspension_points.push(live_locals_here);
         }
     }
 
     // Renumber our liveness_map bitsets to include only the locals we are
     // saving.
     let live_locals_at_suspension_points = live_locals_at_suspension_points
         .iter()
-        .map(|live_locals| renumber_bitset(&live_locals, &set))
+        .map(|live_here| renumber_bitset(&live_here, &live_locals))
         .collect();
 
     let storage_conflicts = compute_storage_conflicts(
         body,
-        &set,
+        &live_locals,
         &ignored,
         storage_live,
         storage_live_analysis);
 
-    (set, live_locals_at_suspension_points, storage_conflicts, storage_liveness_map)
+    LivenessInfo {
+        live_locals,
+        live_locals_at_suspension_points,
+        storage_conflicts,
+        storage_liveness: storage_liveness_map,
+    }
 }
 
 /// For every saved local, looks for which locals are StorageLive at the same
@@ -555,14 +575,11 @@ fn compute_storage_conflicts(
         eligible_storage_live.intersect(&stored_locals);
 
         for local in eligible_storage_live.iter() {
-            let mut overlaps = eligible_storage_live.clone();
-            overlaps.remove(local);
-            local_conflicts[local].union(&overlaps);
+            local_conflicts[local].union(&eligible_storage_live);
+        }
 
-            if !overlaps.is_empty() {
-                trace!("at {:?}, local {:?} conflicts with {:?}",
-                       loc, local, overlaps);
-            }
+        if eligible_storage_live.count() > 1 {
+            trace!("at {:?}, eligible_storage_live={:?}", loc, eligible_storage_live);
         }
     });
 
@@ -617,8 +634,9 @@ fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         FxHashMap<BasicBlock, liveness::LiveVarSet>)
 {
     // Use a liveness analysis to compute locals which are live across a suspension point
-    let (live_locals, live_locals_at_suspension_points, storage_conflicts, storage_liveness) =
-        locals_live_across_suspend_points(tcx, body, source, movable);
+    let LivenessInfo {
+        live_locals, live_locals_at_suspension_points, storage_conflicts, storage_liveness
+    } = locals_live_across_suspend_points(tcx, body, source, movable);
 
     // Erase regions from the types passed in from typeck so we can compare them with
     // MIR types
@@ -673,11 +691,11 @@ fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         let mut fields = IndexVec::new();
         for (idx, saved_local) in live_locals.iter().enumerate() {
             fields.push(saved_local);
-            // Note that if a field is included in multiple variants, it will be
-            // added overwritten here. That's fine; fields do not move around
-            // inside generators, so it doesn't matter which variant index we
-            // access them by.
-            remap.insert(locals[saved_local], (tys[saved_local], variant_index, idx));
+            // Note that if a field is included in multiple variants, we will
+            // just use the first one here. That's fine; fields do not move
+            // around inside generators, so it doesn't matter which variant
+            // index we access them by.
+            remap.entry(locals[saved_local]).or_insert((tys[saved_local], variant_index, idx));
         }
         variant_fields.push(fields);
     }