Support unions in type info reflection

Author: SpriteOvO

compiler/rustc_const_eval/src/const_eval/type_info.rs

@@ -119,12 +119,7 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
                             variant
                         }
                         ty::Adt(adt_def, generics) => {
-                            // TODO(type_info): Handle union
-                            if !adt_def.is_struct() && !adt_def.is_enum() {
-                                self.downcast(&field_dest, sym::Other)?.0
-                            } else {
-                                self.write_adt_type_info(&field_dest, (ty, *adt_def), generics)?
-                            }
+                            self.write_adt_type_info(&field_dest, (ty, *adt_def), generics)?
                         }
                         ty::Bool => {
                             let (variant, _variant_place) =
@@ -389,13 +384,22 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
                 )?;
                 variant
             }
+            AdtKind::Union => {
+                let (variant, variant_place) = self.downcast(place, sym::Union)?;
+                let place = self.project_field(&variant_place, FieldIdx::ZERO)?;
+                self.write_union_type_info(
+                    place,
+                    (adt_ty, adt_def.variant(VariantIdx::ZERO)),
+                    generics,
+                )?;
+                variant
+            }
             AdtKind::Enum => {
                 let (variant, variant_place) = self.downcast(place, sym::Enum)?;
                 let place = self.project_field(&variant_place, FieldIdx::ZERO)?;
                 self.write_enum_type_info(place, adt, generics)?;
                 variant
             }
-            AdtKind::Union => todo!(),
         };
         interp_ok(variant_idx)
     }
@@ -430,6 +434,36 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
         interp_ok(())
     }
 
+    pub(crate) fn write_union_type_info(
+        &mut self,
+        place: impl Writeable<'tcx, CtfeProvenance>,
+        union_: (Ty<'tcx>, &'tcx VariantDef),
+        generics: &'tcx GenericArgs<'tcx>,
+    ) -> InterpResult<'tcx> {
+        let (union_ty, union_def) = union_;
+        let union_layout = self.layout_of(union_ty)?;
+
+        for (field_idx, field) in
+            place.layout().ty.ty_adt_def().unwrap().non_enum_variant().fields.iter_enumerated()
+        {
+            let field_place = self.project_field(&place, field_idx)?;
+
+            match field.name {
+                sym::generics => self.write_generics(field_place, generics)?,
+                sym::fields => {
+                    self.write_variant_fields(field_place, union_def, union_layout, generics)?
+                }
+                sym::non_exhaustive => {
+                    let is_non_exhaustive = union_def.is_field_list_non_exhaustive();
+                    self.write_scalar(Scalar::from_bool(is_non_exhaustive), &field_place)?
+                }
+                other => span_bug!(self.tcx.def_span(field.did), "unimplemented field {other}"),
+            }
+        }
+
+        interp_ok(())
+    }
+
     pub(crate) fn write_enum_type_info(
         &mut self,
         place: impl Writeable<'tcx, CtfeProvenance>,

compiler/rustc_span/src/symbol.rs

@@ -400,6 +400,7 @@ symbols! {
         TyCtxt,
         TyKind,
         Type,
+        Union,
         Unknown,
         Unsize,
         UnsizedConstParamTy,

library/core/src/mem/type_info.rs

@@ -53,6 +53,8 @@ pub enum TypeKind {
     Struct(Struct),
     /// Enums.
     Enum(Enum),
+    /// Unions.
+    Union(Union),
     /// Primitive boolean type.
     Bool(Bool),
     /// Primitive character type.
@@ -156,6 +158,17 @@ pub struct Struct {
     pub non_exhaustive: bool,
 }
 
+/// Compile-time type information about unions.
+#[derive(Debug)]
+#[non_exhaustive]
+#[unstable(feature = "type_info", issue = "146922")]
+pub struct Union {
+    /// Instantiated generics of the union.
+    pub generics: &'static [Generic],
+    /// All fields of the union.
+    pub fields: &'static [Field],
+}
+
 /// Compile-time type information about enums.
 #[derive(Debug)]
 #[non_exhaustive]

library/coretests/tests/mem/type_info.rs

@@ -127,6 +127,47 @@ fn test_structs() {
     }
 }
 
+#[test]
+fn test_unions() {
+    use TypeKind::*;
+
+    const {
+        union TestUnion {
+            first: i16,
+            second: u16,
+        }
+
+        let Type { kind: Union(ty), size, .. } = Type::of::<TestUnion>() else { panic!() };
+        assert!(size == Some(size_of::<TestUnion>()));
+        assert!(ty.fields.len() == 2);
+        assert!(ty.fields[0].name == "first");
+        assert!(ty.fields[0].offset == 0);
+        assert!(ty.fields[1].name == "second");
+        assert!(ty.fields[1].offset == 0);
+    }
+
+    const {
+        union Generics<'a, T: Copy, const C: u64> {
+            a: T,
+            z: &'a (),
+        }
+
+        let Type { kind: Union(ty), .. } = Type::of::<Generics<'static, i32, 1_u64>>() else {
+            panic!()
+        };
+        assert!(ty.fields.len() == 2);
+        assert!(ty.fields[0].offset == 0);
+        assert!(ty.fields[1].offset == 0);
+
+        assert!(ty.generics.len() == 3);
+        let Generic::Lifetime(_) = ty.generics[0] else { panic!() };
+        let Generic::Type(GenericType { ty: generic_ty, .. }) = ty.generics[1] else { panic!() };
+        assert!(generic_ty == TypeId::of::<i32>());
+        let Generic::Const(Const { ty: const_ty, .. }) = ty.generics[2] else { panic!() };
+        assert!(const_ty == TypeId::of::<u64>());
+    }
+}
+
 #[test]
 fn test_enums() {
     use TypeKind::*;

tests/ui/reflection/dump.bit32.run.stdout

@@ -296,6 +296,28 @@ Type {
         12,
     ),
 }
+Type {
+    kind: Union(
+        Union {
+            generics: [],
+            fields: [
+                Field {
+                    name: "first",
+                    ty: TypeId(0xdeb66dd04fa9db03298cc77926096aae),
+                    offset: 0,
+                },
+                Field {
+                    name: "second",
+                    ty: TypeId(0xc50c4a8d8e150aa67101203f1fab1cd7),
+                    offset: 0,
+                },
+            ],
+        },
+    ),
+    size: Some(
+        2,
+    ),
+}
 Type {
     kind: Reference(
         Reference {

tests/ui/reflection/dump.bit64.run.stdout

@@ -296,6 +296,28 @@ Type {
         16,
     ),
 }
+Type {
+    kind: Union(
+        Union {
+            generics: [],
+            fields: [
+                Field {
+                    name: "first",
+                    ty: TypeId(0xdeb66dd04fa9db03298cc77926096aae),
+                    offset: 0,
+                },
+                Field {
+                    name: "second",
+                    ty: TypeId(0xc50c4a8d8e150aa67101203f1fab1cd7),
+                    offset: 0,
+                },
+            ],
+        },
+    ),
+    size: Some(
+        2,
+    ),
+}
 Type {
     kind: Reference(
         Reference {

tests/ui/reflection/dump.rs

@@ -42,6 +42,11 @@ struct Unsized {
     s: str,
 }
 
+union Union {
+    first: i16,
+    second: u16,
+}
+
 macro_rules! dump_types {
     ($($ty:ty),+ $(,)?) => {
         $(println!("{:#?}", const { Type::of::<$ty>() });)+
@@ -54,7 +59,7 @@ fn main() {
         [u8; 2],
         i8, i32, i64, i128, isize,
         u8, u32, u64, u128, usize,
-        Foo, Bar, NonExhaustiveStruct, TupleStruct, Generics<i32, u32, 1>,
+        Foo, Bar, NonExhaustiveStruct, TupleStruct, Generics<i32, u32, 1>, Union,
         &Unsized, &str, &[u8],
         str, [u8],
         &u8, &mut u8,