Support enums in type info reflection

Author: SpriteOvO

compiler/rustc_const_eval/src/const_eval/type_info.rs

@@ -110,15 +110,11 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
                             variant
                         }
                         ty::Adt(adt_def, generics) => {
-                            // TODO(type_info): Handle enum and union
-                            if !adt_def.is_struct() {
+                            // TODO(type_info): Handle union
+                            if !adt_def.is_struct() && !adt_def.is_enum() {
                                 self.downcast(&field_dest, sym::Other)?.0
                             } else {
-                                let (variant, variant_place) =
-                                    self.downcast(&field_dest, sym::Struct)?;
-                                let place = self.project_field(&variant_place, FieldIdx::ZERO)?;
-                                self.write_adt_type_info(place, (ty, *adt_def), generics)?;
-                                variant
+                                self.write_adt_type_info(&field_dest, (ty, *adt_def), generics)?
                             }
                         }
                         ty::Bool => {
@@ -253,6 +249,21 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
         self.write_immediate(ptr, &fields_slice_place)
     }
 
+    // Write fields for struct, enum variants
+    fn write_variant_fields(
+        &mut self,
+        place: impl Writeable<'tcx, CtfeProvenance>,
+        variant_def: &'tcx VariantDef,
+        variant_layout: TyAndLayout<'tcx>,
+        generics: &'tcx GenericArgs<'tcx>,
+    ) -> InterpResult<'tcx> {
+        self.project_write_array(place, variant_def.fields.len() as u64, |this, i, place| {
+            let field_def = &variant_def.fields[FieldIdx::from_usize(i as usize)];
+            let field_ty = field_def.ty(*this.tcx, generics);
+            this.write_field(field_ty, place, variant_layout, Some(field_def.name), i)
+        })
+    }
+
     fn write_field(
         &mut self,
         field_ty: Ty<'tcx>,
@@ -318,20 +329,31 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
     // FIXME(type_info): No semver considerations for now
     pub(crate) fn write_adt_type_info(
         &mut self,
-        place: impl Writeable<'tcx, CtfeProvenance>,
+        place: &(impl Writeable<'tcx, CtfeProvenance> + 'tcx),
         adt: (Ty<'tcx>, AdtDef<'tcx>),
         generics: &'tcx GenericArgs<'tcx>,
-    ) -> InterpResult<'tcx> {
+    ) -> InterpResult<'tcx, VariantIdx> {
         let (adt_ty, adt_def) = adt;
-        match adt_def.adt_kind() {
-            AdtKind::Struct => self.write_struct_type_info(
-                place,
-                (adt_ty, adt_def.variant(VariantIdx::ZERO)),
-                generics,
-            ),
+        let variant_idx = match adt_def.adt_kind() {
+            AdtKind::Struct => {
+                let (variant, variant_place) = self.downcast(place, sym::Struct)?;
+                let place = self.project_field(&variant_place, FieldIdx::ZERO)?;
+                self.write_struct_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!(),
-            AdtKind::Enum => todo!(),
-        }
+        };
+        interp_ok(variant_idx)
     }
 
     pub(crate) fn write_struct_type_info(
@@ -350,15 +372,9 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
 
             match field.name {
                 sym::generics => self.write_generics(field_place, generics)?,
-                sym::fields => self.project_write_array(
-                    field_place,
-                    struct_def.fields.len() as u64,
-                    |this, i, place| {
-                        let field_def = &struct_def.fields[FieldIdx::from_usize(i as usize)];
-                        let field_ty = field_def.ty(*this.tcx, generics);
-                        this.write_field(field_ty, place, struct_layout, Some(field_def.name), i)
-                    },
-                )?,
+                sym::fields => {
+                    self.write_variant_fields(field_place, struct_def, struct_layout, generics)?
+                }
                 sym::non_exhaustive => {
                     let is_non_exhaustive = struct_def.is_field_list_non_exhaustive();
                     self.write_scalar(Scalar::from_bool(is_non_exhaustive), &field_place)?
@@ -371,6 +387,77 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
         interp_ok(())
     }
 
+    pub(crate) fn write_enum_type_info(
+        &mut self,
+        place: impl Writeable<'tcx, CtfeProvenance>,
+        enum_: (Ty<'tcx>, AdtDef<'tcx>),
+        generics: &'tcx GenericArgs<'tcx>,
+    ) -> InterpResult<'tcx> {
+        let (enum_ty, enum_def) = enum_;
+        let enum_layout = self.layout_of(enum_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::variants => {
+                    self.project_write_array(
+                        field_place,
+                        enum_def.variants().len() as u64,
+                        |this, i, place| {
+                            let variant_idx = VariantIdx::from_usize(i as usize);
+                            let variant_def = &enum_def.variants()[variant_idx];
+                            let variant_layout = enum_layout.for_variant(this, variant_idx);
+                            // TODO(type_info): Is it correct to use enum_ty here? If yes, leave some explanation.
+                            this.write_enum_variant(place, (variant_layout, &variant_def), generics)
+                        },
+                    )?;
+                }
+                sym::non_exhaustive => {
+                    let is_non_exhaustive = enum_def.is_variant_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(())
+    }
+
+    fn write_enum_variant(
+        &mut self,
+        place: impl Writeable<'tcx, CtfeProvenance>,
+        variant: (TyAndLayout<'tcx>, &'tcx VariantDef),
+        generics: &'tcx GenericArgs<'tcx>,
+    ) -> InterpResult<'tcx> {
+        let (variant_layout, variant_def) = variant;
+
+        for (field_idx, field_def) 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_def.name {
+                sym::name => {
+                    let name_place = self.allocate_str_dedup(variant_def.name.as_str())?;
+                    let ptr = self.mplace_to_ref(&name_place)?;
+                    self.write_immediate(*ptr, &field_place)?
+                }
+                sym::fields => {
+                    self.write_variant_fields(field_place, &variant_def, variant_layout, generics)?
+                }
+                sym::non_exhaustive => {
+                    let is_non_exhaustive = variant_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_def.did), "unimplemented field {other}"),
+            }
+        }
+        interp_ok(())
+    }
+
     fn write_generics(
         &mut self,
         place: impl Writeable<'tcx, CtfeProvenance>,

compiler/rustc_span/src/symbol.rs

@@ -234,6 +234,7 @@ symbols! {
         Duration,
         Encodable,
         Encoder,
+        Enum,
         Enumerate,
         Eq,
         Equal,
@@ -2469,6 +2470,7 @@ symbols! {
         values,
         var,
         variant_count,
+        variants,
         vec,
         vec_as_mut_slice,
         vec_as_slice,

library/core/src/mem/type_info.rs

@@ -47,6 +47,8 @@ pub enum TypeKind {
     Array(Array),
     /// Structs.
     Struct(Struct),
+    /// Enums.
+    Enum(Enum),
     /// Primitive boolean type.
     Bool(Bool),
     /// Primitive character type.
@@ -109,6 +111,32 @@ pub struct Struct {
     pub non_exhaustive: bool,
 }
 
+/// Compile-time type information about enums.
+#[derive(Debug)]
+#[non_exhaustive]
+#[unstable(feature = "type_info", issue = "146922")]
+pub struct Enum {
+    /// Instantiated generics of the enum.
+    pub generics: &'static [Generic],
+    /// All variants of the enum.
+    pub variants: &'static [Variant],
+    /// Whether the enum variant list is non-exhaustive.
+    pub non_exhaustive: bool,
+}
+
+/// Compile-time type information about variants of enums.
+#[derive(Debug)]
+#[non_exhaustive]
+#[unstable(feature = "type_info", issue = "146922")]
+pub struct Variant {
+    /// The name of the variant.
+    pub name: &'static str,
+    /// All fields of the variant.
+    pub fields: &'static [Field],
+    /// Whether the enum variant fields is non-exhaustive.
+    pub non_exhaustive: bool,
+}
+
 /// Compile-time type information about instantiated generics of structs, enum and union variants.
 #[derive(Debug)]
 #[non_exhaustive]

library/coretests/tests/mem/type_info.rs

@@ -119,6 +119,48 @@ fn test_structs() {
     }
 }
 
+#[test]
+fn test_enums() {
+    use TypeKind::*;
+
+    const {
+        enum E {
+            Some(u32),
+            None,
+            #[non_exhaustive]
+            Foomp {
+                a: (),
+                b: &'static str,
+            },
+        }
+
+        let Type { kind: Enum(ty), size, .. } = Type::of::<E>() else { panic!() };
+        assert!(size == Some(size_of::<E>()));
+        assert!(ty.variants.len() == 3);
+
+        assert!(ty.variants[0].name == "Some");
+        assert!(!ty.variants[0].non_exhaustive);
+        assert!(ty.variants[0].fields.len() == 1);
+
+        assert!(ty.variants[1].name == "None");
+        assert!(!ty.variants[1].non_exhaustive);
+        assert!(ty.variants[1].fields.len() == 0);
+
+        assert!(ty.variants[2].name == "Foomp");
+        assert!(ty.variants[2].non_exhaustive);
+        assert!(ty.variants[2].fields.len() == 2);
+    }
+
+    const {
+        let Type { kind: Enum(ty), size, .. } = Type::of::<Option<i32>>() else { panic!() };
+        assert!(size == Some(size_of::<Option<i32>>()));
+        assert!(ty.variants.len() == 2);
+        assert!(ty.generics.len() == 1);
+        let Generic::Type(GenericType { ty: generic_ty, .. }) = ty.generics[0] else { panic!() };
+        assert!(generic_ty == TypeId::of::<i32>());
+    }
+}
+
 #[test]
 fn test_primitives() {
     use TypeKind::*;

tests/ui/privacy/issue-79593.rs

@@ -10,7 +10,7 @@ mod foo {
         Pub {};
         //~^ ERROR missing field `private` in initializer of `Pub`
         Enum::Variant { x: () };
-        //~^ ERROR missing field `y` in initializer of `Enum`
+        //~^ ERROR missing field `y` in initializer of `foo::Enum`
     }
 }
 
@@ -21,9 +21,9 @@ fn correct() {
 
 fn wrong() {
     foo::Enum::Variant { x: () };
-    //~^ ERROR missing field `y` in initializer of `Enum`
+    //~^ ERROR missing field `y` in initializer of `foo::Enum`
     foo::Enum::Variant { };
-    //~^ ERROR missing fields `x` and `y` in initializer of `Enum`
+    //~^ ERROR missing fields `x` and `y` in initializer of `foo::Enum`
 }
 
 fn main() {}

tests/ui/privacy/issue-79593.stderr

@@ -4,7 +4,7 @@ error[E0063]: missing field `private` in initializer of `Pub`
 LL |         Pub {};
    |         ^^^ missing `private`
 
-error[E0063]: missing field `y` in initializer of `Enum`
+error[E0063]: missing field `y` in initializer of `foo::Enum`
   --> $DIR/issue-79593.rs:12:9
    |
 LL |         Enum::Variant { x: () };
@@ -18,13 +18,13 @@ LL |     foo::Pub {};
    |
    = note: private field `private` that was not provided
 
-error[E0063]: missing field `y` in initializer of `Enum`
+error[E0063]: missing field `y` in initializer of `foo::Enum`
   --> $DIR/issue-79593.rs:23:5
    |
 LL |     foo::Enum::Variant { x: () };
    |     ^^^^^^^^^^^^^^^^^^ missing `y`
 
-error[E0063]: missing fields `x` and `y` in initializer of `Enum`
+error[E0063]: missing fields `x` and `y` in initializer of `foo::Enum`
   --> $DIR/issue-79593.rs:25:5
    |
 LL |     foo::Enum::Variant { };

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

@@ -164,7 +164,46 @@ Type {
     ),
 }
 Type {
-    kind: Other,
+    kind: Enum(
+        Enum {
+            generics: [],
+            variants: [
+                Variant {
+                    name: "Some",
+                    fields: [
+                        Field {
+                            name: "0",
+                            ty: TypeId(0x1378bb1c0a0202683eb65e7c11f2e4d7),
+                            offset: 4,
+                        },
+                    ],
+                    non_exhaustive: false,
+                },
+                Variant {
+                    name: "None",
+                    fields: [],
+                    non_exhaustive: false,
+                },
+                Variant {
+                    name: "Foomp",
+                    fields: [
+                        Field {
+                            name: "a",
+                            ty: TypeId(0x41223169ff28813ba79b7268a2a968d9),
+                            offset: 4,
+                        },
+                        Field {
+                            name: "b",
+                            ty: TypeId(0xb98b1b7157a6417863eb502cd6cb5d6d),
+                            offset: 4,
+                        },
+                    ],
+                    non_exhaustive: true,
+                },
+            ],
+            non_exhaustive: false,
+        },
+    ),
     size: Some(
         12,
     ),