|
| 1 | +use std::cmp::Reverse; |
| 2 | +use std::collections::BTreeSet; |
| 3 | +use std::fmt::Debug; |
| 4 | +use std::ops::Bound; |
| 5 | + |
| 6 | +use rustc_abi::{ |
| 7 | + BackendRepr, FieldIdx, FieldsShape, Integer, Layout, LayoutData, Primitive, Scalar, Size, |
| 8 | + TagEncoding, TyAndLayout, Variants, WrappingRange, |
| 9 | +}; |
| 10 | +use rustc_index::bit_set::DenseBitSet; |
| 11 | +use rustc_index::{Idx, IndexVec}; |
| 12 | +use rustc_middle::mir::CoroutineSavedLocal; |
| 13 | +use rustc_middle::ty::layout::{HasTyCtxt, IntegerExt, LayoutCx, LayoutError, LayoutOf}; |
| 14 | +use rustc_middle::ty::{EarlyBinder, GenericArgsRef, Ty}; |
| 15 | +use tracing::{debug, instrument}; |
| 16 | + |
| 17 | +use super::error; |
| 18 | + |
| 19 | +#[instrument(level = "debug", skip(cx))] |
| 20 | +pub(super) fn coroutine_layout<'tcx>( |
| 21 | + cx: &LayoutCx<'tcx>, |
| 22 | + ty: Ty<'tcx>, |
| 23 | + def_id: rustc_hir::def_id::DefId, |
| 24 | + args: GenericArgsRef<'tcx>, |
| 25 | +) -> Result<Layout<'tcx>, &'tcx LayoutError<'tcx>> { |
| 26 | + let tcx = cx.tcx(); |
| 27 | + let Some(info) = tcx.coroutine_layout(def_id, args.as_coroutine().kind_ty()) else { |
| 28 | + return Err(error(cx, LayoutError::Unknown(ty))); |
| 29 | + }; |
| 30 | + let mut relocated_upvars = info.relocated_upvars.clone(); |
| 31 | + for (&source, &dest) in &info.relocated_upvars { |
| 32 | + relocated_upvars.insert(dest, source); |
| 33 | + } |
| 34 | + |
| 35 | + let tcx = cx.tcx(); |
| 36 | + let field_layouts: IndexVec<CoroutineSavedLocal, _> = info |
| 37 | + .field_tys |
| 38 | + .iter_enumerated() |
| 39 | + .map(|(saved_local, ty)| { |
| 40 | + let ty = EarlyBinder::bind(ty.ty).instantiate(tcx, args); |
| 41 | + cx.spanned_layout_of(ty, info.field_tys[saved_local].source_info.span) |
| 42 | + }) |
| 43 | + .try_collect()?; |
| 44 | + let layouts: IndexVec<CoroutineSavedLocal, _> = |
| 45 | + field_layouts.iter().map(|data| data.layout.clone()).collect(); |
| 46 | + |
| 47 | + let field_sort_keys: IndexVec<CoroutineSavedLocal, _>; |
| 48 | + let mut saved_locals: Vec<_>; |
| 49 | + // ## The heuristic on which saved locals get allocation first ## |
| 50 | + // 1. the alignment |
| 51 | + // Intuitively data with a larger alignment asks for a larger block of contiguous memory. |
| 52 | + // It is easier to get large blocks early in the beginning, but it will get harder to |
| 53 | + // recover them as fragmentation creeps in when data with smaller alignment occupies |
| 54 | + // the large chunks. |
| 55 | + // 2. the size |
| 56 | + // The size also poses restriction on layout, but not as potent as alignment. |
| 57 | + // 3. the degree of conflicts |
| 58 | + // This metric is the number of confliciting saved locals with a given saved local. |
| 59 | + // Preferring allocating highly conflicting data over those that are less and more |
| 60 | + // transient in nature will keep the fragmentation contained in neighbourhoods of a layout. |
| 61 | + (saved_locals, field_sort_keys) = field_layouts |
| 62 | + .iter_enumerated() |
| 63 | + .map(|(saved_local, ty)| { |
| 64 | + ( |
| 65 | + saved_local, |
| 66 | + ( |
| 67 | + Reverse(ty.align.abi), |
| 68 | + Reverse(ty.size), |
| 69 | + Reverse(info.storage_conflicts.count(saved_local)), |
| 70 | + ), |
| 71 | + ) |
| 72 | + }) |
| 73 | + .unzip(); |
| 74 | + let mut uninhabited_or_zst = DenseBitSet::new_empty(saved_locals.len()); |
| 75 | + for (saved_local, ty) in field_layouts.iter_enumerated() { |
| 76 | + if ty.backend_repr.is_uninhabited() || ty.is_zst() { |
| 77 | + uninhabited_or_zst.insert(saved_local); |
| 78 | + } |
| 79 | + } |
| 80 | + saved_locals.sort_by_key(|&idx| &field_sort_keys[idx]); |
| 81 | + let max_discr = (info.variant_fields.len() - 1) as u128; |
| 82 | + let discr_int = Integer::fit_unsigned(max_discr); |
| 83 | + let tag = Scalar::Initialized { |
| 84 | + value: Primitive::Int(discr_int, false), |
| 85 | + valid_range: WrappingRange { start: 0, end: max_discr }, |
| 86 | + }; |
| 87 | + let tag_layout = TyAndLayout { |
| 88 | + ty: discr_int.to_ty(tcx, false), |
| 89 | + layout: tcx.mk_layout(LayoutData::scalar(cx, tag)), |
| 90 | + }; |
| 91 | + // This will be *the* align of the entire coroutine, |
| 92 | + // which is the maximal alignment of all saved locals. |
| 93 | + // We need to also consider the tag layout alignment. |
| 94 | + let align = saved_locals |
| 95 | + .get(0) |
| 96 | + .map(|&idx| layouts[idx].align.max(tag_layout.align)) |
| 97 | + .unwrap_or(tag_layout.align); |
| 98 | + |
| 99 | + // ## The blocked map, or the reservation map ## |
| 100 | + // This map from saved locals to memory layout records the reservation |
| 101 | + // status of the coroutine state memory, down to the byte granularity. |
| 102 | + // `Slot`s are inserted to mark ranges of memory that a particular saved local |
| 103 | + // shall not have overlapping memory allocation, due to the liveness of |
| 104 | + // other conflicting saved locals. |
| 105 | + // Therefore, we can try to make reservation for this saved local |
| 106 | + // by inspecting the gaps before, between, and after those blocked-out memory ranges. |
| 107 | + let mut blocked: IndexVec<CoroutineSavedLocal, BTreeSet<Slot>> = |
| 108 | + IndexVec::from_elem_n(BTreeSet::new(), saved_locals.len()); |
| 109 | + let mut has_assignment = DenseBitSet::new_empty(saved_locals.len()); |
| 110 | + let mut tag_blocked = BTreeSet::new(); |
| 111 | + let mut assignment = IndexVec::from_elem_n(Slot { start: 0, end: 0 }, saved_locals.len()); |
| 112 | + for (idx, ¤t_local) in saved_locals.iter().enumerate() { |
| 113 | + if uninhabited_or_zst.contains(current_local) { |
| 114 | + // Do not bother to compute on uninhabited data. |
| 115 | + // They will not get allocation after all. |
| 116 | + // By default, a ZST occupies the beginning of the coroutine state. |
| 117 | + continue; |
| 118 | + } |
| 119 | + let layout_data = &field_layouts[current_local]; |
| 120 | + |
| 121 | + let candidate = if let Some(&another_local) = relocated_upvars.get(¤t_local) |
| 122 | + && has_assignment.contains(another_local) |
| 123 | + { |
| 124 | + assignment[another_local] |
| 125 | + } else { |
| 126 | + find_lowest_viable_allocation(&layout_data, &blocked[current_local]) |
| 127 | + }; |
| 128 | + // The discriminant is certainly conflicting with all the saved locals |
| 129 | + merge_slot_in(&mut tag_blocked, candidate); |
| 130 | + for &other_local in &saved_locals[idx + 1..] { |
| 131 | + if info.storage_conflicts.contains(current_local, other_local) { |
| 132 | + merge_slot_in(&mut blocked[other_local], candidate); |
| 133 | + } |
| 134 | + } |
| 135 | + // Adjustment to the layout of this field by shifting them into the chosen slot |
| 136 | + assignment[current_local] = candidate; |
| 137 | + has_assignment.insert(current_local); |
| 138 | + } |
| 139 | + debug!(assignment = ?assignment.debug_map_view()); |
| 140 | + |
| 141 | + // Find a slot for discriminant, also known as the tag. |
| 142 | + let tag_candidate = find_lowest_viable_allocation(&tag_layout, &tag_blocked); |
| 143 | + debug!(tag = ?tag_candidate); |
| 144 | + |
| 145 | + // Assemble the layout for each coroutine state |
| 146 | + let variants: IndexVec<_, LayoutData<_, _>> = info |
| 147 | + .variant_fields |
| 148 | + .iter_enumerated() |
| 149 | + .map(|(index, fields)| { |
| 150 | + let size = Size::from_bytes( |
| 151 | + fields |
| 152 | + .iter() |
| 153 | + .map(|&saved_local| assignment[saved_local].end) |
| 154 | + .max() |
| 155 | + .unwrap_or(0) |
| 156 | + .max(tag_candidate.end), |
| 157 | + ) |
| 158 | + .align_to(align.abi); |
| 159 | + let offsets: IndexVec<_, _> = fields |
| 160 | + .iter() |
| 161 | + .map(|&saved_local| Size::from_bytes(assignment[saved_local].start)) |
| 162 | + .collect(); |
| 163 | + let memory_index = IndexVec::from_fn_n(|n: FieldIdx| (n.index() as u32), offsets.len()); |
| 164 | + LayoutData { |
| 165 | + // We are aware of specialized layouts such as scalar pairs but this is still |
| 166 | + // in development. |
| 167 | + // Let us hold off from further optimisation until more information is available. |
| 168 | + fields: FieldsShape::Arbitrary { offsets, memory_index }, |
| 169 | + variants: Variants::Single { index }, |
| 170 | + backend_repr: BackendRepr::Memory { sized: true }, |
| 171 | + largest_niche: None, |
| 172 | + align, |
| 173 | + size, |
| 174 | + max_repr_align: None, |
| 175 | + unadjusted_abi_align: align.abi, |
| 176 | + randomization_seed: 0, |
| 177 | + } |
| 178 | + }) |
| 179 | + .collect(); |
| 180 | + let size = variants |
| 181 | + .iter() |
| 182 | + .map(|data| data.size) |
| 183 | + .max() |
| 184 | + .unwrap_or(Size::ZERO) |
| 185 | + .max(Size::from_bytes(tag_candidate.end)) |
| 186 | + .align_to(align.abi); |
| 187 | + let layout = tcx.mk_layout(LayoutData { |
| 188 | + fields: FieldsShape::Arbitrary { |
| 189 | + offsets: [Size::from_bytes(tag_candidate.start)].into(), |
| 190 | + memory_index: [0].into(), |
| 191 | + }, |
| 192 | + variants: Variants::Multiple { |
| 193 | + tag, |
| 194 | + tag_encoding: TagEncoding::Direct, |
| 195 | + tag_field: 0, |
| 196 | + variants, |
| 197 | + }, |
| 198 | + backend_repr: BackendRepr::Memory { sized: true }, |
| 199 | + // Suppress niches inside coroutines. If the niche is inside a field that is aliased (due to |
| 200 | + // self-referentiality), getting the discriminant can cause aliasing violations. |
| 201 | + // `UnsafeCell` blocks niches for the same reason, but we don't yet have `UnsafePinned` that |
| 202 | + // would do the same for us here. |
| 203 | + // See <https://github.com/rust-lang/rust/issues/63818>, <https://github.com/rust-lang/miri/issues/3780>. |
| 204 | + // FIXME(#125735): Remove when <https://github.com/rust-lang/rust/issues/125735> |
| 205 | + // is implemented and aliased coroutine fields are wrapped in `UnsafePinned`. |
| 206 | + // NOTE(@dingxiangfei2009): I believe there is still niche, which is the tag, |
| 207 | + // but I am not sure how much benefit is there for us to grab. |
| 208 | + largest_niche: None, |
| 209 | + align, |
| 210 | + size, |
| 211 | + max_repr_align: None, |
| 212 | + unadjusted_abi_align: align.abi, |
| 213 | + randomization_seed: 0, |
| 214 | + }); |
| 215 | + debug!("coroutine layout ({:?}): {:#?}", ty, layout); |
| 216 | + Ok(layout) |
| 217 | +} |
| 218 | + |
| 219 | +/// An occupied slot in the coroutine memory at some yield point |
| 220 | +#[derive(PartialEq, Eq, Copy, Clone)] |
| 221 | +struct Slot { |
| 222 | + /// Beginning of the memory slot, inclusive |
| 223 | + start: u64, |
| 224 | + /// End of the memory slot, exclusive or one byte past the data |
| 225 | + end: u64, |
| 226 | +} |
| 227 | + |
| 228 | +impl PartialOrd for Slot { |
| 229 | + fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { |
| 230 | + (self.start, self.end).partial_cmp(&(other.start, other.end)) |
| 231 | + } |
| 232 | +} |
| 233 | + |
| 234 | +impl Ord for Slot { |
| 235 | + fn cmp(&self, other: &Self) -> std::cmp::Ordering { |
| 236 | + (self.start, self.end).cmp(&(other.start, other.end)) |
| 237 | + } |
| 238 | +} |
| 239 | + |
| 240 | +impl Debug for Slot { |
| 241 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 242 | + f.debug_tuple("Slot").field(&self.start).field(&self.end).finish() |
| 243 | + } |
| 244 | +} |
| 245 | + |
| 246 | +impl Slot { |
| 247 | + fn overlap_with(&self, other: &Self) -> bool { |
| 248 | + if self.start == self.end || other.start == other.end { |
| 249 | + return false; |
| 250 | + } |
| 251 | + self.contains_point(other.start) || other.contains_point(self.start) |
| 252 | + } |
| 253 | + fn contains_point(&self, point: u64) -> bool { |
| 254 | + self.start <= point && point < self.end |
| 255 | + } |
| 256 | +} |
| 257 | + |
| 258 | +fn merge_slot_in(slots: &mut BTreeSet<Slot>, slot: Slot) { |
| 259 | + let start = Slot { start: slot.start, end: slot.start }; |
| 260 | + let end = Slot { start: slot.end, end: slot.end }; |
| 261 | + let one_past_end = Slot { start: slot.end + 1, end: slot.end + 1 }; |
| 262 | + let (range_start, replace_start) = if let Some(prev) = slots.range(..start).next_back() |
| 263 | + && (prev.end == slot.start || prev.contains_point(slot.start)) |
| 264 | + { |
| 265 | + (Bound::Included(prev), prev.start) |
| 266 | + } else { |
| 267 | + (Bound::Included(&start), slot.start) |
| 268 | + }; |
| 269 | + let (range_end, replace_end) = if let Some(next) = slots.range(..one_past_end).next_back() |
| 270 | + && next.start == slot.end |
| 271 | + { |
| 272 | + (Bound::Included(next), next.end) |
| 273 | + } else if let Some(prev) = slots.range(..end).next_back() |
| 274 | + && prev.contains_point(slot.end) |
| 275 | + { |
| 276 | + (Bound::Included(prev), prev.end) |
| 277 | + } else { |
| 278 | + (Bound::Included(&end), slot.end) |
| 279 | + }; |
| 280 | + let to_remove: Vec<_> = slots.range((range_start, range_end)).copied().collect(); |
| 281 | + for slot in to_remove { |
| 282 | + slots.remove(&slot); |
| 283 | + } |
| 284 | + slots.insert(Slot { start: replace_start, end: replace_end }); |
| 285 | +} |
| 286 | + |
| 287 | +fn find_lowest_viable_allocation<F: Idx, V: Idx>( |
| 288 | + layout: &LayoutData<F, V>, |
| 289 | + blocked: &BTreeSet<Slot>, |
| 290 | +) -> Slot { |
| 291 | + let size = layout.size.bytes(); |
| 292 | + let align = layout.align.abi; |
| 293 | + let mut candidate = Slot { start: 0, end: size }; |
| 294 | + for slot in blocked { |
| 295 | + if slot.overlap_with(&candidate) { |
| 296 | + let start = Size::from_bytes(slot.end).align_to(align).bytes(); |
| 297 | + candidate = Slot { start, end: start + size }; |
| 298 | + } else { |
| 299 | + break; |
| 300 | + } |
| 301 | + } |
| 302 | + candidate |
| 303 | +} |
0 commit comments