Auto merge of #106934 - DrMeepster:offset_of, r=WaffleLapkin

Add offset_of! macro (RFC 3308)

Implements rust-lang/rfcs#3308 (tracking issue #106655) by adding the built in macro `core::mem::offset_of`. Two of the future possibilities are also implemented:

* Nested field accesses (without array indexing)
* DST support (for `Sized` fields)

I wrote this a few months ago, before the RFC merged. Now that it's merged, I decided to rebase and finish it.

cc `@thomcc` (RFC author)

Author: bors

compiler/rustc_ast/src/ast.rs

@@ -1271,6 +1271,7 @@ impl Expr {
             ExprKind::Continue(..) => ExprPrecedence::Continue,
             ExprKind::Ret(..) => ExprPrecedence::Ret,
             ExprKind::InlineAsm(..) => ExprPrecedence::InlineAsm,
+            ExprKind::OffsetOf(..) => ExprPrecedence::OffsetOf,
             ExprKind::MacCall(..) => ExprPrecedence::Mac,
             ExprKind::Struct(..) => ExprPrecedence::Struct,
             ExprKind::Repeat(..) => ExprPrecedence::Repeat,
@@ -1469,6 +1470,9 @@ pub enum ExprKind {
     /// Output of the `asm!()` macro.
     InlineAsm(P<InlineAsm>),
 
+    /// Output of the `offset_of!()` macro.
+    OffsetOf(P<Ty>, P<[Ident]>),
+
     /// A macro invocation; pre-expansion.
     MacCall(P<MacCall>),
 

compiler/rustc_ast/src/mut_visit.rs

@@ -1456,6 +1456,12 @@ pub fn noop_visit_expr<T: MutVisitor>(
         }
         ExprKind::InlineAsm(asm) => vis.visit_inline_asm(asm),
         ExprKind::FormatArgs(fmt) => vis.visit_format_args(fmt),
+        ExprKind::OffsetOf(container, fields) => {
+            vis.visit_ty(container);
+            for field in fields.iter_mut() {
+                vis.visit_ident(field);
+            }
+        }
         ExprKind::MacCall(mac) => vis.visit_mac_call(mac),
         ExprKind::Struct(se) => {
             let StructExpr { qself, path, fields, rest } = se.deref_mut();

compiler/rustc_ast/src/util/parser.rs

@@ -269,6 +269,7 @@ pub enum ExprPrecedence {
     Index,
     Try,
     InlineAsm,
+    OffsetOf,
     Mac,
     FormatArgs,
 
@@ -335,7 +336,8 @@ impl ExprPrecedence {
             | ExprPrecedence::Try
             | ExprPrecedence::InlineAsm
             | ExprPrecedence::Mac
-            | ExprPrecedence::FormatArgs => PREC_POSTFIX,
+            | ExprPrecedence::FormatArgs
+            | ExprPrecedence::OffsetOf => PREC_POSTFIX,
 
             // Never need parens
             ExprPrecedence::Array

compiler/rustc_ast/src/visit.rs

@@ -909,6 +909,12 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
         ExprKind::Paren(subexpression) => visitor.visit_expr(subexpression),
         ExprKind::InlineAsm(asm) => visitor.visit_inline_asm(asm),
         ExprKind::FormatArgs(f) => visitor.visit_format_args(f),
+        ExprKind::OffsetOf(container, fields) => {
+            visitor.visit_ty(container);
+            for &field in fields {
+                visitor.visit_ident(field);
+            }
+        }
         ExprKind::Yield(optional_expression) => {
             walk_list!(visitor, visit_expr, optional_expression);
         }

compiler/rustc_ast_lowering/src/expr.rs

@@ -289,6 +289,13 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     hir::ExprKind::InlineAsm(self.lower_inline_asm(e.span, asm))
                 }
                 ExprKind::FormatArgs(fmt) => self.lower_format_args(e.span, fmt),
+                ExprKind::OffsetOf(container, fields) => hir::ExprKind::OffsetOf(
+                    self.lower_ty(
+                        container,
+                        &mut ImplTraitContext::Disallowed(ImplTraitPosition::OffsetOf),
+                    ),
+                    self.arena.alloc_from_iter(fields.iter().map(|&ident| self.lower_ident(ident))),
+                ),
                 ExprKind::Struct(se) => {
                     let rest = match &se.rest {
                         StructRest::Base(e) => Some(self.lower_expr(e)),

compiler/rustc_ast_lowering/src/lib.rs

@@ -283,6 +283,7 @@ enum ImplTraitPosition {
     FieldTy,
     Cast,
     ImplSelf,
+    OffsetOf,
 }
 
 impl std::fmt::Display for ImplTraitPosition {
@@ -313,6 +314,7 @@ impl std::fmt::Display for ImplTraitPosition {
             ImplTraitPosition::FieldTy => "field types",
             ImplTraitPosition::Cast => "cast types",
             ImplTraitPosition::ImplSelf => "impl headers",
+            ImplTraitPosition::OffsetOf => "`offset_of!` params",
         };
 
         write!(f, "{name}")

compiler/rustc_ast_pretty/src/pprust/state/expr.rs

@@ -549,6 +549,26 @@ impl<'a> State<'a> {
                 self.end();
                 self.pclose();
             }
+            ast::ExprKind::OffsetOf(container, fields) => {
+                // FIXME: This should have its own syntax, distinct from a macro invocation.
+                self.word("offset_of!");
+                self.popen();
+                self.rbox(0, Inconsistent);
+                self.print_type(container);
+                self.word(",");
+                self.space();
+
+                if let Some((&first, rest)) = fields.split_first() {
+                    self.print_ident(first);
+
+                    for &field in rest {
+                        self.word(".");
+                        self.print_ident(field);
+                    }
+                }
+
+                self.end();
+            }
             ast::ExprKind::MacCall(m) => self.print_mac(m),
             ast::ExprKind::Paren(e) => {
                 self.popen();

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -2306,7 +2306,8 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             Rvalue::AddressOf(..)
             | Rvalue::ThreadLocalRef(..)
             | Rvalue::Len(..)
-            | Rvalue::Discriminant(..) => {}
+            | Rvalue::Discriminant(..)
+            | Rvalue::NullaryOp(NullOp::OffsetOf(..), _) => {}
         }
     }
 

compiler/rustc_builtin_macros/messages.ftl

@@ -149,3 +149,6 @@ builtin_macros_format_pos_mismatch = {$n} positional {$n ->
     [one] argument
     *[more] arguments
     } in format string, but {$desc}
+builtin_macros_offset_of_expected_field = expected field
+
+builtin_macros_offset_of_expected_two_args = expected 2 arguments

compiler/rustc_builtin_macros/src/assert/context.rs

@@ -301,6 +301,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             | ExprKind::If(_, _, _)
             | ExprKind::IncludedBytes(..)
             | ExprKind::InlineAsm(_)
+            | ExprKind::OffsetOf(_, _)
             | ExprKind::Let(_, _, _)
             | ExprKind::Lit(_)
             | ExprKind::Loop(_, _, _)

compiler/rustc_builtin_macros/src/lib.rs

@@ -45,6 +45,7 @@ mod format;
 mod format_foreign;
 mod global_allocator;
 mod log_syntax;
+mod offset_of;
 mod source_util;
 mod test;
 mod trace_macros;
@@ -92,6 +93,7 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
         line: source_util::expand_line,
         log_syntax: log_syntax::expand_log_syntax,
         module_path: source_util::expand_mod,
+        offset_of: offset_of::expand_offset_of,
         option_env: env::expand_option_env,
         core_panic: edition_panic::expand_panic,
         std_panic: edition_panic::expand_panic,

compiler/rustc_builtin_macros/src/offset_of.rs

@@ -0,0 +1,99 @@
+use rustc_ast as ast;
+use rustc_ast::ptr::P;
+use rustc_ast::token;
+use rustc_ast::tokenstream::TokenStream;
+use rustc_errors::PResult;
+use rustc_expand::base::{self, *};
+use rustc_macros::Diagnostic;
+use rustc_parse::parser::Parser;
+use rustc_span::{symbol::Ident, Span};
+
+#[derive(Diagnostic)]
+#[diag(builtin_macros_offset_of_expected_field)]
+struct ExpectedField {
+    #[primary_span]
+    span: Span,
+}
+
+#[derive(Diagnostic)]
+#[diag(builtin_macros_offset_of_expected_two_args)]
+struct ExpectedTwoArgs {
+    #[primary_span]
+    span: Span,
+}
+
+fn parse_field<'a>(cx: &ExtCtxt<'a>, p: &mut Parser<'a>) -> PResult<'a, Ident> {
+    let token = p.token.uninterpolate();
+    let field = match token.kind {
+        token::Ident(name, _) => Ident::new(name, token.span),
+        token::Literal(token::Lit { kind: token::Integer, symbol, suffix: None }) => {
+            Ident::new(symbol, token.span)
+        }
+        _ => return Err(cx.create_err(ExpectedField { span: p.token.span })),
+    };
+
+    p.bump();
+
+    Ok(field)
+}
+
+fn parse_args<'a>(
+    cx: &mut ExtCtxt<'a>,
+    sp: Span,
+    tts: TokenStream,
+) -> PResult<'a, (P<ast::Ty>, P<[Ident]>)> {
+    let mut p = cx.new_parser_from_tts(tts);
+
+    let container = p.parse_ty()?;
+
+    p.expect(&token::Comma)?;
+
+    if p.eat(&token::Eof) {
+        return Err(cx.create_err(ExpectedTwoArgs { span: sp }));
+    }
+
+    let mut fields = Vec::new();
+
+    loop {
+        let field = parse_field(cx, &mut p)?;
+        fields.push(field);
+
+        if p.eat(&token::Dot) {
+            continue;
+        }
+
+        p.eat(&token::Comma);
+
+        if !p.eat(&token::Eof) {
+            return Err(cx.create_err(ExpectedTwoArgs { span: sp }));
+        }
+
+        break;
+    }
+
+    Ok((container, fields.into()))
+}
+
+pub fn expand_offset_of<'cx>(
+    cx: &'cx mut ExtCtxt<'_>,
+    sp: Span,
+    tts: TokenStream,
+) -> Box<dyn base::MacResult + 'cx> {
+    match parse_args(cx, sp, tts) {
+        Ok((container, fields)) => {
+            let expr = P(ast::Expr {
+                id: ast::DUMMY_NODE_ID,
+                kind: ast::ExprKind::OffsetOf(container, fields),
+                span: sp,
+                attrs: ast::AttrVec::new(),
+                tokens: None,
+            });
+
+            MacEager::expr(expr)
+        }
+        Err(mut err) => {
+            err.emit();
+            DummyResult::any(sp)
+        }
+    }
+}

compiler/rustc_codegen_cranelift/src/base.rs

@@ -781,12 +781,15 @@ fn codegen_stmt<'tcx>(
                     let operand = operand.load_scalar(fx);
                     lval.write_cvalue(fx, CValue::by_val(operand, box_layout));
                 }
-                Rvalue::NullaryOp(null_op, ty) => {
+                Rvalue::NullaryOp(ref null_op, ty) => {
                     assert!(lval.layout().ty.is_sized(fx.tcx, ParamEnv::reveal_all()));
                     let layout = fx.layout_of(fx.monomorphize(ty));
                     let val = match null_op {
                         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()
+                        }
                     };
                     let val = CValue::const_val(fx, fx.layout_of(fx.tcx.types.usize), val.into());
                     lval.write_cvalue(fx, val);

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -666,13 +666,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 }
             }
 
-            mir::Rvalue::NullaryOp(null_op, ty) => {
+            mir::Rvalue::NullaryOp(ref null_op, ty) => {
                 let ty = self.monomorphize(ty);
                 assert!(bx.cx().type_is_sized(ty));
                 let layout = bx.cx().layout_of(ty);
                 let val = match null_op {
                     mir::NullOp::SizeOf => layout.size.bytes(),
                     mir::NullOp::AlignOf => layout.align.abi.bytes(),
+                    mir::NullOp::OffsetOf(fields) => {
+                        layout.offset_of_subfield(bx.cx(), fields.iter().map(|f| f.index())).bytes()
+                    }
                 };
                 let val = bx.cx().const_usize(val);
                 let tcx = self.cx.tcx();

compiler/rustc_const_eval/src/interpret/step.rs

@@ -280,20 +280,23 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.write_immediate(*val, &dest)?;
             }
 
-            NullaryOp(null_op, ty) => {
+            NullaryOp(ref null_op, ty) => {
                 let ty = self.subst_from_current_frame_and_normalize_erasing_regions(ty)?;
                 let layout = self.layout_of(ty)?;
-                if layout.is_unsized() {
+                if let mir::NullOp::SizeOf | mir::NullOp::AlignOf = null_op && layout.is_unsized() {
                     // FIXME: This should be a span_bug (#80742)
                     self.tcx.sess.delay_span_bug(
                         self.frame().current_span(),
-                        &format!("Nullary MIR operator called for unsized type {}", ty),
+                        &format!("{null_op:?} MIR operator called for unsized type {ty}"),
                     );
                     throw_inval!(SizeOfUnsizedType(ty));
                 }
                 let val = match null_op {
                     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()
+                    }
                 };
                 self.write_scalar(Scalar::from_target_usize(val, self), &dest)?;
             }

compiler/rustc_const_eval/src/transform/check_consts/check.rs

@@ -558,7 +558,7 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
 
             Rvalue::Cast(_, _, _) => {}
 
-            Rvalue::NullaryOp(NullOp::SizeOf | NullOp::AlignOf, _) => {}
+            Rvalue::NullaryOp(NullOp::SizeOf | NullOp::AlignOf | NullOp::OffsetOf(_), _) => {}
             Rvalue::ShallowInitBox(_, _) => {}
 
             Rvalue::UnaryOp(_, operand) => {

compiler/rustc_const_eval/src/transform/promote_consts.rs

@@ -514,6 +514,7 @@ impl<'tcx> Validator<'_, 'tcx> {
             Rvalue::NullaryOp(op, _) => match op {
                 NullOp::SizeOf => {}
                 NullOp::AlignOf => {}
+                NullOp::OffsetOf(_) => {}
             },
 
             Rvalue::ShallowInitBox(_, _) => return Err(Unpromotable),