Implement asm_const_ptr for global_asm and naked_asm

nbdd0121 · nbdd0121 · commit ffba8e8d1a8f · 2025-03-17T17:55:53.000Z
diff --git a/compiler/rustc_codegen_llvm/src/asm.rs b/compiler/rustc_codegen_llvm/src/asm.rs
@@ -7,7 +7,7 @@ use rustc_codegen_ssa::traits::*;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::layout::TyAndLayout;
-use rustc_middle::{bug, span_bug};
+use rustc_middle::{bug, mir, span_bug};
 use rustc_span::{Pos, Span, Symbol, sym};
 use rustc_target::asm::*;
 use smallvec::SmallVec;
@@ -396,6 +396,111 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
         let intel_syntax = matches!(asm_arch, InlineAsmArch::X86 | InlineAsmArch::X86_64)
             && !options.contains(InlineAsmOptions::ATT_SYNTAX);
 
+        // Convert all operands to string interpolations
+        let converted_operands = operands
+            .iter()
+            .enumerate()
+            .map(|(operand_idx, operand)| {
+                match *operand {
+                    GlobalAsmOperandRef::Interpolate { ref string } => {
+                        // Const operands get injected directly into the
+                        // template. Note that we don't need to escape $
+                        // here unlike normal inline assembly.
+                        string.to_owned()
+                    }
+                    GlobalAsmOperandRef::ConstPointer { value, instance } => {
+                        let (prov, offset) = value.into_parts();
+                        let global_alloc = self.tcx.global_alloc(prov.alloc_id());
+                        let llval = 'llval: {
+                            let alloc = match global_alloc {
+                                mir::interpret::GlobalAlloc::Function { instance } => {
+                                    break 'llval self.get_fn(instance);
+                                }
+                                mir::interpret::GlobalAlloc::VTable(ty, dyn_ty) => self
+                                    .tcx
+                                    .global_alloc(self.tcx.vtable_allocation((
+                                        ty,
+                                        dyn_ty.principal().map(|principal| {
+                                            self.tcx
+                                                .instantiate_bound_regions_with_erased(principal)
+                                        }),
+                                    )))
+                                    .unwrap_memory(),
+                                mir::interpret::GlobalAlloc::Static(def_id) => {
+                                    break 'llval self
+                                        .renamed_statics
+                                        .borrow()
+                                        .get(&def_id)
+                                        .copied()
+                                        .unwrap_or_else(|| self.get_static(def_id));
+                                }
+                                mir::interpret::GlobalAlloc::Memory(alloc) => alloc,
+                            };
+
+                            // For ZSTs directly codegen an aligned pointer.
+                            if alloc.inner().len() == 0 {
+                                assert_eq!(offset.bytes(), 0);
+                                return format!("{}", alloc.inner().align.bytes());
+                            }
+
+                            let asm_name = self.tcx.symbol_name(instance);
+                            let sym_name = format!("{asm_name}.{operand_idx}");
+
+                            let init = crate::consts::const_alloc_to_llvm(
+                                self, alloc, /*static*/ false,
+                            );
+                            let alloc = alloc.inner();
+                            let g = self.static_addr_of_mut(init, alloc.align, None);
+                            if alloc.mutability.is_not() {
+                                // NB: we can't use `static_addr_of_impl` here to avoid sharing
+                                // the global, as we need to set name and linkage.
+                                unsafe { llvm::LLVMSetGlobalConstant(g, llvm::True) };
+                            }
+
+                            llvm::set_value_name(g, sym_name.as_bytes());
+
+                            // `static_addr_of_mut` gives us a private global which can't be
+                            // used by global asm. Update it to a hidden internal global instead.
+                            llvm::set_linkage(g, llvm::Linkage::InternalLinkage);
+                            llvm::set_visibility(g, llvm::Visibility::Hidden);
+                            g
+                        };
+                        self.add_compiler_used_global(llval);
+                        let symbol = llvm::build_string(|s| unsafe {
+                            llvm::LLVMRustGetMangledName(llval, s);
+                        })
+                        .expect("symbol is not valid UTF-8");
+
+                        let offset = offset.bytes();
+                        if offset != 0 { format!("{symbol}+{offset}") } else { symbol }
+                    }
+                    GlobalAsmOperandRef::SymFn { instance } => {
+                        let llval = self.get_fn(instance);
+                        self.add_compiler_used_global(llval);
+                        let symbol = llvm::build_string(|s| unsafe {
+                            llvm::LLVMRustGetMangledName(llval, s);
+                        })
+                        .expect("symbol is not valid UTF-8");
+                        symbol
+                    }
+                    GlobalAsmOperandRef::SymStatic { def_id } => {
+                        let llval = self
+                            .renamed_statics
+                            .borrow()
+                            .get(&def_id)
+                            .copied()
+                            .unwrap_or_else(|| self.get_static(def_id));
+                        self.add_compiler_used_global(llval);
+                        let symbol = llvm::build_string(|s| unsafe {
+                            llvm::LLVMRustGetMangledName(llval, s);
+                        })
+                        .expect("symbol is not valid UTF-8");
+                        symbol
+                    }
+                }
+            })
+            .collect::<Vec<_>>();
+
         // Build the template string
         let mut template_str = String::new();
         if intel_syntax {
@@ -405,37 +510,7 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
             match *piece {
                 InlineAsmTemplatePiece::String(ref s) => template_str.push_str(s),
                 InlineAsmTemplatePiece::Placeholder { operand_idx, modifier: _, span: _ } => {
-                    match operands[operand_idx] {
-                        GlobalAsmOperandRef::Interpolate { ref string } => {
-                            // Const operands get injected directly into the
-                            // template. Note that we don't need to escape $
-                            // here unlike normal inline assembly.
-                            template_str.push_str(string);
-                        }
-                        GlobalAsmOperandRef::SymFn { instance } => {
-                            let llval = self.get_fn(instance);
-                            self.add_compiler_used_global(llval);
-                            let symbol = llvm::build_string(|s| unsafe {
-                                llvm::LLVMRustGetMangledName(llval, s);
-                            })
-                            .expect("symbol is not valid UTF-8");
-                            template_str.push_str(&symbol);
-                        }
-                        GlobalAsmOperandRef::SymStatic { def_id } => {
-                            let llval = self
-                                .renamed_statics
-                                .borrow()
-                                .get(&def_id)
-                                .copied()
-                                .unwrap_or_else(|| self.get_static(def_id));
-                            self.add_compiler_used_global(llval);
-                            let symbol = llvm::build_string(|s| unsafe {
-                                llvm::LLVMRustGetMangledName(llval, s);
-                            })
-                            .expect("symbol is not valid UTF-8");
-                            template_str.push_str(&symbol);
-                        }
-                    }
+                    template_str.push_str(&converted_operands[operand_idx])
                 }
             }
         }
diff --git a/compiler/rustc_codegen_ssa/src/common.rs b/compiler/rustc_codegen_ssa/src/common.rs
@@ -166,7 +166,7 @@ pub fn asm_const_to_str<'tcx>(
     };
     let value = scalar.assert_scalar_int().to_bits(ty_and_layout.size);
     match ty_and_layout.ty.kind() {
-        ty::Uint(_) => value.to_string(),
+        ty::Uint(_) | ty::RawPtr(..) | ty::Ref(..) => value.to_string(),
         ty::Int(int_ty) => match int_ty.normalize(tcx.sess.target.pointer_width) {
             ty::IntTy::I8 => (value as i8).to_string(),
             ty::IntTy::I16 => (value as i16).to_string(),
diff --git a/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs b/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
@@ -2,7 +2,7 @@ use rustc_abi::{BackendRepr, Float, Integer, Primitive, RegKind};
 use rustc_attr_parsing::InstructionSetAttr;
 use rustc_hir::def_id::DefId;
 use rustc_middle::mir::mono::{Linkage, MonoItem, MonoItemData, Visibility};
-use rustc_middle::mir::{Body, InlineAsmOperand};
+use rustc_middle::mir::{self, Body, InlineAsmOperand};
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
 use rustc_middle::ty::{Instance, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug, ty};
@@ -69,14 +69,28 @@ fn inline_to_global_operand<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
                 ty::EarlyBinder::bind(value.ty()),
             );
 
-            let string = common::asm_const_to_str(
-                cx.tcx(),
-                value.span,
-                const_value,
-                cx.layout_of(mono_type),
-            );
-
-            GlobalAsmOperandRef::Interpolate { string }
+            let mir::ConstValue::Scalar(scalar) = const_value else {
+                span_bug!(
+                    value.span,
+                    "expected Scalar for promoted asm const, but got {:#?}",
+                    const_value
+                )
+            };
+            match scalar {
+                mir::interpret::Scalar::Int(_) => {
+                    let string = common::asm_const_to_str(
+                        cx.tcx(),
+                        value.span,
+                        const_value,
+                        cx.layout_of(mono_type),
+                    );
+                    GlobalAsmOperandRef::Interpolate { string }
+                }
+                mir::interpret::Scalar::Ptr(value, _) => GlobalAsmOperandRef::ConstPointer {
+                    value,
+                    instance: Instance::mono(cx.tcx(), instance.def_id()),
+                },
+            }
         }
         InlineAsmOperand::SymFn { value } => {
             let mono_type = instance.instantiate_mir_and_normalize_erasing_regions(
diff --git a/compiler/rustc_codegen_ssa/src/mono_item.rs b/compiler/rustc_codegen_ssa/src/mono_item.rs
@@ -1,10 +1,10 @@
 use rustc_hir as hir;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
-use rustc_middle::mir::interpret::ErrorHandled;
+use rustc_middle::mir::interpret::{ErrorHandled, Scalar};
 use rustc_middle::mir::mono::{Linkage, MonoItem, Visibility};
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
-use rustc_middle::{span_bug, ty};
+use rustc_middle::{mir, span_bug, ty};
 use tracing::debug;
 
 use crate::traits::*;
@@ -48,13 +48,23 @@ impl<'a, 'tcx: 'a> MonoItemExt<'a, 'tcx> for MonoItem<'tcx> {
                                             .tcx()
                                             .typeck_body(anon_const.body)
                                             .node_type(anon_const.hir_id);
-                                        let string = common::asm_const_to_str(
-                                            cx.tcx(),
-                                            *op_sp,
-                                            const_value,
-                                            cx.layout_of(ty),
-                                        );
-                                        GlobalAsmOperandRef::Interpolate { string }
+                                        let mir::ConstValue::Scalar(scalar) = const_value else {
+                                            span_bug!(*op_sp, "expected Scalar for promoted asm const, but got {:#?}", const_value)
+                                        };
+                                        match scalar {
+                                            Scalar::Int(_) => {
+                                                let string = common::asm_const_to_str(
+                                                    cx.tcx(),
+                                                    *op_sp,
+                                                    const_value,
+                                                    cx.layout_of(ty),
+                                                );
+                                                GlobalAsmOperandRef::Interpolate { string }
+                                            }
+                                            Scalar::Ptr(value, _) => {
+                                                GlobalAsmOperandRef::ConstPointer { value, instance: Instance::mono(cx.tcx(), item_id.owner_id.to_def_id()) }
+                                            }
+                                        }
                                     }
                                     Err(ErrorHandled::Reported { .. }) => {
                                         // An error has already been reported and
diff --git a/compiler/rustc_codegen_ssa/src/traits/asm.rs b/compiler/rustc_codegen_ssa/src/traits/asm.rs
@@ -1,5 +1,6 @@
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_hir::def_id::DefId;
+use rustc_middle::mir;
 use rustc_middle::ty::Instance;
 use rustc_span::Span;
 use rustc_target::asm::InlineAsmRegOrRegClass;
@@ -44,9 +45,20 @@ pub enum InlineAsmOperandRef<'tcx, B: BackendTypes + ?Sized> {
 
 #[derive(Debug)]
 pub enum GlobalAsmOperandRef<'tcx> {
-    Interpolate { string: String },
-    SymFn { instance: Instance<'tcx> },
-    SymStatic { def_id: DefId },
+    Interpolate {
+        string: String,
+    },
+    ConstPointer {
+        value: mir::interpret::Pointer,
+        /// Instance that instantiates this const operand.
+        instance: Instance<'tcx>,
+    },
+    SymFn {
+        instance: Instance<'tcx>,
+    },
+    SymStatic {
+        def_id: DefId,
+    },
 }
 
 pub trait AsmBuilderMethods<'tcx>: BackendTypes {
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
@@ -490,8 +490,7 @@ fn collect_items_rec<'tcx>(
                                 Err(ErrorHandled::TooGeneric(..)) => {
                                     span_bug!(*op_sp, "asm const cannot be resolved; too generic")
                                 }
-                                Err(err @ ErrorHandled::Reported(..)) => {
-                                    err.emit_note(tcx);
+                                Err(ErrorHandled::Reported(..)) => {
                                     continue;
                                 }
                             }
diff --git a/tests/assembly/asm/global_asm.rs b/tests/assembly/asm/global_asm.rs
@@ -5,6 +5,7 @@
 //@ compile-flags: -C symbol-mangling-version=v0
 
 #![crate_type = "rlib"]
+#![feature(asm_const_ptr)]
 
 use std::arch::global_asm;
 
@@ -26,6 +27,10 @@ global_asm!("call {}", sym my_func);
 global_asm!("lea rax, [rip + {}]", sym MY_STATIC);
 // CHECK: call _RNvC[[CRATE_IDENT:[a-zA-Z0-9]{12}]]_10global_asm6foobar
 global_asm!("call {}", sym foobar);
+// CHECK: lea rax, [rip + _RNSC[[CRATE_IDENT]]_10global_asms4_10global_asm.0]
+global_asm!("lea rax, [rip + {}]", const &1);
+// CHECK: lea rax, [rip + _RNSC[[CRATE_IDENT]]_10global_asms5_10global_asm.0+4]
+global_asm!("lea rax, [rip + {}]", const &[1; 2][1]);
 // CHECK: _RNvC[[CRATE_IDENT]]_10global_asm6foobar:
 fn foobar() {
     loop {}
diff --git a/tests/assembly/asm/x86-types.rs b/tests/assembly/asm/x86-types.rs
@@ -9,7 +9,7 @@
 //@ compile-flags: -C target-feature=+avx512bw
 //@ compile-flags: -Zmerge-functions=disabled
 
-#![feature(no_core, repr_simd, f16, f128)]
+#![feature(no_core, repr_simd, f16, f128, asm_const_ptr)]
 #![crate_type = "rlib"]
 #![no_core]
 #![allow(asm_sub_register, non_camel_case_types)]
@@ -92,6 +92,18 @@ pub unsafe fn sym_fn() {
     asm!("call {}", sym extern_func);
 }
 
+// NOTE: this only works for x64, as this test is compiled with PIC,
+// and on x86 PIC symbol can't be constant.
+// x86_64-LABEL: const_ptr:
+// x86_64: #APP
+// x86_64: mov al, byte ptr [{{.*}}anon{{.*}}]
+// x86_64: #NO_APP
+#[cfg(x86_64)]
+#[no_mangle]
+pub unsafe fn const_ptr() {
+    asm!("mov al, byte ptr [{}]", const &1);
+}
+
 // CHECK-LABEL: sym_static:
 // CHECK: #APP
 // CHECK: mov al, byte ptr [extern_static]
diff --git a/tests/ui/asm/invalid-const-operand.rs b/tests/ui/asm/invalid-const-operand.rs
@@ -17,16 +17,21 @@ global_asm!("{}", const 0 as *mut u8);
 
 fn test1() {
     unsafe {
-        // Const operands must be integers and must be constants.
+        // Const operands must be integers or thin pointers
 
         asm!("{}", const 0);
         asm!("{}", const 0i32);
         asm!("{}", const 0i128);
         asm!("{}", const 0f32);
         //~^ ERROR invalid type for `const` operand
         asm!("{}", const 0 as *mut u8);
+        asm!("{}", const &0);
         asm!("{}", const b"Foo".as_slice());
         //~^ ERROR invalid type for `const` operand
+
+        asm!("{}", const test1 as fn());
+        asm!("{}", const test1);
+        //~^ ERROR invalid type for `const` operand
     }
 }
 
diff --git a/tests/ui/asm/invalid-const-operand.stderr b/tests/ui/asm/invalid-const-operand.stderr
diff --git a/tests/ui/feature-gates/feature-gate-asm_const_ptr.rs b/tests/ui/feature-gates/feature-gate-asm_const_ptr.rs
diff --git a/tests/ui/feature-gates/feature-gate-asm_const_ptr.stderr b/tests/ui/feature-gates/feature-gate-asm_const_ptr.stderr

Original file line number	Diff line number	Diff line change
`@@ -490,8 +490,7 @@ fn collect_items_rec<'tcx>(`
`490`	`490`	`Err(ErrorHandled::TooGeneric(..)) => {`
`491`	`491`	`span_bug!(*op_sp, "asm const cannot be resolved; too generic")`
`492`	`492`	`}`
`493`		`- Err(err @ ErrorHandled::Reported(..)) => {`
`494`		`- err.emit_note(tcx);`
	`493`	`+ Err(ErrorHandled::Reported(..)) => {`
`495`	`494`	`continue;`
`496`	`495`	`}`
`497`	`496`	`}`