Auto merge of #114208 - GKFX:offset_of_enum, r=wesleywiser

Support enum variants in offset_of!

This MR implements support for navigating through enum variants in `offset_of!`, placing the enum variant name in the second argument to `offset_of!`. The RFC placed it in the first argument, but I think it interacts better with nested field access in the second, as you can then write things like

```rust
offset_of!(Type, field.Variant.field)
```

Alternatively, a syntactic distinction could be made between variants and fields (e.g. `field::Variant.field`) but I'm not convinced this would be helpful.

[RFC 3308 # Enum Support](https://rust-lang.github.io/rfcs/3308-offset_of.html#enum-support-offset_ofsomeenumstructvariant-field_on_variant)
Tracking Issue #106655.

Author: bors

compiler/rustc_codegen_cranelift/src/base.rs

@@ -766,7 +766,7 @@ fn codegen_stmt<'tcx>(
                         NullOp::SizeOf => layout.size.bytes(),
                         NullOp::AlignOf => layout.align.abi.bytes(),
                         NullOp::OffsetOf(fields) => {
-                            layout.offset_of_subfield(fx, fields.iter().map(|f| f.index())).bytes()
+                            layout.offset_of_subfield(fx, fields.iter()).bytes()
                         }
                     };
                     let val = CValue::by_val(

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -680,7 +680,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         layout.align.abi.bytes()
                     }
                     mir::NullOp::OffsetOf(fields) => {
-                        layout.offset_of_subfield(bx.cx(), fields.iter().map(|f| f.index())).bytes()
+                        layout.offset_of_subfield(bx.cx(), fields.iter()).bytes()
                     }
                 };
                 let val = bx.cx().const_usize(val);

compiler/rustc_const_eval/src/interpret/step.rs

@@ -275,7 +275,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                     mir::NullOp::SizeOf => layout.size.bytes(),
                     mir::NullOp::AlignOf => layout.align.abi.bytes(),
                     mir::NullOp::OffsetOf(fields) => {
-                        layout.offset_of_subfield(self, fields.iter().map(|f| f.index())).bytes()
+                        layout.offset_of_subfield(self, fields.iter()).bytes()
                     }
                 };
                 self.write_scalar(Scalar::from_target_usize(val, self), &dest)?;

compiler/rustc_const_eval/src/transform/validate.rs

@@ -1056,16 +1056,23 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     }
                 }
             }
-            Rvalue::NullaryOp(NullOp::OffsetOf(fields), container) => {
+            Rvalue::NullaryOp(NullOp::OffsetOf(indices), container) => {
                 let fail_out_of_bounds = |this: &mut Self, location, field, ty| {
                     this.fail(location, format!("Out of bounds field {field:?} for {ty:?}"));
                 };
 
                 let mut current_ty = *container;
 
-                for field in fields.iter() {
+                for (variant, field) in indices.iter() {
                     match current_ty.kind() {
                         ty::Tuple(fields) => {
+                            if variant != FIRST_VARIANT {
+                                self.fail(
+                                    location,
+                                    format!("tried to get variant {variant:?} of tuple"),
+                                );
+                                return;
+                            }
                             let Some(&f_ty) = fields.get(field.as_usize()) else {
                                 fail_out_of_bounds(self, location, field, current_ty);
                                 return;
@@ -1074,15 +1081,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             current_ty = self.tcx.normalize_erasing_regions(self.param_env, f_ty);
                         }
                         ty::Adt(adt_def, args) => {
-                            if adt_def.is_enum() {
-                                self.fail(
-                                    location,
-                                    format!("Cannot get field offset from enum {current_ty:?}"),
-                                );
-                                return;
-                            }
-
-                            let Some(field) = adt_def.non_enum_variant().fields.get(field) else {
+                            let Some(field) = adt_def.variant(variant).fields.get(field) else {
                                 fail_out_of_bounds(self, location, field, current_ty);
                                 return;
                             };
@@ -1093,7 +1092,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                         _ => {
                             self.fail(
                                 location,
-                                format!("Cannot get field offset from non-adt type {current_ty:?}"),
+                                format!("Cannot get offset ({variant:?}, {field:?}) from type {current_ty:?}"),
                             );
                             return;
                         }

compiler/rustc_error_codes/src/error_codes.rs

@@ -514,6 +514,7 @@ E0791: include_str!("./error_codes/E0791.md"),
 E0792: include_str!("./error_codes/E0792.md"),
 E0793: include_str!("./error_codes/E0793.md"),
 E0794: include_str!("./error_codes/E0794.md"),
+E0795: include_str!("./error_codes/E0795.md"),
 }
 
 // Undocumented removed error codes. Note that many removed error codes are kept in the list above

compiler/rustc_error_codes/src/error_codes/E0795.md

@@ -0,0 +1,28 @@
+Invalid argument for the `offset_of!` macro.
+
+Erroneous code example:
+
+```compile_fail,E0795
+#![feature(offset_of)]
+
+let x = std::mem::offset_of!(Option<u8>, Some);
+```
+
+The `offset_of!` macro gives the offset of a field within a type. It can
+navigate through enum variants, but the final component of its second argument
+must be a field and not a variant.
+
+The offset of the contained `u8` in the `Option<u8>` can be found by specifying
+the field name `0`:
+
+```
+#![feature(offset_of)]
+
+let x: usize = std::mem::offset_of!(Option<u8>, Some.0);
+```
+
+The discriminant of an enumeration may be read with `core::mem::discriminant`,
+but this is not always a value physically present within the enum.
+
+Further information about enum layout may be found at
+https://rust-lang.github.io/unsafe-code-guidelines/layout/enums.html.

compiler/rustc_hir_typeck/src/expr.rs

@@ -54,7 +54,7 @@ use rustc_span::edit_distance::find_best_match_for_name;
 use rustc_span::hygiene::DesugaringKind;
 use rustc_span::source_map::{Span, Spanned};
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
-use rustc_target::abi::FieldIdx;
+use rustc_target::abi::{FieldIdx, FIRST_VARIANT};
 use rustc_target::spec::abi::Abi::RustIntrinsic;
 use rustc_trait_selection::infer::InferCtxtExt;
 use rustc_trait_selection::traits::error_reporting::TypeErrCtxtExt;
@@ -3107,12 +3107,81 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
         let mut field_indices = Vec::with_capacity(fields.len());
         let mut current_container = container;
+        let mut fields = fields.into_iter();
 
-        for &field in fields {
+        while let Some(&field) = fields.next() {
             let container = self.structurally_resolve_type(expr.span, current_container);
 
             match container.kind() {
-                ty::Adt(container_def, args) if !container_def.is_enum() => {
+                ty::Adt(container_def, args) if container_def.is_enum() => {
+                    let block = self.tcx.hir().local_def_id_to_hir_id(self.body_id);
+                    let (ident, _def_scope) =
+                        self.tcx.adjust_ident_and_get_scope(field, container_def.did(), block);
+
+                    let Some((index, variant)) = container_def.variants()
+                        .iter_enumerated()
+                        .find(|(_, v)| v.ident(self.tcx).normalize_to_macros_2_0() == ident) else {
+                        let mut err = type_error_struct!(
+                            self.tcx().sess,
+                            ident.span,
+                            container,
+                            E0599,
+                            "no variant named `{ident}` found for enum `{container}`",
+                            );
+                        err.span_label(field.span, "variant not found");
+                        err.emit();
+                        break;
+                    };
+                    let Some(&subfield) = fields.next() else {
+                        let mut err = type_error_struct!(
+                            self.tcx().sess,
+                            ident.span,
+                            container,
+                            E0795,
+                            "`{ident}` is an enum variant; expected field at end of `offset_of`",
+                            );
+                        err.span_label(field.span, "enum variant");
+                        err.emit();
+                        break;
+                    };
+                    let (subident, sub_def_scope) =
+                        self.tcx.adjust_ident_and_get_scope(subfield, variant.def_id, block);
+
+                    let Some((subindex, field)) = variant.fields
+                        .iter_enumerated()
+                        .find(|(_, f)| f.ident(self.tcx).normalize_to_macros_2_0() == subident) else {
+                        let mut err = type_error_struct!(
+                            self.tcx().sess,
+                            ident.span,
+                            container,
+                            E0609,
+                            "no field named `{subfield}` on enum variant `{container}::{ident}`",
+                            );
+                        err.span_label(field.span, "this enum variant...");
+                        err.span_label(subident.span, "...does not have this field");
+                        err.emit();
+                        break;
+                    };
+
+                    let field_ty = self.field_ty(expr.span, field, args);
+
+                    // FIXME: DSTs with static alignment should be allowed
+                    self.require_type_is_sized(field_ty, expr.span, traits::MiscObligation);
+
+                    if field.vis.is_accessible_from(sub_def_scope, self.tcx) {
+                        self.tcx.check_stability(field.did, Some(expr.hir_id), expr.span, None);
+                    } else {
+                        self.private_field_err(ident, container_def.did()).emit();
+                    }
+
+                    // Save the index of all fields regardless of their visibility in case
+                    // of error recovery.
+                    field_indices.push((index, subindex));
+                    current_container = field_ty;
+
+                    continue;
+                }
+                ty::Adt(container_def, args) => {
                     let block = self.tcx.hir().local_def_id_to_hir_id(self.body_id);
                     let (ident, def_scope) =
                         self.tcx.adjust_ident_and_get_scope(field, container_def.did(), block);
@@ -3135,7 +3204,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
                         // Save the index of all fields regardless of their visibility in case
                         // of error recovery.
-                        field_indices.push(index);
+                        field_indices.push((FIRST_VARIANT, index));
                         current_container = field_ty;
 
                         continue;
@@ -3149,7 +3218,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                             self.require_type_is_sized(ty, expr.span, traits::MiscObligation);
                         }
                         if let Some(&field_ty) = tys.get(index) {
-                            field_indices.push(index.into());
+                            field_indices.push((FIRST_VARIANT, index.into()));
                             current_container = field_ty;
 
                             continue;

compiler/rustc_middle/src/mir/syntax.rs

@@ -1354,7 +1354,7 @@ pub enum NullOp<'tcx> {
     /// Returns the minimum alignment of a type
     AlignOf,
     /// Returns the offset of a field
-    OffsetOf(&'tcx List<FieldIdx>),
+    OffsetOf(&'tcx List<(VariantIdx, FieldIdx)>),
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]

compiler/rustc_middle/src/thir.rs

@@ -492,7 +492,7 @@ pub enum ExprKind<'tcx> {
     /// Field offset (`offset_of!`)
     OffsetOf {
         container: Ty<'tcx>,
-        fields: &'tcx List<FieldIdx>,
+        fields: &'tcx List<(VariantIdx, FieldIdx)>,
     },
     /// An expression taking a reference to a thread local.
     ThreadLocalRef(DefId),

compiler/rustc_middle/src/ty/codec.rs

@@ -19,7 +19,7 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::ty::TyCtxt;
 use rustc_serialize::{Decodable, Encodable};
 use rustc_span::Span;
-use rustc_target::abi::FieldIdx;
+use rustc_target::abi::{FieldIdx, VariantIdx};
 pub use rustc_type_ir::{TyDecoder, TyEncoder};
 use std::hash::Hash;
 use std::intrinsics;
@@ -414,6 +414,17 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D> for ty::List<Fi
     }
 }
 
+impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
+    for ty::List<(VariantIdx, FieldIdx)>
+{
+    fn decode(decoder: &mut D) -> &'tcx Self {
+        let len = decoder.read_usize();
+        decoder.interner().mk_offset_of_from_iter(
+            (0..len).map::<(VariantIdx, FieldIdx), _>(|_| Decodable::decode(decoder)),
+        )
+    }
+}
+
 impl_decodable_via_ref! {
     &'tcx ty::TypeckResults<'tcx>,
     &'tcx ty::List<Ty<'tcx>>,
@@ -426,6 +437,7 @@ impl_decodable_via_ref! {
     &'tcx ty::List<ty::BoundVariableKind>,
     &'tcx ty::List<ty::Clause<'tcx>>,
     &'tcx ty::List<FieldIdx>,
+    &'tcx ty::List<(VariantIdx, FieldIdx)>,
 }
 
 #[macro_export]