Skip to content

Commit 0500a9e

Browse files
committed
Auto merge of rust-lang#128004 - folkertdev:naked-fn-asm, r=<try>
codegen `#[naked]` functions using global asm tracking issue: rust-lang#90957 Fixes rust-lang#124375 This implements the approach suggested in the tracking issue: use the existing global assembly infrastructure to emit the body of `#[naked]` functions. The main advantage is that we now have full control over what gets generated, and are no longer dependent on LLVM not sneakily messing with our output (inlining, adding extra instructions, etc). I discussed this approach with `@Amanieu` and while I think the general direction is correct, there is probably a bunch of stuff that needs to change or move around here. I'll leave some inline comments on things that I'm not sure about. Combined with rust-lang#127853, if both accepted, I think that resolves all steps from the tracking issue. r? `@Amanieu` try-job: x86_64-gnu-nopt try-job: x86_64-apple-1
2 parents a224f38 + 10faf24 commit 0500a9e

File tree

20 files changed

+642
-107
lines changed

20 files changed

+642
-107
lines changed

compiler/rustc_attr/src/builtin.rs

+10-1
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,21 @@ pub enum InlineAttr {
4949
Never,
5050
}
5151

52-
#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq, HashStable_Generic)]
52+
#[derive(Copy, Clone, Encodable, Decodable, Debug, PartialEq, Eq, HashStable_Generic)]
5353
pub enum InstructionSetAttr {
5454
ArmA32,
5555
ArmT32,
5656
}
5757

58+
impl InstructionSetAttr {
59+
pub fn as_str(self) -> &'static str {
60+
match self {
61+
Self::ArmA32 => sym::a32.as_str(),
62+
Self::ArmT32 => sym::t32.as_str(),
63+
}
64+
}
65+
}
66+
5867
#[derive(Clone, Encodable, Decodable, Debug, HashStable_Generic)]
5968
pub enum OptimizeAttr {
6069
None,

compiler/rustc_codegen_gcc/src/asm.rs

+7
Original file line numberDiff line numberDiff line change
@@ -867,6 +867,13 @@ impl<'gcc, 'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
867867
template_str.push_str("\n.popsection");
868868
self.context.add_top_level_asm(None, &template_str);
869869
}
870+
871+
fn mangled_name(&self, instance: Instance<'tcx>) -> String {
872+
// TODO(@Amanieu): Additional mangling is needed on
873+
// some targets to add a leading underscore (Mach-O)
874+
// or byte count suffixes (x86 Windows).
875+
self.tcx.symbol_name(instance).name.to_string()
876+
}
870877
}
871878

872879
fn modifier_to_gcc(

compiler/rustc_codegen_llvm/src/asm.rs

+8
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,14 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
442442
);
443443
}
444444
}
445+
446+
fn mangled_name(&self, instance: Instance<'tcx>) -> String {
447+
let llval = self.get_fn(instance);
448+
llvm::build_string(|s| unsafe {
449+
llvm::LLVMRustGetMangledName(llval, s);
450+
})
451+
.expect("symbol is not valid UTF-8")
452+
}
445453
}
446454

447455
pub(crate) fn inline_asm_call<'ll>(

compiler/rustc_codegen_llvm/src/attributes.rs

+3-11
Original file line numberDiff line numberDiff line change
@@ -395,17 +395,9 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
395395
to_add.push(MemoryEffects::None.create_attr(cx.llcx));
396396
}
397397
if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
398-
to_add.push(AttributeKind::Naked.create_attr(cx.llcx));
399-
// HACK(jubilee): "indirect branch tracking" works by attaching prologues to functions.
400-
// And it is a module-level attribute, so the alternative is pulling naked functions into
401-
// new LLVM modules. Otherwise LLVM's "naked" functions come with endbr prefixes per
402-
// https://github.com/rust-lang/rust/issues/98768
403-
to_add.push(AttributeKind::NoCfCheck.create_attr(cx.llcx));
404-
if llvm_util::get_version() < (19, 0, 0) {
405-
// Prior to LLVM 19, branch-target-enforcement was disabled by setting the attribute to
406-
// the string "false". Now it is disabled by absence of the attribute.
407-
to_add.push(llvm::CreateAttrStringValue(cx.llcx, "branch-target-enforcement", "false"));
408-
}
398+
// do nothing; a naked function is converted into an extern function
399+
// and a global assembly block. LLVM's support for naked functions is
400+
// not used.
409401
} else {
410402
// Do not set sanitizer attributes for naked functions.
411403
to_add.extend(sanitize_attrs(cx, codegen_fn_attrs.no_sanitize));

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

+7-4
Original file line numberDiff line numberDiff line change
@@ -542,6 +542,13 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
542542
}
543543
});
544544

545+
// naked function MUST NOT be inlined! This attribute is required for the rust compiler itself,
546+
// but not for the code generation backend because at that point the naked function will just be
547+
// a declaration, with a definition provided in global assembly.
548+
if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
549+
codegen_fn_attrs.inline = InlineAttr::Never;
550+
}
551+
545552
codegen_fn_attrs.optimize = attrs.iter().fold(OptimizeAttr::None, |ia, attr| {
546553
if !attr.has_name(sym::optimize) {
547554
return ia;
@@ -626,10 +633,6 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
626633
}
627634
}
628635

629-
if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
630-
codegen_fn_attrs.inline = InlineAttr::Never;
631-
}
632-
633636
// Weak lang items have the same semantics as "std internal" symbols in the
634637
// sense that they're preserved through all our LTO passes and only
635638
// strippable by the linker.

compiler/rustc_codegen_ssa/src/mir/mod.rs

+6
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ mod coverageinfo;
2020
pub mod debuginfo;
2121
mod intrinsic;
2222
mod locals;
23+
mod naked_asm;
2324
pub mod operand;
2425
pub mod place;
2526
mod rvalue;
@@ -176,6 +177,11 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
176177
let fn_abi = cx.fn_abi_of_instance(instance, ty::List::empty());
177178
debug!("fn_abi: {:?}", fn_abi);
178179

180+
if cx.tcx().codegen_fn_attrs(instance.def_id()).flags.contains(CodegenFnAttrFlags::NAKED) {
181+
crate::mir::naked_asm::codegen_naked_asm::<Bx>(cx, &mir, instance);
182+
return;
183+
}
184+
179185
let debug_context = cx.create_function_debug_context(instance, fn_abi, llfn, mir);
180186

181187
let start_llbb = Bx::append_block(cx, llfn, "start");
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
use rustc_attr::InstructionSetAttr;
2+
use rustc_middle::mir::mono::{Linkage, MonoItem, MonoItemData, Visibility};
3+
use rustc_middle::mir::{Body, InlineAsmOperand};
4+
use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf};
5+
use rustc_middle::ty::{Instance, TyCtxt};
6+
use rustc_middle::{bug, ty};
7+
use rustc_span::sym;
8+
9+
use crate::common;
10+
use crate::traits::{AsmCodegenMethods, BuilderMethods, GlobalAsmOperandRef, MiscCodegenMethods};
11+
12+
pub(crate) fn codegen_naked_asm<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
13+
cx: &'a Bx::CodegenCx,
14+
mir: &Body<'tcx>,
15+
instance: Instance<'tcx>,
16+
) {
17+
let rustc_middle::mir::TerminatorKind::InlineAsm {
18+
asm_macro: _,
19+
template,
20+
ref operands,
21+
options,
22+
line_spans,
23+
targets: _,
24+
unwind: _,
25+
} = mir.basic_blocks.iter().next().unwrap().terminator().kind
26+
else {
27+
bug!("#[naked] functions should always terminate with an asm! block")
28+
};
29+
30+
let operands: Vec<_> =
31+
operands.iter().map(|op| inline_to_global_operand::<Bx>(cx, instance, op)).collect();
32+
33+
let item_data = cx.codegen_unit().items().get(&MonoItem::Fn(instance)).unwrap();
34+
let name = cx.mangled_name(instance);
35+
let (begin, end) = prefix_and_suffix(cx.tcx(), instance, &name, item_data);
36+
37+
let mut template_vec = Vec::new();
38+
template_vec.push(rustc_ast::ast::InlineAsmTemplatePiece::String(begin.into()));
39+
template_vec.extend(template.iter().cloned());
40+
template_vec.push(rustc_ast::ast::InlineAsmTemplatePiece::String(end.into()));
41+
42+
cx.codegen_global_asm(&template_vec, &operands, options, line_spans);
43+
}
44+
45+
fn inline_to_global_operand<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
46+
cx: &'a Bx::CodegenCx,
47+
instance: Instance<'tcx>,
48+
op: &InlineAsmOperand<'tcx>,
49+
) -> GlobalAsmOperandRef<'tcx> {
50+
match op {
51+
InlineAsmOperand::Const { value } => {
52+
let const_value = instance
53+
.instantiate_mir_and_normalize_erasing_regions(
54+
cx.tcx(),
55+
cx.typing_env(),
56+
ty::EarlyBinder::bind(value.const_),
57+
)
58+
.eval(cx.tcx(), cx.typing_env(), value.span)
59+
.expect("erroneous constant missed by mono item collection");
60+
61+
let mono_type = instance.instantiate_mir_and_normalize_erasing_regions(
62+
cx.tcx(),
63+
cx.typing_env(),
64+
ty::EarlyBinder::bind(value.ty()),
65+
);
66+
67+
let string = common::asm_const_to_str(
68+
cx.tcx(),
69+
value.span,
70+
const_value,
71+
cx.layout_of(mono_type),
72+
);
73+
74+
GlobalAsmOperandRef::Const { string }
75+
}
76+
InlineAsmOperand::SymFn { value } => {
77+
let mono_type = instance.instantiate_mir_and_normalize_erasing_regions(
78+
cx.tcx(),
79+
cx.typing_env(),
80+
ty::EarlyBinder::bind(value.ty()),
81+
);
82+
83+
let instance = match mono_type.kind() {
84+
&ty::FnDef(def_id, args) => Instance::new(def_id, args),
85+
_ => bug!("asm sym is not a function"),
86+
};
87+
88+
GlobalAsmOperandRef::SymFn { instance }
89+
}
90+
InlineAsmOperand::SymStatic { def_id } => {
91+
GlobalAsmOperandRef::SymStatic { def_id: *def_id }
92+
}
93+
InlineAsmOperand::In { .. }
94+
| InlineAsmOperand::Out { .. }
95+
| InlineAsmOperand::InOut { .. }
96+
| InlineAsmOperand::Label { .. } => {
97+
bug!("invalid operand type for naked_asm!")
98+
}
99+
}
100+
}
101+
102+
enum AsmBinaryFormat {
103+
Elf,
104+
Macho,
105+
Coff,
106+
}
107+
108+
impl AsmBinaryFormat {
109+
fn from_target(target: &rustc_target::spec::Target) -> Self {
110+
if target.is_like_windows {
111+
Self::Coff
112+
} else if target.is_like_osx {
113+
Self::Macho
114+
} else {
115+
Self::Elf
116+
}
117+
}
118+
}
119+
120+
fn prefix_and_suffix<'tcx>(
121+
tcx: TyCtxt<'tcx>,
122+
instance: Instance<'tcx>,
123+
asm_name: &str,
124+
item_data: &MonoItemData,
125+
) -> (String, String) {
126+
use std::fmt::Write;
127+
128+
let is_arm = tcx.sess.target.arch == "arm";
129+
let is_thumb = tcx.sess.unstable_target_features.contains(&sym::thumb_mode);
130+
131+
let attrs = tcx.codegen_fn_attrs(instance.def_id());
132+
let link_section = attrs.link_section.map(|symbol| symbol.as_str().to_string());
133+
let align = attrs.alignment.map(|a| a.bytes()).unwrap_or(4);
134+
135+
// See https://sourceware.org/binutils/docs/as/ARM-Directives.html for info on these directives.
136+
// In particular, `.arm` can also be written `.code 32` and `.thumb` as `.code 16`.
137+
let (arch_prefix, arch_suffix) = if is_arm {
138+
(
139+
match attrs.instruction_set {
140+
None => match is_thumb {
141+
true => ".thumb\n.thumb_func",
142+
false => ".arm",
143+
},
144+
Some(InstructionSetAttr::ArmT32) => ".thumb\n.thumb_func",
145+
Some(InstructionSetAttr::ArmA32) => ".arm",
146+
},
147+
match is_thumb {
148+
true => ".thumb",
149+
false => ".arm",
150+
},
151+
)
152+
} else {
153+
("", "")
154+
};
155+
156+
let emit_fatal = |msg| tcx.dcx().span_fatal(tcx.def_span(instance.def_id()), msg);
157+
158+
// see https://godbolt.org/z/cPK4sxKor.
159+
// None means the default, which corresponds to internal linkage
160+
let linkage = match item_data.linkage {
161+
Linkage::External => Some(".globl"),
162+
Linkage::LinkOnceAny => Some(".weak"),
163+
Linkage::LinkOnceODR => Some(".weak"),
164+
Linkage::WeakAny => Some(".weak"),
165+
Linkage::WeakODR => Some(".weak"),
166+
Linkage::Internal => None,
167+
Linkage::Private => None,
168+
Linkage::Appending => emit_fatal("Only global variables can have appending linkage!"),
169+
Linkage::Common => emit_fatal("Functions may not have common linkage"),
170+
Linkage::AvailableExternally => {
171+
// this would make the function equal an extern definition
172+
emit_fatal("Functions may not have available_externally linkage")
173+
}
174+
Linkage::ExternalWeak => {
175+
// FIXME: actually this causes a SIGILL in LLVM
176+
emit_fatal("Functions may not have external weak linkage")
177+
}
178+
};
179+
180+
let mut begin = String::new();
181+
let mut end = String::new();
182+
match AsmBinaryFormat::from_target(&tcx.sess.target) {
183+
AsmBinaryFormat::Elf => {
184+
let section = link_section.unwrap_or(format!(".text.{asm_name}"));
185+
186+
let progbits = match is_arm {
187+
true => "%progbits",
188+
false => "@progbits",
189+
};
190+
191+
let function = match is_arm {
192+
true => "%function",
193+
false => "@function",
194+
};
195+
196+
writeln!(begin, ".pushsection {section},\"ax\", {progbits}").unwrap();
197+
writeln!(begin, ".balign {align}").unwrap();
198+
if let Some(linkage) = linkage {
199+
writeln!(begin, "{linkage} {asm_name}").unwrap();
200+
}
201+
if let Visibility::Hidden = item_data.visibility {
202+
writeln!(begin, ".hidden {asm_name}").unwrap();
203+
}
204+
writeln!(begin, ".type {asm_name}, {function}").unwrap();
205+
if !arch_prefix.is_empty() {
206+
writeln!(begin, "{}", arch_prefix).unwrap();
207+
}
208+
writeln!(begin, "{asm_name}:").unwrap();
209+
210+
writeln!(end).unwrap();
211+
writeln!(end, ".size {asm_name}, . - {asm_name}").unwrap();
212+
writeln!(end, ".popsection").unwrap();
213+
if !arch_suffix.is_empty() {
214+
writeln!(end, "{}", arch_suffix).unwrap();
215+
}
216+
}
217+
AsmBinaryFormat::Macho => {
218+
let section = link_section.unwrap_or("__TEXT,__text".to_string());
219+
writeln!(begin, ".pushsection {},regular,pure_instructions", section).unwrap();
220+
writeln!(begin, ".balign {align}").unwrap();
221+
if let Some(linkage) = linkage {
222+
writeln!(begin, "{linkage} {asm_name}").unwrap();
223+
}
224+
if let Visibility::Hidden = item_data.visibility {
225+
writeln!(begin, ".private_extern {asm_name}").unwrap();
226+
}
227+
writeln!(begin, "{asm_name}:").unwrap();
228+
229+
writeln!(end).unwrap();
230+
writeln!(end, ".popsection").unwrap();
231+
if !arch_suffix.is_empty() {
232+
writeln!(end, "{}", arch_suffix).unwrap();
233+
}
234+
}
235+
AsmBinaryFormat::Coff => {
236+
let section = link_section.unwrap_or(format!(".text.{asm_name}"));
237+
writeln!(begin, ".pushsection {},\"xr\"", section).unwrap();
238+
writeln!(begin, ".balign {align}").unwrap();
239+
if let Some(linkage) = linkage {
240+
writeln!(begin, "{linkage} {asm_name}").unwrap();
241+
}
242+
writeln!(begin, ".def {asm_name}").unwrap();
243+
writeln!(begin, ".scl 2").unwrap();
244+
writeln!(begin, ".type 32").unwrap();
245+
writeln!(begin, ".endef {asm_name}").unwrap();
246+
writeln!(begin, "{asm_name}:").unwrap();
247+
248+
writeln!(end).unwrap();
249+
writeln!(end, ".popsection").unwrap();
250+
if !arch_suffix.is_empty() {
251+
writeln!(end, "{}", arch_suffix).unwrap();
252+
}
253+
}
254+
}
255+
256+
(begin, end)
257+
}

compiler/rustc_codegen_ssa/src/mono_item.rs

+8-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use rustc_hir as hir;
2+
use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
23
use rustc_middle::mir::interpret::ErrorHandled;
34
use rustc_middle::mir::mono::{Linkage, MonoItem, Visibility};
45
use rustc_middle::ty::Instance;
@@ -135,7 +136,13 @@ impl<'a, 'tcx: 'a> MonoItemExt<'a, 'tcx> for MonoItem<'tcx> {
135136
cx.predefine_static(def_id, linkage, visibility, symbol_name);
136137
}
137138
MonoItem::Fn(instance) => {
138-
cx.predefine_fn(instance, linkage, visibility, symbol_name);
139+
let attrs = cx.tcx().codegen_fn_attrs(instance.def_id());
140+
141+
if attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
142+
// do not define this function; it will become a global assembly block
143+
} else {
144+
cx.predefine_fn(instance, linkage, visibility, symbol_name);
145+
};
139146
}
140147
MonoItem::GlobalAsm(..) => {}
141148
}

0 commit comments

Comments
 (0)