From 2f1ced73e12561aaff630a022ce5e9cb3370f959 Mon Sep 17 00:00:00 2001
From: Florian Schmiderer <florian.schmiderer@posteo.net>
Date: Sun, 22 Sep 2024 14:07:06 +0200
Subject: [PATCH] Added tests to validate functions are hotpatchable

---
 .../src/external_deps/llvm.rs                 |  7 ++
 tests/run-make/hotpatch/lib.rs                | 64 +++++++++++++++++++
 tests/run-make/hotpatch/rmake.rs              | 44 +++++++++++++
 3 files changed, 115 insertions(+)
 create mode 100644 tests/run-make/hotpatch/lib.rs
 create mode 100644 tests/run-make/hotpatch/rmake.rs

diff --git a/src/tools/run-make-support/src/external_deps/llvm.rs b/src/tools/run-make-support/src/external_deps/llvm.rs
index 38a9ac923b4dc..d644488548251 100644
--- a/src/tools/run-make-support/src/external_deps/llvm.rs
+++ b/src/tools/run-make-support/src/external_deps/llvm.rs
@@ -254,6 +254,13 @@ impl LlvmFilecheck {
         self
     }
 
+    /// Specify the prefix (without :) for patterns to match. By default, these patterns are prefixed with "CHECK:".
+    pub fn check_prefix(&mut self, prefix: &str) -> &mut Self {
+        self.cmd.arg("--check-prefix");
+        self.cmd.arg(prefix);
+        self
+    }
+
     /// `--input-file` option.
     pub fn input_file<P: AsRef<Path>>(&mut self, input_file: P) -> &mut Self {
         self.cmd.arg("--input-file");
diff --git a/tests/run-make/hotpatch/lib.rs b/tests/run-make/hotpatch/lib.rs
new file mode 100644
index 0000000000000..6a0d0ce550823
--- /dev/null
+++ b/tests/run-make/hotpatch/lib.rs
@@ -0,0 +1,64 @@
+// hotpatch has two requirements:
+// 1. the first instruction of a functin must be at least two bytes long
+// 2. there must not be a jump to the first instruction
+
+// the hotpatch flag should insert nops as needed to fullfil the requirements,
+// but only if the the function does not already fulfill them.
+// Over 99% of function in regular codebases already fulfill the conditions,
+//  so its important to check that those are
+// unneccessarily affected
+
+// ------------------------------------------------------------------------------------------------
+
+// regularly this tailcall would jump to the first instruction the function
+// CHECK-LABEL: <tailcall_fn>:
+// CHECK: jne 0x0 <tailcall_fn>
+
+// hotpatch insert nops so that the tailcall will not jump to the first instruction of the function
+// HOTPATCH-LABEL: <tailcall_fn>:
+// HOTPATCH-NOT: jne 0x0 <tailcall_fn>
+
+#[no_mangle]
+pub fn tailcall_fn() {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    static COUNT: AtomicUsize = AtomicUsize::new(0);
+    if COUNT.fetch_sub(1, Ordering::Relaxed) != 0 {
+        tailcall_fn()
+    }
+}
+
+// ------------------------------------------------------------------------------------------------
+
+// empty_fn just returns. Note that 'ret' is a single byte instruction, but hotpatch requires a two
+// or more byte instructions to be at the start of the functions.
+// Preferably we would also tests a different single byte instruction,
+// but I was not able to make rustc emit anything but 'ret'.
+
+// CHECK-LABEL: <empty_fn>:
+// CHECK-NEXT: ret
+
+// HOTPATCH-LABEL: <empty_fn>:
+// HOTPATCH-NOT: ret
+// HOTPATCH: ret
+
+#[no_mangle]
+#[inline(never)]
+pub fn empty_fn() {}
+
+// ------------------------------------------------------------------------------------------------
+
+// return_42 should not be affected by hotpatch
+
+// CHECK-LABEL: <return_42>:
+// CHECK-NEXT: 0:
+// CHECK-NEXT: ret
+
+// HOTPATCH-LABEL: <return_42>:
+// HOTPATCH-NEXT: 0:
+// HOTPATCH-NEXT: ret
+
+#[no_mangle]
+#[inline(never)]
+pub fn return_42() -> i32 {
+    42
+}
diff --git a/tests/run-make/hotpatch/rmake.rs b/tests/run-make/hotpatch/rmake.rs
new file mode 100644
index 0000000000000..65635b3afcba9
--- /dev/null
+++ b/tests/run-make/hotpatch/rmake.rs
@@ -0,0 +1,44 @@
+// Check if hotpatch only makes the functions hotpachable that were not,
+// but leaving the other functions untouched
+// More details in lib.rs
+
+//@ revisions: x32 x64
+//@[x32] only-x86
+//@[x64] only-x86_64
+// Reason: hotpatch is only implemented for X86
+
+use run_make_support::{llvm, rfs, rustc};
+
+fn main() {
+    {
+        rustc().input("lib.rs").crate_name("regular").crate_type("lib").opt_level("3").run();
+
+        let regular_dump = llvm::llvm_objdump()
+            .arg("--disassemble-symbols=tailcall_fn,empty_fn,return_42")
+            .input("libregular.rlib")
+            .run();
+
+        llvm::llvm_filecheck().patterns("lib.rs").stdin_buf(regular_dump.stdout_utf8()).run();
+    }
+
+    {
+        rustc()
+            .input("lib.rs")
+            .crate_name("hotpatch")
+            .crate_type("lib")
+            .opt_level("3")
+            .arg("-Zhotpatch")
+            .run();
+
+        let hotpatch_dump = llvm::llvm_objdump()
+            .arg("--disassemble-symbols=tailcall_fn,empty_fn,return_42")
+            .input("libhotpatch.rlib")
+            .run();
+
+        llvm::llvm_filecheck()
+            .patterns("lib.rs")
+            .check_prefix("HOTPATCH")
+            .stdin_buf(hotpatch_dump.stdout_utf8())
+            .run();
+    }
+}