WIP

Author: Kobzol

compiler/rustc_codegen_llvm/src/type_of.rs

@@ -118,6 +118,13 @@ fn struct_llfields<'a, 'tcx>(
             target_offset,
             effective_field_align.bytes()
         );
+        if target_offset < offset {
+            println!("BUG!");
+            tracing::info!("BUG LAYOUT: {:#?}", layout);
+            tracing::info!("i={}, offset={:?}, target_offset={:?}", i, offset, target_offset);
+            tracing::info!("field0={:#?}", layout.field(cx, 0));
+            tracing::info!("field1={:#?}", layout.field(cx, 1));
+        }
         assert!(target_offset >= offset);
         let padding = target_offset - offset;
         if padding != Size::ZERO {

compiler/rustc_middle/src/ty/layout.rs

@@ -1418,11 +1418,15 @@ enum SavedLocalEligibility {
 // Also included in the layout are the upvars and the discriminant.
 // These are included as fields on the "outer" layout; they are not part
 // of any variant.
+// TODO: change documentation
 impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
     /// Compute the eligibility and assignment of each local.
+    /// Saved locals corresponding to upvars will not be returned in the return value, as they are
+    /// already stored (saved) in the generator implicitly.
     fn generator_saved_local_eligibility(
         &self,
         info: &GeneratorLayout<'tcx>,
+        upvar_count: usize,
     ) -> (BitSet<GeneratorSavedLocal>, IndexVec<GeneratorSavedLocal, SavedLocalEligibility>) {
         use SavedLocalEligibility::*;
 
@@ -1498,7 +1502,9 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
             let mut used_variants = BitSet::new_empty(info.variant_fields.len());
             for assignment in &assignments {
                 if let Assigned(idx) = assignment {
-                    used_variants.insert(*idx);
+                    if *idx != VariantIdx::new(0) {
+                        used_variants.insert(*idx);
+                    }
                 }
             }
             if used_variants.count() < 2 {
@@ -1511,8 +1517,14 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
 
         // Write down the order of our locals that will be promoted to the prefix.
         {
+            for upvar in 0..upvar_count {
+                let local = GeneratorSavedLocal::new(upvar);
+                ineligible_locals.remove(local);
+                assignments[local] = Ineligible(Some(upvar as u32));
+            }
             for (idx, local) in ineligible_locals.iter().enumerate() {
-                assignments[local] = Ineligible(Some(idx as u32));
+                // Skip over tag and upvars
+                assignments[local] = Ineligible(Some((upvar_count + 1 + idx) as u32));
             }
         }
         debug!("generator saved local assignments: {:?}", assignments);
@@ -1535,12 +1547,15 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
             None => return Err(LayoutError::Unknown(ty)),
             Some(info) => info,
         };
-        let (ineligible_locals, assignments) = self.generator_saved_local_eligibility(&info);
+        let upvar_count = substs.as_generator().prefix_tys().count();
+
+        let (ineligible_locals, assignments) =
+            self.generator_saved_local_eligibility(&info, upvar_count);
 
         // Build a prefix layout, including "promoting" all ineligible
         // locals as part of the prefix. We compute the layout of all of
         // these fields at once to get optimal packing.
-        let tag_index = substs.as_generator().prefix_tys().count();
+        let tag_index = upvar_count;
 
         // `info.variant_fields` already accounts for the reserved variants, so no need to add them.
         let max_discr = (info.variant_fields.len() - 1) as u128;
@@ -1558,6 +1573,12 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
             .map(|local| subst_field(info.field_tys[local]))
             .map(|ty| tcx.mk_maybe_uninit(ty))
             .map(|ty| self.layout_of(ty));
+
+        // The upvars are stored at the beginning of the generators
+        // They are also included in the 0th (unresumed) variant
+        // The promoted locals are placed directly after the upvars.
+        // Because of this the rest of the code can handle upvars and promoted locals in a generic
+        // way.
         let prefix_layouts = substs
             .as_generator()
             .prefix_tys()
@@ -1580,35 +1601,57 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
         // GeneratorLayout.
         debug!("prefix = {:#?}", prefix);
         let (outer_fields, promoted_offsets, promoted_memory_index) = match prefix.fields {
-            FieldsShape::Arbitrary { mut offsets, memory_index } => {
-                let mut inverse_memory_index = invert_mapping(&memory_index);
-
-                // "a" (`0..b_start`) and "b" (`b_start..`) correspond to
-                // "outer" and "promoted" fields respectively.
-                let b_start = (tag_index + 1) as u32;
-                let offsets_b = offsets.split_off(b_start as usize);
-                let offsets_a = offsets;
-
-                // Disentangle the "a" and "b" components of `inverse_memory_index`
-                // by preserving the order but keeping only one disjoint "half" each.
-                // FIXME(eddyb) build a better abstraction for permutations, if possible.
-                let inverse_memory_index_b: Vec<_> =
-                    inverse_memory_index.iter().filter_map(|&i| i.checked_sub(b_start)).collect();
-                inverse_memory_index.retain(|&i| i < b_start);
-                let inverse_memory_index_a = inverse_memory_index;
-
-                // Since `inverse_memory_index_{a,b}` each only refer to their
+            FieldsShape::Arbitrary { offsets, memory_index } => {
+                let inverse_memory_index = invert_mapping(&memory_index);
+
+                // Outer fields - upvars and tag
+                let after_tag = tag_index + 1;
+                let offsets_outer: Vec<_> = offsets.iter().take(after_tag).cloned().collect();
+                let mut inverse_memory_index_outer = inverse_memory_index.clone();
+                inverse_memory_index_outer.retain(|&i| i < after_tag as u32);
+
+                // Promoted fields - upvars and promoted locals
+                let offsets_promoted = offsets;
+                let inverse_memory_index_promoted = inverse_memory_index;
+                /*let mut offsets_promoted = offsets;
+                offsets_promoted.remove(tag_index);
+                // Keep indices lower than tag, remove tag, and lower indices higher than tag
+                let tag_source_index = inverse_memory_index[tag_index];
+                let inverse_memory_index_promoted: Vec<_> = inverse_memory_index
+                    .into_iter()
+                    .filter_map(|source_index| {
+                        if source_index != tag_source_index {
+                            if source_index >= tag_source_index {
+                                Some(source_index - 1)
+                            } else {
+                                Some(source_index)
+                            }
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();*/
+
+                // Since `inverse_memory_index_{outer,promoted}` each only refer to their
                 // respective fields, they can be safely inverted
-                let memory_index_a = invert_mapping(&inverse_memory_index_a);
-                let memory_index_b = invert_mapping(&inverse_memory_index_b);
+                let memory_index_outer = invert_mapping(&inverse_memory_index_outer);
+                let memory_index_promoted = invert_mapping(&inverse_memory_index_promoted);
 
-                let outer_fields =
-                    FieldsShape::Arbitrary { offsets: offsets_a, memory_index: memory_index_a };
-                (outer_fields, offsets_b, memory_index_b)
+                let outer_fields = FieldsShape::Arbitrary {
+                    offsets: offsets_outer,
+                    memory_index: memory_index_outer,
+                };
+                (outer_fields, offsets_promoted, memory_index_promoted)
             }
             _ => bug!(),
         };
 
+        for (index, layout) in prefix_layouts.iter().enumerate() {
+            info!("LAYOUT, prefix_layout_{}={:#?}", index, layout);
+        }
+        info!("LAYOUT, tag_index={:?}, upvar_count={:?}, prefix_layouts.len()={:?}, prefix_size={:?}", tag_index, upvar_count, prefix_layouts.len(), prefix_size);
+        info!("LAYOUT, outer_fields={:#?}, promoted_offsets={:#?}, promoted_memory_index={:#?}", outer_fields, promoted_offsets, promoted_memory_index);
+
         let mut size = prefix.size;
         let mut align = prefix.align;
         let variants = info
@@ -1626,11 +1669,14 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                     })
                     .map(|local| subst_field(info.field_tys[*local]));
 
+                let x = variant_only_tys
+                    .map(|ty| self.layout_of(ty))
+                    .collect::<Result<Vec<_>, _>>()?;
+                info!("VARIANT, index={:?}, variant_only_tys={:#?}", index, x);
+
                 let mut variant = self.univariant_uninterned(
                     ty,
-                    &variant_only_tys
-                        .map(|ty| self.layout_of(ty))
-                        .collect::<Result<Vec<_>, _>>()?,
+                    &x,
                     &ReprOptions::default(),
                     StructKind::Prefixed(prefix_size, prefix_align.abi),
                 )?;
@@ -1713,7 +1759,7 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
             size,
             align,
         });
-        debug!("generator layout ({:?}): {:#?}", ty, layout);
+        info!("generator layout ({:?}): {:#?}", ty, layout);
         Ok(layout)
     }
 

compiler/rustc_mir_build/src/build/expr/as_rvalue.rs

@@ -256,7 +256,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let result = match substs {
                     UpvarSubsts::Generator(substs) => {
                         // We implicitly set the discriminant to 0. See
-                        // librustc_mir/transform/deaggregator.rs for details.
+                        // rustc_mir_transform/src/deaggregator.rs for details.
                         let movability = movability.unwrap();
                         Box::new(AggregateKind::Generator(closure_id, substs, movability))
                     }

compiler/rustc_mir_transform/src/generator.rs

@@ -425,6 +425,20 @@ fn replace_local<'tcx>(
     new_local
 }
 
+/// Upvars are stored in locals directly after the return place and generator arguments.
+/// _0 = return place
+/// _1 = generator type
+/// _2 = resume type
+/// _3..._n = upvars
+/// TODO: where is this guaranteed?
+fn upvar_locals<'tcx>(upvars: &[Ty<'tcx>]) -> impl Iterator<Item = Local> {
+    (0..upvars.len()).map(|index| Local::new(3 + index))
+}
+
+fn is_upvar_local<'tcx>(upvars: &[Ty<'tcx>], local: Local) -> bool {
+    local >= Local::new(3) && local < Local::new(3 + upvars.len())
+}
+
 struct LivenessInfo {
     /// Which locals are live across any suspension point.
     saved_locals: GeneratorSavedLocals,
@@ -449,6 +463,7 @@ fn locals_live_across_suspend_points(
     tcx: TyCtxt<'tcx>,
     body: &Body<'tcx>,
     always_live_locals: &storage::AlwaysLiveLocals,
+    upvars: &[Ty<'tcx>],
     movable: bool,
 ) -> LivenessInfo {
     let body_ref: &Body<'_> = &body;
@@ -490,6 +505,11 @@ fn locals_live_across_suspend_points(
     let mut source_info_at_suspension_points = Vec::new();
     let mut live_locals_at_any_suspension_point = BitSet::new_empty(body.local_decls.len());
 
+    // Upvars are always stored in the generator
+    for local in upvar_locals(upvars) {
+        live_locals_at_any_suspension_point.insert(local);
+    }
+
     for (block, data) in body.basic_blocks().iter_enumerated() {
         if let TerminatorKind::Yield { .. } = data.terminator().kind {
             let loc = Location { block, statement_index: data.statements.len() };
@@ -762,6 +782,7 @@ fn sanitize_witness<'tcx>(
 fn compute_layout<'tcx>(
     liveness: LivenessInfo,
     body: &mut Body<'tcx>,
+    upvars: &[Ty<'tcx>],
 ) -> (
     FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
     GeneratorLayout<'tcx>,
@@ -775,34 +796,44 @@ fn compute_layout<'tcx>(
         storage_liveness,
     } = liveness;
 
-    // Gather live local types and their indices.
-    let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
-    let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
-    for (saved_local, local) in saved_locals.iter_enumerated() {
-        locals.push(local);
-        tys.push(body.local_decls[local].ty);
-        debug!("generator saved local {:?} => {:?}", saved_local, local);
-    }
-
-    // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states.
+    // Leave variants for the UNRESUMED, RETURNED, and POISONED states.
     // In debuginfo, these will correspond to the beginning (UNRESUMED) or end
     // (RETURNED, POISONED) of the function.
     const RESERVED_VARIANTS: usize = 3;
     let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
-    let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = [
+    let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = std::array::IntoIter::new([
         SourceInfo::outermost(body_span.shrink_to_lo()),
         SourceInfo::outermost(body_span.shrink_to_hi()),
         SourceInfo::outermost(body_span.shrink_to_hi()),
-    ]
-    .iter()
-    .copied()
+    ])
     .collect();
 
     // Build the generator variant field list.
     // Create a map from local indices to generator struct indices.
     let mut variant_fields: IndexVec<VariantIdx, IndexVec<Field, GeneratorSavedLocal>> =
         iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect();
+
+    // Gather live local types and their indices.
+    // Store upvars in the UNRESUMED variant.
+    let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
+    let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
+    let unresumed_idx = VariantIdx::new(UNRESUMED);
     let mut remap = FxHashMap::default();
+
+    for (saved_local, local) in saved_locals.iter_enumerated() {
+        locals.push(local);
+        tys.push(body.local_decls[local].ty);
+
+        if is_upvar_local(upvars, local) {
+            variant_fields[unresumed_idx].push(saved_local);
+            remap.entry(locals[saved_local]).or_insert((
+                tys[saved_local],
+                unresumed_idx,
+                local.as_usize() - 3,
+            ));
+        }
+    }
+
     for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() {
         let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx);
         let mut fields = IndexVec::new();
@@ -1249,7 +1280,7 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
             ty::Generator(_, substs, movability) => {
                 let substs = substs.as_generator();
                 (
-                    substs.upvar_tys().collect(),
+                    substs.upvar_tys().collect::<Vec<_>>(),
                     substs.witness(),
                     substs.discr_ty(tcx),
                     movability == hir::Movability::Movable,
@@ -1297,9 +1328,9 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         let always_live_locals = storage::AlwaysLiveLocals::new(&body);
 
         let liveness_info =
-            locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
+            locals_live_across_suspend_points(tcx, body, &always_live_locals, &upvars, movable);
 
-        sanitize_witness(tcx, body, interior, upvars, &liveness_info.saved_locals);
+        sanitize_witness(tcx, body, interior, upvars.clone(), &liveness_info.saved_locals);
 
         if tcx.sess.opts.debugging_opts.validate_mir {
             let mut vis = EnsureGeneratorFieldAssignmentsNeverAlias {
@@ -1314,7 +1345,7 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         // Extract locals which are live across suspension point into `layout`
         // `remap` gives a mapping from local indices onto generator struct indices
         // `storage_liveness` tells us which locals have live storage at suspension points
-        let (remap, layout, storage_liveness) = compute_layout(liveness_info, body);
+        let (remap, layout, storage_liveness) = compute_layout(liveness_info, body, &upvars);
 
         let can_return = can_return(tcx, body, tcx.param_env(body.source.def_id()));
 

src/test/ui/async-await/async-fn-size-parameters.rs

@@ -0,0 +1,35 @@
+// run-pass
+// edition:2018
+
+async fn parameter_no_drop_no_use(_a: u32) {
+    std::future::ready(()).await;
+}
+
+async fn parameter_explicit_drop(a: u32) {
+    std::future::ready(()).await;
+    drop(a);
+}
+
+struct HasDrop(u32);
+
+impl Drop for HasDrop {
+    fn drop(&mut self) {
+        println!("drop");
+    }
+}
+
+async fn parameter_implicit_drop(_a: HasDrop) {
+    std::future::ready(()).await;
+}
+
+async fn parameter_rebind(a: HasDrop) {
+    let _x = a;
+    std::future::ready(()).await;
+}
+
+fn main() {
+    assert_eq!(8, std::mem::size_of_val(&parameter_no_drop_no_use(0)));
+    assert_eq!(8, std::mem::size_of_val(&parameter_explicit_drop(0)));
+    assert_eq!(8, std::mem::size_of_val(&parameter_implicit_drop(HasDrop(0))));
+    assert_eq!(16, std::mem::size_of_val(&parameter_rebind(HasDrop(0))));
+}