Add BTF relocation builtins

compiler/rustc_ast/src/ast.rs

@@ -1621,6 +1621,7 @@ impl Expr {
             | ExprKind::While(..)
             | ExprKind::Yield(YieldKind::Postfix(..))
             | ExprKind::DirectConstArg(..)
+            | ExprKind::BtfFieldInfo(..)
             | ExprKind::Err(_)
             | ExprKind::Dummy => prefix_attrs_precedence(&self.attrs),
         }
@@ -1920,6 +1921,9 @@ pub enum ExprKind {
     /// An mGCA `direct_const_arg!()` expression.
     DirectConstArg(Box<Expr>),
 
+    /// A BTF field metadata query.
+    BtfFieldInfo(BtfFieldInfoKind, Box<Ty>, ThinVec<Ident>),
+
     /// Placeholder for an expression that wasn't syntactically well formed in some way.
     Err(ErrorGuaranteed),
 
@@ -2185,6 +2189,27 @@ impl YieldKind {
     }
 }
 
+/// The kind of BTF field metadata query.
+#[derive(Clone, Copy, Encodable, Decodable, Debug, Eq, PartialEq, StableHash, Walkable)]
+pub enum BtfFieldInfoKind {
+    /// Offset of the field.
+    ByteOffset,
+    /// Size of the field.
+    ByteSize,
+    /// Whether the field exists.
+    Exists,
+}
+
+impl BtfFieldInfoKind {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::ByteOffset => "btf_field_byte_offset",
+            Self::ByteSize => "btf_field_byte_size",
+            Self::Exists => "btf_field_exists",
+        }
+    }
+}
+
 /// A literal in a meta item.
 #[derive(Clone, Copy, Encodable, Decodable, Debug, StableHash)]
 pub struct MetaItemLit {

compiler/rustc_ast/src/util/classify.rs

@@ -159,6 +159,7 @@ pub fn leading_labeled_expr(mut expr: &ast::Expr) -> bool {
             | Yield(..)
             | UnsafeBinderCast(..)
             | DirectConstArg(..)
+            | BtfFieldInfo(..)
             | Err(..)
             | Dummy => return false,
         }
@@ -218,7 +219,7 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<TrailingBrace<'_>> {
                 break (mac.args.delim == Delimiter::Brace).then_some(TrailingBrace::MacCall(mac));
             }
 
-            InlineAsm(_) | OffsetOf(_, _) | IncludedBytes(_) | FormatArgs(_) => {
+            InlineAsm(_) | OffsetOf(_, _) | IncludedBytes(_) | FormatArgs(_) | BtfFieldInfo(..) => {
                 // These should have been denied pre-expansion.
                 break None;
             }

compiler/rustc_ast/src/visit.rs

@@ -493,6 +493,7 @@ macro_rules! common_visitor_and_walkers {
             YieldKind,
             EiiDecl,
             EiiImpl,
+            BtfFieldInfoKind,
         );
 
         /// Each method of this trait is a hook to be potentially
@@ -1074,6 +1075,8 @@ macro_rules! common_visitor_and_walkers {
                     visit_visitable!($($mut)? vis, kind, expr, ty),
                 ExprKind::DirectConstArg(expr) =>
                     visit_visitable!($($mut)? vis, expr),
+                ExprKind::BtfFieldInfo(kind, container, fields) =>
+                    visit_visitable!($($mut)? vis, kind, container, fields),
                 ExprKind::Err(_guar) => {}
                 ExprKind::Dummy => {}
             }

compiler/rustc_ast_lowering/src/expr.rs

@@ -483,6 +483,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 let e = self.emit_bad_direct_const_arg(e.span, expr, "expression");
                 hir::ExprKind::Err(e)
             }
+
+            ExprKind::BtfFieldInfo(kind, container, fields) => hir::ExprKind::BtfFieldInfo(
+                *kind,
+                self.lower_ty_alloc(
+                    container,
+                    ImplTraitContext::Disallowed(ImplTraitPosition::BtfFieldInfo),
+                ),
+                self.arena.alloc_from_iter(fields.iter().map(|&ident| self.lower_ident(ident))),
+            ),
         };
 
         hir::Expr { hir_id: expr_hir_id, kind, span }

compiler/rustc_ast_lowering/src/lib.rs

@@ -437,6 +437,7 @@ enum ImplTraitPosition {
     Cast,
     ImplSelf,
     OffsetOf,
+    BtfFieldInfo,
 }
 
 impl std::fmt::Display for ImplTraitPosition {
@@ -463,6 +464,7 @@ impl std::fmt::Display for ImplTraitPosition {
             ImplTraitPosition::Cast => "cast expression types",
             ImplTraitPosition::ImplSelf => "impl headers",
             ImplTraitPosition::OffsetOf => "`offset_of!` parameters",
+            ImplTraitPosition::BtfFieldInfo => "BTF field info query parameters",
         };
 
         write!(f, "{name}")

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

@@ -888,6 +888,24 @@ impl<'a> State<'a> {
                 self.print_expr(expr, FixupContext::default());
                 self.pclose()
             }
+            ast::ExprKind::BtfFieldInfo(kind, container, fields) => {
+                self.word("builtin # ");
+                self.word(kind.as_str());
+                self.popen();
+                let ib = self.ibox(0);
+                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(ib);
+                self.pclose();
+            }
         }
 
         self.ann.post(self, AnnNode::Expr(expr));

compiler/rustc_attr_ir/src/data_structures.rs

@@ -1002,6 +1002,9 @@ pub enum AttributeKind {
     /// Represents `#[automatically_derived]`
     AutomaticallyDerived,
 
+    /// Represents `#[btf_relocatable]`.
+    BtfRelocatable(Span),
+
     /// Represents the trace attribute of `#[cfg_attr]`
     CfgAttrTrace(ThinVec<(CfgEntry, Span)>),
 

compiler/rustc_attr_ir/src/encode_cross_crate.rs

@@ -21,6 +21,7 @@ impl AttributeKind {
             AllowInternalUnsafe(..) => Yes,
             AllowInternalUnstable(..) => Yes,
             AutomaticallyDerived => Yes,
+            BtfRelocatable(..) => Yes,
             CfgAttrTrace(..) => Yes,
             CfgTrace(..) => Yes,
             CfiEncoding { .. } => Yes,

compiler/rustc_attr_parsing/src/attributes/btf_relocatable.rs

@@ -0,0 +1,22 @@
+use rustc_feature::AttributeStability;
+use rustc_target::spec::Arch;
+
+use super::prelude::*;
+use crate::diagnostics::BtfRelocatableOnNonBpfArch;
+
+pub(crate) struct BtfRelocatableParser;
+
+impl NoArgsAttributeParser for BtfRelocatableParser {
+    const PATH: &[Symbol] = &[sym::btf_relocatable];
+    const ALLOWED_TARGETS: AllowedTargets<'_> =
+        AllowedTargets::AllowList(&[Allow(Target::Struct), Allow(Target::Union)]);
+    const STABILITY: AttributeStability = unstable!(btf_relocations);
+    const CREATE: fn(Span) -> AttributeKind = AttributeKind::BtfRelocatable;
+
+    fn finalize_check(cx: &FinalizeCheckContext<'_, '_>, attr_span: Span) {
+        // `#[btf_relocatable]` may be only applied on BPF architecture.
+        if cx.shared.cx.sess().target.arch != Arch::Bpf {
+            cx.shared.cx.dcx().emit_err(BtfRelocatableOnNonBpfArch { span: attr_span });
+        }
+    }
+}

compiler/rustc_attr_parsing/src/attributes/mod.rs

@@ -38,6 +38,7 @@ mod prelude;
 pub(crate) mod allow_unstable;
 pub(crate) mod autodiff;
 pub(crate) mod body;
+pub(crate) mod btf_relocatable;
 pub(crate) mod cfg;
 pub(crate) mod cfg_select;
 pub(crate) mod cfi_encoding;

compiler/rustc_attr_parsing/src/context.rs

@@ -23,6 +23,7 @@ use rustc_span::{ErrorGuaranteed, Ident, Span, Symbol};
 use crate::attributes::allow_unstable::*;
 use crate::attributes::autodiff::*;
 use crate::attributes::body::*;
+use crate::attributes::btf_relocatable::*;
 use crate::attributes::cfi_encoding::*;
 use crate::attributes::codegen_attrs::*;
 use crate::attributes::confusables::*;
@@ -261,6 +262,7 @@ attribute_parsers!(
         Single<WindowsSubsystemParser>,
         Single<WithoutArgs<AllowInternalUnsafeParser>>,
         Single<WithoutArgs<AutomaticallyDerivedParser>>,
+        Single<WithoutArgs<BtfRelocatableParser>>,
         Single<WithoutArgs<ColdParser>>,
         Single<WithoutArgs<CompilerBuiltinsParser>>,
         Single<WithoutArgs<ComptimeParser>>,

compiler/rustc_attr_parsing/src/diagnostics.rs

@@ -2035,3 +2035,10 @@ pub(crate) struct UnusedDuplicate {
     )]
     pub warning: bool,
 }
+
+#[derive(Diagnostic)]
+#[diag("the `btf_relocatable` attribute can only be used on BPF architecture")]
+pub(crate) struct BtfRelocatableOnNonBpfArch {
+    #[primary_span]
+    pub span: Span,
+}

compiler/rustc_borrowck/src/lib.rs

@@ -1580,7 +1580,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
                 );
             }
 
-            Rvalue::ThreadLocalRef(_) => {}
+            Rvalue::ThreadLocalRef(_) | Rvalue::BtfFieldInfo { .. } => {}
 
             Rvalue::Use(operand, _)
             | Rvalue::Repeat(operand, _)

compiler/rustc_borrowck/src/polonius/legacy/loan_invalidations.rs

@@ -298,7 +298,7 @@ impl<'a, 'tcx> LoanInvalidationsGenerator<'a, 'tcx> {
                 self.access_place(location, place, access_kind, LocalMutationIsAllowed::No);
             }
 
-            Rvalue::ThreadLocalRef(_) => {}
+            Rvalue::ThreadLocalRef(_) | Rvalue::BtfFieldInfo { .. } => {}
 
             Rvalue::Use(operand, _)
             | Rvalue::Repeat(operand, _)

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1702,7 +1702,8 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
             | Rvalue::BinaryOp(..)
             | Rvalue::RawPtr(..)
             | Rvalue::ThreadLocalRef(..)
-            | Rvalue::Discriminant(..) => {}
+            | Rvalue::Discriminant(..)
+            | Rvalue::BtfFieldInfo { .. } => {}
         }
     }
 
@@ -2280,7 +2281,8 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             | Rvalue::CopyForDeref(..)
             | Rvalue::UnaryOp(..)
             | Rvalue::Discriminant(..)
-            | Rvalue::WrapUnsafeBinder(..) => None,
+            | Rvalue::WrapUnsafeBinder(..)
+            | Rvalue::BtfFieldInfo { .. } => None,
 
             Rvalue::Aggregate(aggregate, _) => match **aggregate {
                 AggregateKind::Adt(_, _, _, user_ty, _) => user_ty,

compiler/rustc_builtin_macros/src/assert/context.rs

@@ -324,7 +324,8 @@ impl<'cx, 'a> Context<'cx, 'a> {
             | ExprKind::Become(_)
             | ExprKind::Yield(_)
             | ExprKind::DirectConstArg(_)
-            | ExprKind::UnsafeBinderCast(..) => {}
+            | ExprKind::UnsafeBinderCast(..)
+            | ExprKind::BtfFieldInfo(..) => {}
         }
     }
 

compiler/rustc_codegen_llvm/src/builder.rs

@@ -22,6 +22,7 @@ use rustc_middle::ty::layout::{
     TyAndLayout,
 };
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
+use rustc_middle::{bug, mir};
 use rustc_sanitizers::{cfi, kcfi};
 use rustc_session::config::OptLevel;
 use rustc_span::Span;
@@ -34,6 +35,7 @@ use crate::abi::FnAbiLlvmExt;
 use crate::attributes;
 use crate::common::Funclet;
 use crate::context::{CodegenCx, FullCx, GenericCx, SCx};
+use crate::debuginfo::metadata::type_di_node;
 use crate::llvm::{
     self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, FromGeneric, GEPNoWrapFlags, Metadata, TRUE,
     ToLlvmBool, Type, Value,
@@ -1545,6 +1547,83 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         let cold_inline = llvm::AttributeKind::Cold.create_attr(self.llcx);
         attributes::apply_to_callsite(llret, llvm::AttributePlace::Function, &[cold_inline]);
     }
+
+    // BTF relocations
+    fn btf_field_info(
+        &mut self,
+        base: &'ll Value,
+        path: impl Iterator<Item = mir::BtfFieldStep<'tcx>>,
+        kind: mir::BtfFieldInfoKind,
+    ) -> &'ll Value {
+        fn llvm_struct_field_index<'ll, 'tcx>(
+            bx: &Builder<'_, 'll, 'tcx>,
+            layout: TyAndLayout<'tcx>,
+            field_index: usize,
+        ) -> usize {
+            let mut llvm_index = 0;
+            let mut offset = Size::ZERO;
+
+            for i in layout.fields.index_by_increasing_offset() {
+                let target_offset = layout.fields.offset(i as usize);
+                if target_offset != offset {
+                    llvm_index += 1;
+                }
+                if i as usize == field_index {
+                    return llvm_index;
+                }
+
+                let field = layout.field(bx.cx(), i);
+                llvm_index += 1;
+                offset = target_offset + field.size;
+            }
+
+            bug!("field index {field_index} not found in layout {layout:#?}")
+        }
+
+        let layout_cx = ty::layout::LayoutCx::new(self.tcx, self.typing_env());
+        let mut field_ptr = base;
+        for step in path {
+            let layout = self.layout_of(step.container_ty);
+            let layout = layout.for_variant(&layout_cx, step.variant);
+            field_ptr = match step.container_ty.kind() {
+                ty::Adt(adt, _) if adt.is_union() => {
+                    let dbg_info: &'ll Metadata = type_di_node(self.cx, step.container_ty);
+                    unsafe {
+                        llvm::LLVMRustBuildPreserveUnionAccessIndex(
+                            self.llbuilder,
+                            field_ptr,
+                            step.field.index() as c_uint,
+                            Some(dbg_info),
+                        )
+                    }
+                }
+                ty::Adt(..) | ty::Tuple(..) => {
+                    let llvm_index = llvm_struct_field_index(self, layout, step.field.index());
+                    let dbg_info: &'ll Metadata = type_di_node(self.cx, step.container_ty);
+                    unsafe {
+                        llvm::LLVMRustBuildPreserveStructAccessIndex(
+                            self.llbuilder,
+                            self.cx().backend_type(layout),
+                            field_ptr,
+                            llvm_index as c_uint,
+                            step.field.index() as c_uint,
+                            Some(dbg_info),
+                        )
+                    }
+                }
+                _ => bug!(
+                    "BTF field info queries are unsupported for container: {:?}",
+                    step.container_ty
+                ),
+            };
+        }
+
+        self.call_intrinsic(
+            "llvm.bpf.preserve.field.info",
+            &[self.val_ty(field_ptr)],
+            &[field_ptr, self.const_u64(kind.as_u64())],
+        )
+    }
 }
 
 impl<'ll> StaticBuilderMethods for Builder<'_, 'll, '_> {

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1711,6 +1711,22 @@ unsafe extern "C" {
         NumBundles: c_uint,
         Name: *const c_char,
     ) -> &'a Value;
+
+    // BTF relocations
+    pub(crate) fn LLVMRustBuildPreserveUnionAccessIndex<'a>(
+        B: &Builder<'a>,
+        Base: &'a Value,
+        FieldIndex: c_uint,
+        DbgInfo: Option<&'a Metadata>,
+    ) -> &'a Value;
+    pub(crate) fn LLVMRustBuildPreserveStructAccessIndex<'a>(
+        B: &Builder<'a>,
+        ElTy: &'a Type,
+        Base: &'a Value,
+        Index: c_uint,
+        FieldIndex: c_uint,
+        DbgInfo: Option<&'a Metadata>,
+    ) -> &'a Value;
 }
 
 // FFI bindings for `DIBuilder` functions in the LLVM-C API.