naked functions: respect function-sections

compiler/rustc_codegen_ssa/src/mir/naked_asm.rs

@@ -8,7 +8,7 @@ use rustc_middle::ty::{Instance, Ty, TyCtxt, TypeVisitableExt};
 use rustc_middle::{bug, ty};
 use rustc_span::sym;
 use rustc_target::callconv::{ArgAbi, FnAbi, PassMode};
-use rustc_target::spec::{Arch, BinaryFormat};
+use rustc_target::spec::{Arch, BinaryFormat, Env, Os};
 
 use crate::common;
 use crate::mir::AsmCodegenMethods;
@@ -128,6 +128,8 @@ fn prefix_and_suffix<'tcx>(
 
     let is_arm = tcx.sess.target.arch == Arch::Arm;
     let is_thumb = tcx.sess.unstable_target_features.contains(&sym::thumb_mode);
+    let function_sections =
+        tcx.sess.opts.unstable_opts.function_sections.unwrap_or(tcx.sess.target.function_sections);
 
     // If we're compiling the compiler-builtins crate, e.g., the equivalent of
     // compiler-rt, then we want to implicitly compile everything with hidden
@@ -218,8 +220,6 @@ fn prefix_and_suffix<'tcx>(
     let mut end = String::new();
     match asm_binary_format {
         BinaryFormat::Elf => {
-            let section = link_section.unwrap_or_else(|| format!(".text.{asm_name}"));
-
             let progbits = match is_arm {
                 true => "%progbits",
                 false => "@progbits",
@@ -230,7 +230,13 @@ fn prefix_and_suffix<'tcx>(
                 false => "@function",
             };
 
-            writeln!(begin, ".pushsection {section},\"ax\", {progbits}").unwrap();
+            if let Some(section) = &link_section {
+                writeln!(begin, ".pushsection {section},\"ax\", {progbits}").unwrap();
+            } else if function_sections {
+                writeln!(begin, ".pushsection .text.{asm_name},\"ax\", {progbits}").unwrap();
+            } else {
+                writeln!(begin, ".text").unwrap();
+            }
             writeln!(begin, ".balign {align_bytes}").unwrap();
             write_linkage(&mut begin).unwrap();
             match visibility {
@@ -249,14 +255,22 @@ fn prefix_and_suffix<'tcx>(
             // pattern match on assembly generated by LLVM.
             writeln!(end, ".Lfunc_end_{asm_name}:").unwrap();
             writeln!(end, ".size {asm_name}, . - {asm_name}").unwrap();
-            writeln!(end, ".popsection").unwrap();
+            if link_section.is_some() || function_sections {
+                writeln!(end, ".popsection").unwrap();
+            }
             if !arch_suffix.is_empty() {
                 writeln!(end, "{}", arch_suffix).unwrap();
             }
         }
         BinaryFormat::MachO => {
-            let section = link_section.unwrap_or_else(|| "__TEXT,__text".to_string());
-            writeln!(begin, ".pushsection {},regular,pure_instructions", section).unwrap();
+            // NOTE: LLVM ignores `-Zfunction-sections` on macos. Instead the Mach-O symbol
+            // subsection splitting feature is used, which can be enabled with the
+            // `.subsections_via_symbols` global directive. LLVM already enables this directive.
+            if let Some(section) = &link_section {
+                writeln!(begin, ".pushsection {section},regular,pure_instructions").unwrap();
+            } else {
+                writeln!(begin, ".section __TEXT,__text,regular,pure_instructions").unwrap();
+            }
             writeln!(begin, ".balign {align_bytes}").unwrap();
             write_linkage(&mut begin).unwrap();
             match visibility {
@@ -267,7 +281,9 @@ fn prefix_and_suffix<'tcx>(
 
             writeln!(end).unwrap();
             writeln!(end, ".Lfunc_end_{asm_name}:").unwrap();
-            writeln!(end, ".popsection").unwrap();
+            if link_section.is_some() {
+                writeln!(end, ".popsection").unwrap();
+            }
             if !arch_suffix.is_empty() {
                 writeln!(end, "{}", arch_suffix).unwrap();
             }
@@ -278,8 +294,36 @@ fn prefix_and_suffix<'tcx>(
             writeln!(begin, ".type 32").unwrap();
             writeln!(begin, ".endef").unwrap();
 
-            let section = link_section.unwrap_or_else(|| format!(".text.{asm_name}"));
-            writeln!(begin, ".pushsection {},\"xr\"", section).unwrap();
+            if let Some(section) = &link_section {
+                writeln!(begin, ".section {section},\"xr\"").unwrap()
+            } else if !function_sections {
+                // Function sections are enabled by default on MSVC and windows-gnullvm,
+                // but disabled by default on GNU.
+                writeln!(begin, ".text").unwrap();
+            } else {
+                // LLVM uses an extension to the section directive to support defining multiple
+                // sections with the same name and comdat. It adds `unique,<id>` at the end of the
+                // `.section` directive. We have no way of generating that unique ID here, so don't
+                // emit it.
+                //
+                // See https://llvm.org/docs/Extensions.html#id2.
+                match &tcx.sess.target.options.env {
+                    Env::Gnu => {
+                        writeln!(begin, ".section .text${asm_name},\"xr\",one_only,{asm_name}")
+                            .unwrap();
+                    }
+                    Env::Msvc => {
+                        writeln!(begin, ".section .text,\"xr\",one_only,{asm_name}").unwrap();
+                    }
+                    Env::Unspecified => match &tcx.sess.target.options.os {
+                        Os::Uefi => {
+                            writeln!(begin, ".section .text,\"xr\",one_only,{asm_name}").unwrap();
+                        }
+                        _ => bug!("unexpected coff target {}", tcx.sess.target.llvm_target),
+                    },
+                    other => bug!("unexpected coff env {other:?}"),
+                }
+            }
             write_linkage(&mut begin).unwrap();
             writeln!(begin, ".balign {align_bytes}").unwrap();
             writeln!(begin, "{asm_name}:").unwrap();

src/tools/run-make-support/src/external_deps/rustc.rs

@@ -449,6 +449,16 @@ impl Rustc {
         self.cmd.arg("-Zcodegen-source-order");
         self
     }
+
+    /// Specify `-Z function-sections={yes, no}`.
+    pub fn function_sections(&mut self, enable: bool) -> &mut Self {
+        let flag = match enable {
+            true => "-Zfunction-sections=yes",
+            false => "-Zfunction-sections=no",
+        };
+        self.cmd.arg(flag);
+        self
+    }
 }
 
 /// Query the sysroot path corresponding `rustc --print=sysroot`.

tests/assembly-llvm/naked-functions/function-sections.rs

@@ -0,0 +1,97 @@
+//@ add-minicore
+//@ assembly-output: emit-asm
+//
+//@ revisions: linux-x86-gnu-fs-true linux-x86-gnu-fs-false
+//@[linux-x86-gnu-fs-true] compile-flags: --target x86_64-unknown-linux-gnu -Zfunction-sections=true
+//@[linux-x86-gnu-fs-true] needs-llvm-components: x86
+//@[linux-x86-gnu-fs-false] compile-flags: --target x86_64-unknown-linux-gnu -Zfunction-sections=false
+//@[linux-x86-gnu-fs-false] needs-llvm-components: x86
+//
+//@ revisions: macos-aarch64-fs-true macos-aarch64-fs-false
+//@[macos-aarch64-fs-true] compile-flags: --target aarch64-apple-darwin -Zfunction-sections=true
+//@[macos-aarch64-fs-true] needs-llvm-components: aarch64
+//@[macos-aarch64-fs-false] compile-flags: --target aarch64-apple-darwin -Zfunction-sections=false
+//@[macos-aarch64-fs-false] needs-llvm-components: aarch64
+//
+//@ revisions: windows-x86-gnu-fs-true windows-x86-gnu-fs-false
+//@[windows-x86-gnu-fs-true] compile-flags: --target x86_64-pc-windows-gnu -Zfunction-sections=true
+//@[windows-x86-gnu-fs-true] needs-llvm-components: x86
+//@[windows-x86-gnu-fs-false] compile-flags: --target x86_64-pc-windows-gnu -Zfunction-sections=false
+//@[windows-x86-gnu-fs-false] needs-llvm-components: x86
+//
+//@ revisions: windows-x86-msvc-fs-true windows-x86-msvc-fs-false
+//@[windows-x86-msvc-fs-true] compile-flags: --target x86_64-pc-windows-msvc -Zfunction-sections=true
+//@[windows-x86-msvc-fs-true] needs-llvm-components: x86
+//@[windows-x86-msvc-fs-false] compile-flags: --target x86_64-pc-windows-msvc -Zfunction-sections=false
+//@[windows-x86-msvc-fs-false] needs-llvm-components: x86
+//
+//@ revisions: x86-uefi-fs-true x86-uefi-fs-false
+//@[x86-uefi-fs-true] compile-flags: --target x86_64-unknown-uefi -Zfunction-sections=true
+//@[x86-uefi-fs-true] needs-llvm-components: x86
+//@[x86-uefi-fs-false] compile-flags: --target x86_64-unknown-uefi -Zfunction-sections=false
+//@[x86-uefi-fs-false] needs-llvm-components: x86
+
+#![crate_type = "lib"]
+#![feature(no_core)]
+#![no_core]
+
+// Tests that naked and non-naked functions emit the same directives when (not) using
+// -Zfunction-sections. This setting is ignored on macos, off by default on windows gnu,
+// and on by default in the remaining revisions tested here.
+
+extern crate minicore;
+use minicore::*;
+
+#[unsafe(naked)]
+#[unsafe(no_mangle)]
+extern "C" fn naked_ret() {
+    // linux-x86-gnu-fs-true: .section .text.naked_ret,"ax",@progbits
+    // linux-x86-gnu-fs-false: .text
+    //
+    // macos-aarch64-fs-true:  .section __TEXT,__text,regular,pure_instructions
+    // macos-aarch64-fs-false: .section __TEXT,__text,regular,pure_instructions
+    //
+    // NOTE: the regular function below adds `unique,0` at the end, but we have no way of generating
+    // the unique ID to use there, so don't emit that part.
+    //
+    // windows-x86-gnu-fs-true: .section .text$naked_ret,"xr",one_only,naked_ret
+    // windows-x86-msvc-fs-true: .section .text,"xr",one_only,naked_ret
+    // x86-uefi-fs-true: .section .text,"xr",one_only,naked_ret
+    //
+    // windows-x86-gnu-fs-false: .text
+    // windows-x86-msvc-fs-false: .text
+    // x86-uefi-fs-false: .text
+    //
+    // CHECK-LABEL: naked_ret:
+    naked_asm!("ret")
+}
+
+// Use a different section here so that `regular_ret` has to explicitly specify the section.
+#[link_section = cfg_select!(
+    target_os = "macos" => "__FOO,bar",
+    _ => ".bar",
+)]
+#[unsafe(no_mangle)]
+extern "C" fn omarker() -> i32 {
+    // CHECK-LABEL: omarker:
+    32
+}
+
+#[unsafe(no_mangle)]
+extern "C" fn regular_ret() {
+    // linux-x86-gnu-fs-true: .section .text.regular_ret,"ax",@progbits
+    // linux-x86-gnu-fs-false: .text
+    //
+    // macos-aarch64-fs-true:  .section __TEXT,__text,regular,pure_instructions
+    // macos-aarch64-fs-false: .section __TEXT,__text,regular,pure_instructions
+    //
+    // windows-x86-gnu-fs-true: .section .text$regular_ret,"xr",one_only,regular_ret,unique,0
+    // windows-x86-msvc-fs-true: .section .text,"xr",one_only,regular_ret,unique,0
+    // x86-uefi-fs-true: .section .text,"xr",one_only,regular_ret,unique,0
+    //
+    // windows-x86-gnu-fs-false: .text
+    // windows-x86-msvc-fs-false: .text
+    // x86-uefi-fs-false: .text
+    //
+    // CHECK-LABEL: regular_ret:
+}

tests/assembly-llvm/naked-functions/link-section-windows.rs

@@ -1,4 +1,4 @@
-//@ revisions: windows-x86-gnu windows-x86-msvc
+//@ revisions: windows-x86-gnu windows-x86-msvc x86-uefi
 //@ add-minicore
 //@ assembly-output: emit-asm
 //
@@ -7,6 +7,9 @@
 //
 //@[windows-x86-msvc] compile-flags: --target x86_64-pc-windows-msvc
 //@[windows-x86-msvc] needs-llvm-components: x86
+//
+//@[x86-uefi] compile-flags: --target x86_64-unknown-uefi
+//@[x86-uefi] needs-llvm-components: x86
 
 #![crate_type = "lib"]
 #![feature(no_core)]