rustc_transmute/layout/
tree.rs

1use std::ops::{ControlFlow, RangeInclusive};
2
3use super::{Byte, Def, Ref};
4
5#[cfg(test)]
6mod tests;
7
8/// A tree-based representation of a type layout.
9///
10/// Invariants:
11/// 1. All paths through the layout have the same length (in bytes).
12///
13/// Nice-to-haves:
14/// 1. An `Alt` is never directly nested beneath another `Alt`.
15/// 2. A `Seq` is never directly nested beneath another `Seq`.
16/// 3. `Seq`s and `Alt`s with a single member do not exist.
17#[derive(Clone, Debug, Hash, PartialEq, Eq)]
18pub(crate) enum Tree<D, R>
19where
20    D: Def,
21    R: Ref,
22{
23    /// A sequence of successive layouts.
24    Seq(Vec<Self>),
25    /// A choice between alternative layouts.
26    Alt(Vec<Self>),
27    /// A definition node.
28    Def(D),
29    /// A reference node.
30    Ref(R),
31    /// A byte node.
32    Byte(Byte),
33}
34
35#[derive(Debug, Copy, Clone, Eq, PartialEq)]
36pub(crate) enum Endian {
37    Little,
38    Big,
39}
40
41#[cfg(feature = "rustc")]
42impl From<rustc_abi::Endian> for Endian {
43    fn from(order: rustc_abi::Endian) -> Endian {
44        match order {
45            rustc_abi::Endian::Little => Endian::Little,
46            rustc_abi::Endian::Big => Endian::Big,
47        }
48    }
49}
50
51impl<D, R> Tree<D, R>
52where
53    D: Def,
54    R: Ref,
55{
56    /// A `Tree` consisting only of a definition node.
57    pub(crate) fn def(def: D) -> Self {
58        Self::Def(def)
59    }
60
61    /// A `Tree` representing an uninhabited type.
62    pub(crate) fn uninhabited() -> Self {
63        Self::Alt(vec![])
64    }
65
66    /// A `Tree` representing a zero-sized type.
67    pub(crate) fn unit() -> Self {
68        Self::Seq(Vec::new())
69    }
70
71    /// A `Tree` containing a single, uninitialized byte.
72    pub(crate) fn uninit() -> Self {
73        Self::Byte(Byte::uninit())
74    }
75
76    /// A `Tree` representing the layout of `bool`.
77    pub(crate) fn bool() -> Self {
78        Self::byte(0x00..=0x01)
79    }
80
81    /// A `Tree` whose layout matches that of a `u8`.
82    pub(crate) fn u8() -> Self {
83        Self::byte(0x00..=0xFF)
84    }
85
86    /// A `Tree` whose layout matches that of a `char`.
87    pub(crate) fn char(order: Endian) -> Self {
88        // `char`s can be in the following ranges:
89        // - [0, 0xD7FF]
90        // - [0xE000, 10FFFF]
91        //
92        // All other `char` values are illegal. We can thus represent a `char`
93        // as a union of three possible layouts:
94        // - 00 00 [00, D7] XX
95        // - 00 00 [E0, FF] XX
96        // - 00 [01, 10] XX XX
97
98        const _0: RangeInclusive<u8> = 0..=0;
99        const BYTE: RangeInclusive<u8> = 0x00..=0xFF;
100        let x = Self::from_big_endian(order, [_0, _0, 0x00..=0xD7, BYTE]);
101        let y = Self::from_big_endian(order, [_0, _0, 0xE0..=0xFF, BYTE]);
102        let z = Self::from_big_endian(order, [_0, 0x01..=0x10, BYTE, BYTE]);
103        Self::alt([x, y, z])
104    }
105
106    /// A `Tree` whose layout matches `std::num::NonZeroXxx`.
107    #[allow(dead_code)]
108    pub(crate) fn nonzero(width_in_bytes: u64) -> Self {
109        const BYTE: RangeInclusive<u8> = 0x00..=0xFF;
110        const NONZERO: RangeInclusive<u8> = 0x01..=0xFF;
111
112        (0..width_in_bytes)
113            .map(|nz_idx| {
114                (0..width_in_bytes)
115                    .map(|pos| Self::byte(if pos == nz_idx { NONZERO } else { BYTE }))
116                    .fold(Self::unit(), Self::then)
117            })
118            .fold(Self::uninhabited(), Self::or)
119    }
120
121    pub(crate) fn bytes<const N: usize, B: Into<Byte>>(bytes: [B; N]) -> Self {
122        Self::seq(bytes.map(B::into).map(Self::Byte))
123    }
124
125    pub(crate) fn byte(byte: impl Into<Byte>) -> Self {
126        Self::Byte(byte.into())
127    }
128
129    /// A `Tree` whose layout is a number of the given width.
130    pub(crate) fn number(width_in_bytes: u64) -> Self {
131        Self::Seq(vec![Self::u8(); width_in_bytes.try_into().unwrap()])
132    }
133
134    /// A `Tree` whose layout is entirely padding of the given width.
135    pub(crate) fn padding(width_in_bytes: usize) -> Self {
136        Self::Seq(vec![Self::uninit(); width_in_bytes])
137    }
138
139    /// Remove all `Def` nodes, and all branches of the layout for which `f`
140    /// produces `true`.
141    pub(crate) fn prune<F>(self, f: &F) -> Tree<!, R>
142    where
143        F: Fn(D) -> bool,
144    {
145        match self {
146            Self::Seq(elts) => match elts.into_iter().map(|elt| elt.prune(f)).try_fold(
147                Tree::unit(),
148                |elts, elt| {
149                    if elt == Tree::uninhabited() {
150                        ControlFlow::Break(Tree::uninhabited())
151                    } else {
152                        ControlFlow::Continue(elts.then(elt))
153                    }
154                },
155            ) {
156                ControlFlow::Break(node) | ControlFlow::Continue(node) => node,
157            },
158            Self::Alt(alts) => alts
159                .into_iter()
160                .map(|alt| alt.prune(f))
161                .fold(Tree::uninhabited(), |alts, alt| alts.or(alt)),
162            Self::Byte(b) => Tree::Byte(b),
163            Self::Ref(r) => Tree::Ref(r),
164            Self::Def(d) => {
165                if f(d) {
166                    Tree::uninhabited()
167                } else {
168                    Tree::unit()
169                }
170            }
171        }
172    }
173
174    /// Produces `true` if `Tree` is an inhabited type; otherwise false.
175    pub(crate) fn is_inhabited(&self) -> bool {
176        match self {
177            Self::Seq(elts) => elts.into_iter().all(|elt| elt.is_inhabited()),
178            Self::Alt(alts) => alts.into_iter().any(|alt| alt.is_inhabited()),
179            Self::Byte(..) | Self::Ref(..) | Self::Def(..) => true,
180        }
181    }
182
183    /// Produces a `Tree` which represents a sequence of bytes stored in
184    /// `order`.
185    ///
186    /// `bytes` is taken to be in big-endian byte order, and its order will be
187    /// swapped if `order == Endian::Little`.
188    pub(crate) fn from_big_endian<const N: usize, B: Into<Byte>>(
189        order: Endian,
190        mut bytes: [B; N],
191    ) -> Self {
192        if order == Endian::Little {
193            (&mut bytes[..]).reverse();
194        }
195
196        Self::bytes(bytes)
197    }
198
199    /// Produces a `Tree` where each of the trees in `trees` are sequenced one
200    /// after another.
201    pub(crate) fn seq<const N: usize>(trees: [Tree<D, R>; N]) -> Self {
202        trees.into_iter().fold(Tree::unit(), Self::then)
203    }
204
205    /// Produces a `Tree` where each of the trees in `trees` are accepted as
206    /// alternative layouts.
207    pub(crate) fn alt<const N: usize>(trees: [Tree<D, R>; N]) -> Self {
208        trees.into_iter().fold(Tree::uninhabited(), Self::or)
209    }
210
211    /// Produces a new `Tree` where `other` is sequenced after `self`.
212    pub(crate) fn then(self, other: Self) -> Self {
213        match (self, other) {
214            (Self::Seq(elts), other) | (other, Self::Seq(elts)) if elts.len() == 0 => other,
215            (Self::Seq(mut lhs), Self::Seq(mut rhs)) => {
216                lhs.append(&mut rhs);
217                Self::Seq(lhs)
218            }
219            (Self::Seq(mut lhs), rhs) => {
220                lhs.push(rhs);
221                Self::Seq(lhs)
222            }
223            (lhs, Self::Seq(mut rhs)) => {
224                rhs.insert(0, lhs);
225                Self::Seq(rhs)
226            }
227            (lhs, rhs) => Self::Seq(vec![lhs, rhs]),
228        }
229    }
230
231    /// Produces a new `Tree` accepting either `self` or `other` as alternative layouts.
232    pub(crate) fn or(self, other: Self) -> Self {
233        match (self, other) {
234            (Self::Alt(alts), other) | (other, Self::Alt(alts)) if alts.len() == 0 => other,
235            (Self::Alt(mut lhs), Self::Alt(rhs)) => {
236                lhs.extend(rhs);
237                Self::Alt(lhs)
238            }
239            (Self::Alt(mut alts), alt) | (alt, Self::Alt(mut alts)) => {
240                alts.push(alt);
241                Self::Alt(alts)
242            }
243            (lhs, rhs) => Self::Alt(vec![lhs, rhs]),
244        }
245    }
246}
247
248#[cfg(feature = "rustc")]
249pub(crate) mod rustc {
250    use rustc_abi::{
251        FieldIdx, FieldsShape, Layout, Size, TagEncoding, TyAndLayout, VariantIdx, Variants,
252    };
253    use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, LayoutError};
254    use rustc_middle::ty::{self, AdtDef, AdtKind, List, ScalarInt, Ty, TyCtxt, TypeVisitableExt};
255    use rustc_span::ErrorGuaranteed;
256
257    use super::Tree;
258    use crate::layout::rustc::{Def, Ref, layout_of};
259
260    #[derive(Debug, Copy, Clone)]
261    pub(crate) enum Err {
262        /// The layout of the type is not yet supported.
263        NotYetSupported,
264        /// This error will be surfaced elsewhere by rustc, so don't surface it.
265        UnknownLayout,
266        /// Overflow size
267        SizeOverflow,
268        TypeError(ErrorGuaranteed),
269    }
270
271    impl<'tcx> From<&LayoutError<'tcx>> for Err {
272        fn from(err: &LayoutError<'tcx>) -> Self {
273            match err {
274                LayoutError::Unknown(..)
275                | LayoutError::ReferencesError(..)
276                | LayoutError::TooGeneric(..)
277                | LayoutError::NormalizationFailure(..) => Self::UnknownLayout,
278                LayoutError::SizeOverflow(..) => Self::SizeOverflow,
279                LayoutError::Cycle(err) => Self::TypeError(*err),
280            }
281        }
282    }
283
284    impl<'tcx> Tree<Def<'tcx>, Ref<'tcx>> {
285        pub(crate) fn from_ty(ty: Ty<'tcx>, cx: LayoutCx<'tcx>) -> Result<Self, Err> {
286            use rustc_abi::HasDataLayout;
287            let layout = layout_of(cx, ty)?;
288
289            if let Err(e) = ty.error_reported() {
290                return Err(Err::TypeError(e));
291            }
292
293            let target = cx.data_layout();
294            let pointer_size = target.pointer_size;
295
296            match ty.kind() {
297                ty::Bool => Ok(Self::bool()),
298
299                ty::Float(nty) => {
300                    let width = nty.bit_width() / 8;
301                    Ok(Self::number(width.try_into().unwrap()))
302                }
303
304                ty::Int(nty) => {
305                    let width = nty.normalize(pointer_size.bits() as _).bit_width().unwrap() / 8;
306                    Ok(Self::number(width.try_into().unwrap()))
307                }
308
309                ty::Uint(nty) => {
310                    let width = nty.normalize(pointer_size.bits() as _).bit_width().unwrap() / 8;
311                    Ok(Self::number(width.try_into().unwrap()))
312                }
313
314                ty::Tuple(members) => Self::from_tuple((ty, layout), members, cx),
315
316                ty::Array(inner_ty, _len) => {
317                    let FieldsShape::Array { stride, count } = &layout.fields else {
318                        return Err(Err::NotYetSupported);
319                    };
320                    let inner_layout = layout_of(cx, *inner_ty)?;
321                    assert_eq!(*stride, inner_layout.size);
322                    let elt = Tree::from_ty(*inner_ty, cx)?;
323                    Ok(std::iter::repeat(elt)
324                        .take(*count as usize)
325                        .fold(Tree::unit(), |tree, elt| tree.then(elt)))
326                }
327
328                ty::Adt(adt_def, _args_ref) if !ty.is_box() => {
329                    let (lo, hi) = cx.tcx().layout_scalar_valid_range(adt_def.did());
330
331                    use core::ops::Bound::*;
332                    let is_transparent = adt_def.repr().transparent();
333                    match (adt_def.adt_kind(), lo, hi) {
334                        (AdtKind::Struct, Unbounded, Unbounded) => {
335                            Self::from_struct((ty, layout), *adt_def, cx)
336                        }
337                        (AdtKind::Struct, Included(1), Included(_hi)) if is_transparent => {
338                            // FIXME(@joshlf): Support `NonZero` types:
339                            // - Check to make sure that the first field is
340                            //   numerical
341                            // - Check to make sure that the upper bound is the
342                            //   maximum value for the field's type
343                            // - Construct `Self::nonzero`
344                            Err(Err::NotYetSupported)
345                        }
346                        (AdtKind::Enum, Unbounded, Unbounded) => {
347                            Self::from_enum((ty, layout), *adt_def, cx)
348                        }
349                        (AdtKind::Union, Unbounded, Unbounded) => {
350                            Self::from_union((ty, layout), *adt_def, cx)
351                        }
352                        _ => Err(Err::NotYetSupported),
353                    }
354                }
355
356                ty::Ref(lifetime, ty, mutability) => {
357                    let layout = layout_of(cx, *ty)?;
358                    let align = layout.align.abi.bytes_usize();
359                    let size = layout.size.bytes_usize();
360                    Ok(Tree::Ref(Ref {
361                        lifetime: *lifetime,
362                        ty: *ty,
363                        mutability: *mutability,
364                        align,
365                        size,
366                    }))
367                }
368
369                ty::Char => Ok(Self::char(cx.tcx().data_layout.endian.into())),
370
371                _ => Err(Err::NotYetSupported),
372            }
373        }
374
375        /// Constructs a `Tree` from a tuple.
376        fn from_tuple(
377            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
378            members: &'tcx List<Ty<'tcx>>,
379            cx: LayoutCx<'tcx>,
380        ) -> Result<Self, Err> {
381            match &layout.fields {
382                FieldsShape::Primitive => {
383                    assert_eq!(members.len(), 1);
384                    let inner_ty = members[0];
385                    Self::from_ty(inner_ty, cx)
386                }
387                FieldsShape::Arbitrary { offsets, .. } => {
388                    assert_eq!(offsets.len(), members.len());
389                    Self::from_variant(Def::Primitive, None, (ty, layout), layout.size, cx)
390                }
391                FieldsShape::Array { .. } | FieldsShape::Union(_) => Err(Err::NotYetSupported),
392            }
393        }
394
395        /// Constructs a `Tree` from a struct.
396        ///
397        /// # Panics
398        ///
399        /// Panics if `def` is not a struct definition.
400        fn from_struct(
401            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
402            def: AdtDef<'tcx>,
403            cx: LayoutCx<'tcx>,
404        ) -> Result<Self, Err> {
405            assert!(def.is_struct());
406            let def = Def::Adt(def);
407            Self::from_variant(def, None, (ty, layout), layout.size, cx)
408        }
409
410        /// Constructs a `Tree` from an enum.
411        ///
412        /// # Panics
413        ///
414        /// Panics if `def` is not an enum definition.
415        fn from_enum(
416            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
417            def: AdtDef<'tcx>,
418            cx: LayoutCx<'tcx>,
419        ) -> Result<Self, Err> {
420            assert!(def.is_enum());
421
422            // Computes the layout of a variant.
423            let layout_of_variant =
424                |index, encoding: Option<TagEncoding<VariantIdx>>| -> Result<Self, Err> {
425                    let variant_layout = ty_variant(cx, (ty, layout), index);
426                    if variant_layout.is_uninhabited() {
427                        return Ok(Self::uninhabited());
428                    }
429                    let tag = cx.tcx().tag_for_variant((cx.tcx().erase_regions(ty), index));
430                    let variant_def = Def::Variant(def.variant(index));
431                    Self::from_variant(
432                        variant_def,
433                        tag.map(|tag| (tag, index, encoding.unwrap())),
434                        (ty, variant_layout),
435                        layout.size,
436                        cx,
437                    )
438                };
439
440            match layout.variants() {
441                Variants::Empty => Ok(Self::uninhabited()),
442                Variants::Single { index } => {
443                    // `Variants::Single` on enums with variants denotes that
444                    // the enum delegates its layout to the variant at `index`.
445                    layout_of_variant(*index, None)
446                }
447                Variants::Multiple { tag: _, tag_encoding, tag_field, .. } => {
448                    // `Variants::Multiple` denotes an enum with multiple
449                    // variants. The layout of such an enum is the disjunction
450                    // of the layouts of its tagged variants.
451
452                    // For enums (but not coroutines), the tag field is
453                    // currently always the first field of the layout.
454                    assert_eq!(*tag_field, 0);
455
456                    let variants = def.discriminants(cx.tcx()).try_fold(
457                        Self::uninhabited(),
458                        |variants, (idx, _discriminant)| {
459                            let variant = layout_of_variant(idx, Some(tag_encoding.clone()))?;
460                            Result::<Self, Err>::Ok(variants.or(variant))
461                        },
462                    )?;
463
464                    Ok(Self::def(Def::Adt(def)).then(variants))
465                }
466            }
467        }
468
469        /// Constructs a `Tree` from a 'variant-like' layout.
470        ///
471        /// A 'variant-like' layout includes those of structs and, of course,
472        /// enum variants. Pragmatically speaking, this method supports anything
473        /// with `FieldsShape::Arbitrary`.
474        ///
475        /// Note: This routine assumes that the optional `tag` is the first
476        /// field, and enum callers should check that `tag_field` is, in fact,
477        /// `0`.
478        fn from_variant(
479            def: Def<'tcx>,
480            tag: Option<(ScalarInt, VariantIdx, TagEncoding<VariantIdx>)>,
481            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
482            total_size: Size,
483            cx: LayoutCx<'tcx>,
484        ) -> Result<Self, Err> {
485            // This constructor does not support non-`FieldsShape::Arbitrary`
486            // layouts.
487            let FieldsShape::Arbitrary { offsets, memory_index } = layout.fields() else {
488                return Err(Err::NotYetSupported);
489            };
490
491            // When this function is invoked with enum variants,
492            // `ty_and_layout.size` does not encompass the entire size of the
493            // enum. We rely on `total_size` for this.
494            assert!(layout.size <= total_size);
495
496            let mut size = Size::ZERO;
497            let mut struct_tree = Self::def(def);
498
499            // If a `tag` is provided, place it at the start of the layout.
500            if let Some((tag, index, encoding)) = &tag {
501                match encoding {
502                    TagEncoding::Direct => {
503                        size += tag.size();
504                    }
505                    TagEncoding::Niche { niche_variants, .. } => {
506                        if !niche_variants.contains(index) {
507                            size += tag.size();
508                        }
509                    }
510                }
511                struct_tree = struct_tree.then(Self::from_tag(*tag, cx.tcx()));
512            }
513
514            // Append the fields, in memory order, to the layout.
515            let inverse_memory_index = memory_index.invert_bijective_mapping();
516            for &field_idx in inverse_memory_index.iter() {
517                // Add interfield padding.
518                let padding_needed = offsets[field_idx] - size;
519                let padding = Self::padding(padding_needed.bytes_usize());
520
521                let field_ty = ty_field(cx, (ty, layout), field_idx);
522                let field_layout = layout_of(cx, field_ty)?;
523                let field_tree = Self::from_ty(field_ty, cx)?;
524
525                struct_tree = struct_tree.then(padding).then(field_tree);
526
527                size += padding_needed + field_layout.size;
528            }
529
530            // Add trailing padding.
531            let padding_needed = total_size - size;
532            let trailing_padding = Self::padding(padding_needed.bytes_usize());
533
534            Ok(struct_tree.then(trailing_padding))
535        }
536
537        /// Constructs a `Tree` representing the value of a enum tag.
538        fn from_tag(tag: ScalarInt, tcx: TyCtxt<'tcx>) -> Self {
539            use rustc_abi::Endian;
540            let size = tag.size();
541            let bits = tag.to_bits(size);
542            let bytes: [u8; 16];
543            let bytes = match tcx.data_layout.endian {
544                Endian::Little => {
545                    bytes = bits.to_le_bytes();
546                    &bytes[..size.bytes_usize()]
547                }
548                Endian::Big => {
549                    bytes = bits.to_be_bytes();
550                    &bytes[bytes.len() - size.bytes_usize()..]
551                }
552            };
553            Self::Seq(bytes.iter().map(|&b| Self::byte(b)).collect())
554        }
555
556        /// Constructs a `Tree` from a union.
557        ///
558        /// # Panics
559        ///
560        /// Panics if `def` is not a union definition.
561        fn from_union(
562            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
563            def: AdtDef<'tcx>,
564            cx: LayoutCx<'tcx>,
565        ) -> Result<Self, Err> {
566            assert!(def.is_union());
567
568            // This constructor does not support non-`FieldsShape::Union`
569            // layouts. Fields of this shape are all placed at offset 0.
570            let FieldsShape::Union(_fields) = layout.fields() else {
571                return Err(Err::NotYetSupported);
572            };
573
574            let fields = &def.non_enum_variant().fields;
575            let fields = fields.iter_enumerated().try_fold(
576                Self::uninhabited(),
577                |fields, (idx, _field_def)| {
578                    let field_ty = ty_field(cx, (ty, layout), idx);
579                    let field_layout = layout_of(cx, field_ty)?;
580                    let field = Self::from_ty(field_ty, cx)?;
581                    let trailing_padding_needed = layout.size - field_layout.size;
582                    let trailing_padding = Self::padding(trailing_padding_needed.bytes_usize());
583                    let field_and_padding = field.then(trailing_padding);
584                    Result::<Self, Err>::Ok(fields.or(field_and_padding))
585                },
586            )?;
587
588            Ok(Self::def(Def::Adt(def)).then(fields))
589        }
590    }
591
592    fn ty_field<'tcx>(
593        cx: LayoutCx<'tcx>,
594        (ty, layout): (Ty<'tcx>, Layout<'tcx>),
595        i: FieldIdx,
596    ) -> Ty<'tcx> {
597        // We cannot use `ty_and_layout_field` to retrieve the field type, since
598        // `ty_and_layout_field` erases regions in the returned type. We must
599        // not erase regions here, since we may need to ultimately emit outlives
600        // obligations as a consequence of the transmutability analysis.
601        match ty.kind() {
602            ty::Adt(def, args) => {
603                match layout.variants {
604                    Variants::Single { index } => {
605                        let field = &def.variant(index).fields[i];
606                        field.ty(cx.tcx(), args)
607                    }
608                    Variants::Empty => panic!("there is no field in Variants::Empty types"),
609                    // Discriminant field for enums (where applicable).
610                    Variants::Multiple { tag, .. } => {
611                        assert_eq!(i.as_usize(), 0);
612                        ty::layout::PrimitiveExt::to_ty(&tag.primitive(), cx.tcx())
613                    }
614                }
615            }
616            ty::Tuple(fields) => fields[i.as_usize()],
617            kind => unimplemented!(
618                "only a subset of `Ty::ty_and_layout_field`'s functionality is implemented. implementation needed for {:?}",
619                kind
620            ),
621        }
622    }
623
624    fn ty_variant<'tcx>(
625        cx: LayoutCx<'tcx>,
626        (ty, layout): (Ty<'tcx>, Layout<'tcx>),
627        i: VariantIdx,
628    ) -> Layout<'tcx> {
629        let ty = cx.tcx().erase_regions(ty);
630        TyAndLayout { ty, layout }.for_variant(&cx, i).layout
631    }
632}