1use std::ops::ControlFlow;
2
3use super::{Byte, Def, Ref};
4
5#[cfg(test)]
6mod tests;
7
8#[derive(Clone, Debug, Hash, PartialEq, Eq)]
18pub(crate) enum Tree<D, R>
19where
20 D: Def,
21 R: Ref,
22{
23 Seq(Vec<Self>),
25 Alt(Vec<Self>),
27 Def(D),
29 Ref(R),
31 Byte(Byte),
33}
34
35impl<D, R> Tree<D, R>
36where
37 D: Def,
38 R: Ref,
39{
40 pub(crate) fn def(def: D) -> Self {
42 Self::Def(def)
43 }
44
45 pub(crate) fn uninhabited() -> Self {
47 Self::Alt(vec![])
48 }
49
50 pub(crate) fn unit() -> Self {
52 Self::Seq(Vec::new())
53 }
54
55 pub(crate) fn uninit() -> Self {
57 Self::Byte(Byte::Uninit)
58 }
59
60 pub(crate) fn bool() -> Self {
62 Self::from_bits(0x00).or(Self::from_bits(0x01))
63 }
64
65 pub(crate) fn u8() -> Self {
67 Self::Alt((0u8..=255).map(Self::from_bits).collect())
68 }
69
70 pub(crate) fn from_bits(bits: u8) -> Self {
72 Self::Byte(Byte::Init(bits))
73 }
74
75 pub(crate) fn number(width_in_bytes: usize) -> Self {
77 Self::Seq(vec![Self::u8(); width_in_bytes])
78 }
79
80 pub(crate) fn padding(width_in_bytes: usize) -> Self {
82 Self::Seq(vec![Self::uninit(); width_in_bytes])
83 }
84
85 pub(crate) fn prune<F>(self, f: &F) -> Tree<!, R>
88 where
89 F: Fn(D) -> bool,
90 {
91 match self {
92 Self::Seq(elts) => match elts.into_iter().map(|elt| elt.prune(f)).try_fold(
93 Tree::unit(),
94 |elts, elt| {
95 if elt == Tree::uninhabited() {
96 ControlFlow::Break(Tree::uninhabited())
97 } else {
98 ControlFlow::Continue(elts.then(elt))
99 }
100 },
101 ) {
102 ControlFlow::Break(node) | ControlFlow::Continue(node) => node,
103 },
104 Self::Alt(alts) => alts
105 .into_iter()
106 .map(|alt| alt.prune(f))
107 .fold(Tree::uninhabited(), |alts, alt| alts.or(alt)),
108 Self::Byte(b) => Tree::Byte(b),
109 Self::Ref(r) => Tree::Ref(r),
110 Self::Def(d) => {
111 if f(d) {
112 Tree::uninhabited()
113 } else {
114 Tree::unit()
115 }
116 }
117 }
118 }
119
120 pub(crate) fn is_inhabited(&self) -> bool {
122 match self {
123 Self::Seq(elts) => elts.into_iter().all(|elt| elt.is_inhabited()),
124 Self::Alt(alts) => alts.into_iter().any(|alt| alt.is_inhabited()),
125 Self::Byte(..) | Self::Ref(..) | Self::Def(..) => true,
126 }
127 }
128}
129
130impl<D, R> Tree<D, R>
131where
132 D: Def,
133 R: Ref,
134{
135 pub(crate) fn then(self, other: Self) -> Self {
137 match (self, other) {
138 (Self::Seq(elts), other) | (other, Self::Seq(elts)) if elts.len() == 0 => other,
139 (Self::Seq(mut lhs), Self::Seq(mut rhs)) => {
140 lhs.append(&mut rhs);
141 Self::Seq(lhs)
142 }
143 (Self::Seq(mut lhs), rhs) => {
144 lhs.push(rhs);
145 Self::Seq(lhs)
146 }
147 (lhs, Self::Seq(mut rhs)) => {
148 rhs.insert(0, lhs);
149 Self::Seq(rhs)
150 }
151 (lhs, rhs) => Self::Seq(vec![lhs, rhs]),
152 }
153 }
154
155 pub(crate) fn or(self, other: Self) -> Self {
157 match (self, other) {
158 (Self::Alt(alts), other) | (other, Self::Alt(alts)) if alts.len() == 0 => other,
159 (Self::Alt(mut lhs), Self::Alt(rhs)) => {
160 lhs.extend(rhs);
161 Self::Alt(lhs)
162 }
163 (Self::Alt(mut alts), alt) | (alt, Self::Alt(mut alts)) => {
164 alts.push(alt);
165 Self::Alt(alts)
166 }
167 (lhs, rhs) => Self::Alt(vec![lhs, rhs]),
168 }
169 }
170}
171
172#[cfg(feature = "rustc")]
173pub(crate) mod rustc {
174 use rustc_abi::{
175 FieldIdx, FieldsShape, Layout, Size, TagEncoding, TyAndLayout, VariantIdx, Variants,
176 };
177 use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, LayoutError};
178 use rustc_middle::ty::{self, AdtDef, AdtKind, List, ScalarInt, Ty, TyCtxt, TypeVisitableExt};
179 use rustc_span::ErrorGuaranteed;
180
181 use super::Tree;
182 use crate::layout::rustc::{Def, Ref, layout_of};
183
184 #[derive(Debug, Copy, Clone)]
185 pub(crate) enum Err {
186 NotYetSupported,
188 UnknownLayout,
190 SizeOverflow,
192 TypeError(ErrorGuaranteed),
193 }
194
195 impl<'tcx> From<&LayoutError<'tcx>> for Err {
196 fn from(err: &LayoutError<'tcx>) -> Self {
197 match err {
198 LayoutError::Unknown(..)
199 | LayoutError::ReferencesError(..)
200 | LayoutError::TooGeneric(..)
201 | LayoutError::NormalizationFailure(..) => Self::UnknownLayout,
202 LayoutError::SizeOverflow(..) => Self::SizeOverflow,
203 LayoutError::Cycle(err) => Self::TypeError(*err),
204 }
205 }
206 }
207
208 impl<'tcx> Tree<Def<'tcx>, Ref<'tcx>> {
209 pub(crate) fn from_ty(ty: Ty<'tcx>, cx: LayoutCx<'tcx>) -> Result<Self, Err> {
210 use rustc_abi::HasDataLayout;
211 let layout = layout_of(cx, ty)?;
212
213 if let Err(e) = ty.error_reported() {
214 return Err(Err::TypeError(e));
215 }
216
217 let target = cx.data_layout();
218 let pointer_size = target.pointer_size;
219
220 match ty.kind() {
221 ty::Bool => Ok(Self::bool()),
222
223 ty::Float(nty) => {
224 let width = nty.bit_width() / 8;
225 Ok(Self::number(width as _))
226 }
227
228 ty::Int(nty) => {
229 let width = nty.normalize(pointer_size.bits() as _).bit_width().unwrap() / 8;
230 Ok(Self::number(width as _))
231 }
232
233 ty::Uint(nty) => {
234 let width = nty.normalize(pointer_size.bits() as _).bit_width().unwrap() / 8;
235 Ok(Self::number(width as _))
236 }
237
238 ty::Tuple(members) => Self::from_tuple((ty, layout), members, cx),
239
240 ty::Array(inner_ty, _len) => {
241 let FieldsShape::Array { stride, count } = &layout.fields else {
242 return Err(Err::NotYetSupported);
243 };
244 let inner_layout = layout_of(cx, *inner_ty)?;
245 assert_eq!(*stride, inner_layout.size);
246 let elt = Tree::from_ty(*inner_ty, cx)?;
247 Ok(std::iter::repeat(elt)
248 .take(*count as usize)
249 .fold(Tree::unit(), |tree, elt| tree.then(elt)))
250 }
251
252 ty::Adt(adt_def, _args_ref) if !ty.is_box() => match adt_def.adt_kind() {
253 AdtKind::Struct => Self::from_struct((ty, layout), *adt_def, cx),
254 AdtKind::Enum => Self::from_enum((ty, layout), *adt_def, cx),
255 AdtKind::Union => Self::from_union((ty, layout), *adt_def, cx),
256 },
257
258 ty::Ref(lifetime, ty, mutability) => {
259 let layout = layout_of(cx, *ty)?;
260 let align = layout.align.abi.bytes_usize();
261 let size = layout.size.bytes_usize();
262 Ok(Tree::Ref(Ref {
263 lifetime: *lifetime,
264 ty: *ty,
265 mutability: *mutability,
266 align,
267 size,
268 }))
269 }
270
271 _ => Err(Err::NotYetSupported),
272 }
273 }
274
275 fn from_tuple(
277 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
278 members: &'tcx List<Ty<'tcx>>,
279 cx: LayoutCx<'tcx>,
280 ) -> Result<Self, Err> {
281 match &layout.fields {
282 FieldsShape::Primitive => {
283 assert_eq!(members.len(), 1);
284 let inner_ty = members[0];
285 Self::from_ty(inner_ty, cx)
286 }
287 FieldsShape::Arbitrary { offsets, .. } => {
288 assert_eq!(offsets.len(), members.len());
289 Self::from_variant(Def::Primitive, None, (ty, layout), layout.size, cx)
290 }
291 FieldsShape::Array { .. } | FieldsShape::Union(_) => Err(Err::NotYetSupported),
292 }
293 }
294
295 fn from_struct(
301 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
302 def: AdtDef<'tcx>,
303 cx: LayoutCx<'tcx>,
304 ) -> Result<Self, Err> {
305 assert!(def.is_struct());
306 let def = Def::Adt(def);
307 Self::from_variant(def, None, (ty, layout), layout.size, cx)
308 }
309
310 fn from_enum(
316 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
317 def: AdtDef<'tcx>,
318 cx: LayoutCx<'tcx>,
319 ) -> Result<Self, Err> {
320 assert!(def.is_enum());
321
322 let layout_of_variant =
324 |index, encoding: Option<TagEncoding<VariantIdx>>| -> Result<Self, Err> {
325 let variant_layout = ty_variant(cx, (ty, layout), index);
326 if variant_layout.is_uninhabited() {
327 return Ok(Self::uninhabited());
328 }
329 let tag = cx.tcx().tag_for_variant((cx.tcx().erase_regions(ty), index));
330 let variant_def = Def::Variant(def.variant(index));
331 Self::from_variant(
332 variant_def,
333 tag.map(|tag| (tag, index, encoding.unwrap())),
334 (ty, variant_layout),
335 layout.size,
336 cx,
337 )
338 };
339
340 match layout.variants() {
341 Variants::Empty => Ok(Self::uninhabited()),
342 Variants::Single { index } => {
343 layout_of_variant(*index, None)
346 }
347 Variants::Multiple { tag: _, tag_encoding, tag_field, .. } => {
348 assert_eq!(*tag_field, 0);
355
356 let variants = def.discriminants(cx.tcx()).try_fold(
357 Self::uninhabited(),
358 |variants, (idx, _discriminant)| {
359 let variant = layout_of_variant(idx, Some(tag_encoding.clone()))?;
360 Result::<Self, Err>::Ok(variants.or(variant))
361 },
362 )?;
363
364 Ok(Self::def(Def::Adt(def)).then(variants))
365 }
366 }
367 }
368
369 fn from_variant(
379 def: Def<'tcx>,
380 tag: Option<(ScalarInt, VariantIdx, TagEncoding<VariantIdx>)>,
381 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
382 total_size: Size,
383 cx: LayoutCx<'tcx>,
384 ) -> Result<Self, Err> {
385 let FieldsShape::Arbitrary { offsets, memory_index } = layout.fields() else {
388 return Err(Err::NotYetSupported);
389 };
390
391 assert!(layout.size <= total_size);
395
396 let mut size = Size::ZERO;
397 let mut struct_tree = Self::def(def);
398
399 if let Some((tag, index, encoding)) = &tag {
401 match encoding {
402 TagEncoding::Direct => {
403 size += tag.size();
404 }
405 TagEncoding::Niche { niche_variants, .. } => {
406 if !niche_variants.contains(index) {
407 size += tag.size();
408 }
409 }
410 }
411 struct_tree = struct_tree.then(Self::from_tag(*tag, cx.tcx()));
412 }
413
414 let inverse_memory_index = memory_index.invert_bijective_mapping();
416 for &field_idx in inverse_memory_index.iter() {
417 let padding_needed = offsets[field_idx] - size;
419 let padding = Self::padding(padding_needed.bytes_usize());
420
421 let field_ty = ty_field(cx, (ty, layout), field_idx);
422 let field_layout = layout_of(cx, field_ty)?;
423 let field_tree = Self::from_ty(field_ty, cx)?;
424
425 struct_tree = struct_tree.then(padding).then(field_tree);
426
427 size += padding_needed + field_layout.size;
428 }
429
430 let padding_needed = total_size - size;
432 let trailing_padding = Self::padding(padding_needed.bytes_usize());
433
434 Ok(struct_tree.then(trailing_padding))
435 }
436
437 fn from_tag(tag: ScalarInt, tcx: TyCtxt<'tcx>) -> Self {
439 use rustc_abi::Endian;
440 let size = tag.size();
441 let bits = tag.to_bits(size);
442 let bytes: [u8; 16];
443 let bytes = match tcx.data_layout.endian {
444 Endian::Little => {
445 bytes = bits.to_le_bytes();
446 &bytes[..size.bytes_usize()]
447 }
448 Endian::Big => {
449 bytes = bits.to_be_bytes();
450 &bytes[bytes.len() - size.bytes_usize()..]
451 }
452 };
453 Self::Seq(bytes.iter().map(|&b| Self::from_bits(b)).collect())
454 }
455
456 fn from_union(
462 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
463 def: AdtDef<'tcx>,
464 cx: LayoutCx<'tcx>,
465 ) -> Result<Self, Err> {
466 assert!(def.is_union());
467
468 let FieldsShape::Union(_fields) = layout.fields() else {
471 return Err(Err::NotYetSupported);
472 };
473
474 let fields = &def.non_enum_variant().fields;
475 let fields = fields.iter_enumerated().try_fold(
476 Self::uninhabited(),
477 |fields, (idx, _field_def)| {
478 let field_ty = ty_field(cx, (ty, layout), idx);
479 let field_layout = layout_of(cx, field_ty)?;
480 let field = Self::from_ty(field_ty, cx)?;
481 let trailing_padding_needed = layout.size - field_layout.size;
482 let trailing_padding = Self::padding(trailing_padding_needed.bytes_usize());
483 let field_and_padding = field.then(trailing_padding);
484 Result::<Self, Err>::Ok(fields.or(field_and_padding))
485 },
486 )?;
487
488 Ok(Self::def(Def::Adt(def)).then(fields))
489 }
490 }
491
492 fn ty_field<'tcx>(
493 cx: LayoutCx<'tcx>,
494 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
495 i: FieldIdx,
496 ) -> Ty<'tcx> {
497 match ty.kind() {
502 ty::Adt(def, args) => {
503 match layout.variants {
504 Variants::Single { index } => {
505 let field = &def.variant(index).fields[i];
506 field.ty(cx.tcx(), args)
507 }
508 Variants::Empty => panic!("there is no field in Variants::Empty types"),
509 Variants::Multiple { tag, .. } => {
511 assert_eq!(i.as_usize(), 0);
512 ty::layout::PrimitiveExt::to_ty(&tag.primitive(), cx.tcx())
513 }
514 }
515 }
516 ty::Tuple(fields) => fields[i.as_usize()],
517 kind @ _ => unimplemented!(
518 "only a subset of `Ty::ty_and_layout_field`'s functionality is implemented. implementation needed for {:?}",
519 kind
520 ),
521 }
522 }
523
524 fn ty_variant<'tcx>(
525 cx: LayoutCx<'tcx>,
526 (ty, layout): (Ty<'tcx>, Layout<'tcx>),
527 i: VariantIdx,
528 ) -> Layout<'tcx> {
529 let ty = cx.tcx().erase_regions(ty);
530 TyAndLayout { ty, layout }.for_variant(&cx, i).layout
531 }
532}