Support `#[patchable_function_entries]`

See [RFC](https://github.com/maurer/rust-rfcs/blob/patchable-function-entry/text/0000-patchable-function-entry.md) (yet to be numbered)

TODO before submission:
* Needs an RFC
* Improve error reporting for malformed attributes

9 files changed, 83 insertions, 4 deletions

diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
index 7cf789ab5d7..cd82894af18 100644
--- a/compiler/rustc_codegen_llvm/src/attributes.rs
+++ b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -2,7 +2,7 @@
 
 use rustc_codegen_ssa::traits::*;
 use rustc_hir::def_id::DefId;
-use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, PatchableFunctionEntry};
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_session::config::{FunctionReturn, OptLevel};
 use rustc_span::symbol::sym;
@@ -56,9 +56,12 @@ fn inline_attr<'ll>(cx: &CodegenCx<'ll, '_>, inline: InlineAttr) -> Option<&'ll
 #[inline]
 fn patchable_function_entry_attrs<'ll>(
     cx: &CodegenCx<'ll, '_>,
+    attr: Option<PatchableFunctionEntry>,
 ) -> SmallVec<[&'ll Attribute; 2]> {
     let mut attrs = SmallVec::new();
-    let patchable_spec = cx.tcx.sess.opts.unstable_opts.patchable_function_entry;
+    let patchable_spec = attr.unwrap_or_else(|| {
+        PatchableFunctionEntry::from_config(cx.tcx.sess.opts.unstable_opts.patchable_function_entry)
+    });
     let entry = patchable_spec.entry();
     let prefix = patchable_spec.prefix();
     if entry > 0 {
@@ -446,7 +449,7 @@ pub fn from_fn_attrs<'ll, 'tcx>(
         llvm::set_alignment(llfn, align);
     }
     to_add.extend(sanitize_attrs(cx, codegen_fn_attrs.no_sanitize));
-    to_add.extend(patchable_function_entry_attrs(cx));
+    to_add.extend(patchable_function_entry_attrs(cx, codegen_fn_attrs.patchable_function_entry));
 
     // Always annotate functions with the target-cpu they are compiled for.
     // Without this, ThinLTO won't inline Rust functions into Clang generated
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
index fb71cdaa8ff..8924ddb2aed 100644
--- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
+++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -5,7 +5,9 @@ use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId, LOCAL_CRATE};
 use rustc_hir::{lang_items, weak_lang_items::WEAK_LANG_ITEMS, LangItem};
-use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
+use rustc_middle::middle::codegen_fn_attrs::{
+    CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry,
+};
 use rustc_middle::mir::mono::Linkage;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::{self as ty, TyCtxt};
@@ -463,6 +465,28 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                     None
                 };
             }
+            sym::patchable_function_entry => {
+                codegen_fn_attrs.patchable_function_entry = attr.meta_item_list().and_then(|l| {
+                    let mut prefix = 0;
+                    let mut entry = 0;
+                    for item in l {
+                        if let Some((sym, lit)) = item.name_value_literal() {
+                            let val = match lit.kind {
+                                // FIXME emit error if too many nops requested
+                                rustc_ast::LitKind::Int(i, _) => i as u8,
+                                _ => continue,
+                            };
+                            match sym {
+                                sym::prefix => prefix = val,
+                                sym::entry => entry = val,
+                                // FIXME possibly emit error here?
+                                _ => continue,
+                            }
+                        }
+                    }
+                    Some(PatchableFunctionEntry::from_prefix_and_entry(prefix, entry))
+                })
+            }
             _ => {}
         }
     }
diff --git a/compiler/rustc_feature/src/builtin_attrs.rs b/compiler/rustc_feature/src/builtin_attrs.rs
index c53bf965139..b5f9f2c715f 100644
--- a/compiler/rustc_feature/src/builtin_attrs.rs
+++ b/compiler/rustc_feature/src/builtin_attrs.rs
@@ -584,6 +584,13 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         pointee, Normal, template!(Word), ErrorFollowing,
         EncodeCrossCrate::No, derive_smart_pointer, experimental!(pointee)
     ),
+    
+    // FIXME RFC
+    // `#[patchable_function_entry(prefix(n), entry(n))]`
+    gated!(
+        patchable_function_entry, Normal, template!(List: "prefix(n), entry(n)"), ErrorPreceding,
+        experimental!(patchable_function_entry)
+    ),
 
     // ==========================================================================
     // Internal attributes: Stability, deprecation, and unsafe:
diff --git a/compiler/rustc_feature/src/unstable.rs b/compiler/rustc_feature/src/unstable.rs
index 2dfaac8f6e7..796475e766f 100644
--- a/compiler/rustc_feature/src/unstable.rs
+++ b/compiler/rustc_feature/src/unstable.rs
@@ -565,6 +565,9 @@ declare_features! (
     (unstable, offset_of_slice, "CURRENT_RUSTC_VERSION", Some(126151)),
     /// Allows using `#[optimize(X)]`.
     (unstable, optimize_attribute, "1.34.0", Some(54882)),
+    /// Allows specifying nop padding on functions for dynamic patching.
+    // FIXME this needs an RFC #
+    (unstable, patchable_function_entry, "CURRENT_RUSTC_VERSION", Some(9999)),
     /// Allows postfix match `expr.match { ... }`
     (unstable, postfix_match, "1.79.0", Some(121618)),
     /// Allows `use<'a, 'b, A, B>` in `impl Trait + use<...>` for precise capture of generic args.
diff --git a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
index 0fce26dcbbd..23fe72c537a 100644
--- a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
+++ b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
@@ -45,6 +45,9 @@ pub struct CodegenFnAttrs {
     /// The `#[repr(align(...))]` attribute. Indicates the value of which the function should be
     /// aligned to.
     pub alignment: Option<Align>,
+    /// The `#[patchable_function_entry(...)]` attribute. Indicates how many nops should be around
+    /// the function entry.
+    pub patchable_function_entry: Option<PatchableFunctionEntry>,
 }
 
 #[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, HashStable)]
@@ -147,6 +150,7 @@ impl CodegenFnAttrs {
             no_sanitize: SanitizerSet::empty(),
             instruction_set: None,
             alignment: None,
+            patchable_function_entry: None,
         }
     }
 
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 6d4a8c29bc9..ea7fe7e76c4 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -768,6 +768,7 @@ symbols! {
         enable,
         encode,
         end,
+        entry,
         enumerate_method,
         env,
         env_CFG_RELEASE: env!("CFG_RELEASE"),
@@ -1383,6 +1384,7 @@ symbols! {
         passes,
         pat,
         pat_param,
+        patchable_function_entry,
         path,
         pattern_complexity,
         pattern_parentheses,
@@ -1421,6 +1423,7 @@ symbols! {
         prefetch_read_instruction,
         prefetch_write_data,
         prefetch_write_instruction,
+        prefix,
         preg,
         prelude,
         prelude_import,
diff --git a/tests/codegen/patchable-function-entry.rs b/tests/codegen/patchable-function-entry.rs
index f06739303d2..dc20c0a2c6d 100644
--- a/tests/codegen/patchable-function-entry.rs
+++ b/tests/codegen/patchable-function-entry.rs
@@ -1,8 +1,28 @@
+#![feature(patchable_function_entry)]
 // compile-flags: -Z patchable-function-entry=15,10
 
 #![crate_type = "lib"]
 
+// This should have the default, as set by the compile flags
 #[no_mangle]
 pub fn foo() {}
+
+// The attribute should override the compile flags
+#[no_mangle]
+#[patchable_function_entry(prefix(1), entry(2))]
+pub fn bar() {}
+
+// If we override an attribute to 0 or unset, the attribute should go away
+#[no_mangle]
+#[patchable_function_entry(entry(0))]
+pub fn baz() {}
+
 // CHECK: @foo() unnamed_addr #0
+// CHECK: @bar() unnamed_addr #1
+// CHECK: @baz() unnamed_addr #2
+
 // CHECK: attributes #0 = { {{.*}}"patchable-function-entry"="5"{{.*}}"patchable-function-prefix"="10" {{.*}} }
+// CHECK: attributes #1 = { {{.*}}"patchable-function-entry"="2"{{.*}}"patchable-function-prefix"="1" {{.*}} }
+// CHECK-NOT: attributes #2 = { {{.*}}patchable-function-entry{{.*}} }
+// CHECK-NOT: attributes #2 = { {{.*}}patchable-function-prefix{{.*}} }
+// CHECK: attributes #2 = { {{.*}} }
diff --git a/tests/ui/feature-gates/feature-gate-patchable-function-entry.rs b/tests/ui/feature-gates/feature-gate-patchable-function-entry.rs
new file mode 100644
index 00000000000..0e16e873a1e
--- /dev/null
+++ b/tests/ui/feature-gates/feature-gate-patchable-function-entry.rs
@@ -0,0 +1,3 @@
+#[patchable_function_entry(entry(1), prefix(1))]
+//~^ ERROR: the `#[patchable_function_entry]` attribute is an experimental feature
+fn main() {}
diff --git a/tests/ui/feature-gates/feature-gate-patchable-function-entry.stderr b/tests/ui/feature-gates/feature-gate-patchable-function-entry.stderr
new file mode 100644
index 00000000000..c4d57d774e3
--- /dev/null
+++ b/tests/ui/feature-gates/feature-gate-patchable-function-entry.stderr
@@ -0,0 +1,12 @@
+error[E0658]: the `#[patchable_function_entry]` attribute is an experimental feature
+  --> $DIR/feature-gate-patchable-function-entry.rs:1:1
+   |
+LL | #[patchable_function_entry(entry(1), prefix(1))]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: see issue #9999 <https://github.com/rust-lang/rust/issues/9999> for more information
+   = help: add `#![feature(patchable_function_entry)]` to the crate attributes to enable
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0658`.