kernel: add indirect function support

* Use it to determine the fastest implementation of `memcpy` and `memset`.

Signed-off-by: Anhad Singh <[email protected]>

Author: Andy-Python-Programmer

src/aero_kernel/src/arch/x86_64/mem.rs

@@ -15,7 +15,16 @@
 // You should have received a copy of the GNU General Public License
 // along with Aero. If not, see <https://www.gnu.org/licenses/>.
 
-#[no_mangle]
+fn should_store_by_byte() -> bool {
+    let cpuid = raw_cpuid::CpuId::new();
+    if let Some(features) = cpuid.get_extended_feature_info() {
+        // Check if "Enhanced" or "Fast Short" optimizations are available.
+        features.has_rep_movsb_stosb()
+    } else {
+        false
+    }
+}
+
 #[naked]
 unsafe extern "C" fn memcpy_movsq(dest: *mut u8, src: *const u8, len: usize) -> *mut u8 {
     // Registers used:
@@ -39,7 +48,33 @@ unsafe extern "C" fn memcpy_movsq(dest: *mut u8, src: *const u8, len: usize) ->
     );
 }
 
-#[no_mangle]
+#[naked]
+unsafe extern "C" fn memcpy_movsb(dest: *mut u8, src: *const u8, len: usize) -> *mut u8 {
+    // Registers used:
+    //
+    // %rdi = argument 1, `dest`
+    // %rsi = argument 2, `src`
+    // %rdx = argument 3, `len`
+    asm!(
+        // Save the return value.
+        "mov rax, rdi",
+        // Copy!
+        "mov rcx, rdx",
+        "rep movsb",
+        "ret",
+        options(noreturn)
+    )
+}
+
+#[indirect]
+extern "C" fn memcpy() -> fn(*mut u8, *const u8, usize) {
+    if should_store_by_byte() {
+        memcpy_movsb
+    } else {
+        memcpy_movsq
+    }
+}
+
 #[naked]
 unsafe extern "C" fn memset_stosq(dest: *mut u8, byte: i32, len: usize) -> *mut u8 {
     // Registers used:
@@ -70,6 +105,34 @@ unsafe extern "C" fn memset_stosq(dest: *mut u8, byte: i32, len: usize) -> *mut
     )
 }
 
+#[naked]
+unsafe extern "C" fn memset_stosb(dest: *mut u8, byte: i32, len: usize) -> *mut u8 {
+    // Registers used:
+    //
+    // %rdi = argument 1, `dest`
+    // %rsi = argument 2, `byte`
+    // %rdx = argument 3, `len`
+    asm!(
+        // Save the return value.
+        "mov r11, rdi",
+        "mov al, sil",
+        "mov rcx, rdx",
+        "rep stosb",
+        "mov rax, r11",
+        "ret",
+        options(noreturn)
+    )
+}
+
+#[indirect]
+extern "C" fn memset() -> fn(*mut u8, i32, usize) {
+    if should_store_by_byte() {
+        memset_stosb
+    } else {
+        memset_stosq
+    }
+}
+
 #[no_mangle]
 #[naked]
 unsafe extern "C" fn memmove_erms(dest: *mut u8, src: *const u8, len: usize) -> *mut u8 {
@@ -108,18 +171,6 @@ unsafe extern "C" fn memmove_erms(dest: *mut u8, src: *const u8, len: usize) ->
     )
 }
 
-// FIXME(andypython): pick the best implementation for the current CPU using indirect functions.
-
-#[no_mangle]
-extern "C" fn memcpy(dest: *mut u8, src: *const u8, len: usize) -> *mut u8 {
-    unsafe { memcpy_movsq(dest, src, len) }
-}
-
-#[no_mangle]
-extern "C" fn memset(dest: *mut u8, byte: i32, len: usize) -> *mut u8 {
-    unsafe { memset_stosq(dest, byte, len) }
-}
-
 #[no_mangle]
 extern "C" fn memmove(dest: *mut u8, src: *const u8, len: usize) -> *mut u8 {
     unsafe { memmove_erms(dest, src, len) }

src/aero_kernel/src/arch/x86_64/mod.rs

@@ -59,25 +59,6 @@ static HHDM: HhdmRequest = HhdmRequest::new(0);
 
 #[no_mangle]
 extern "C" fn arch_aero_main() -> ! {
-    unsafe {
-        core::ptr::read_volatile(STACK.get_response().as_ptr().unwrap());
-    }
-
-    // SAFETY: We have exclusive access to the memory map.
-    let memmap = MEMMAP
-        .get_response()
-        .get_mut()
-        .expect("limine: invalid memmap response")
-        .memmap_mut();
-
-    unsafe {
-        interrupts::disable_interrupts();
-    }
-
-    unsafe {
-        crate::PHYSICAL_MEMORY_OFFSET = VirtAddr::new(HHDM.get_response().get().unwrap().offset);
-    }
-
     let kernel_file_resp = KERNEL_FILE
         .get_response()
         .get()
@@ -107,6 +88,27 @@ extern "C" fn arch_aero_main() -> ! {
         UnwindInfo::new(elf)
     });
 
+    crate::relocate_self();
+
+    unsafe {
+        core::ptr::read_volatile(STACK.get_response().as_ptr().unwrap());
+    }
+
+    // SAFETY: We have exclusive access to the memory map.
+    let memmap = MEMMAP
+        .get_response()
+        .get_mut()
+        .expect("limine: invalid memmap response")
+        .memmap_mut();
+
+    unsafe {
+        interrupts::disable_interrupts();
+    }
+
+    unsafe {
+        crate::PHYSICAL_MEMORY_OFFSET = VirtAddr::new(HHDM.get_response().get().unwrap().offset);
+    }
+
     // Now that we have unwind info, we can initialize the COM ports. This
     // will be used to print panic messages/logs before the debug renderer is
     // initialized to the serial output (if available).

src/aero_kernel/src/main.rs

@@ -112,6 +112,32 @@ static mut PHYSICAL_MEMORY_OFFSET: VirtAddr = VirtAddr::zero();
 
 const IO_VIRTUAL_BASE: VirtAddr = VirtAddr::new(0xffffff0000000000);
 
+const STT_GNU_IFUNC: u32 = 37;
+
+pub fn relocate_self() {
+    use xmas_elf::sections::SectionData;
+
+    let unwind_info = unwind::UNWIND_INFO.get().unwrap();
+    let kernel_elf = &unwind_info.kernel_elf;
+
+    for section in kernel_elf.section_iter() {
+        if let Ok(SectionData::Rela64(rela)) = section.get_data(kernel_elf) {
+            for item in rela {
+                if !item.get_type() == STT_GNU_IFUNC {
+                    continue;
+                }
+
+                let offset = unsafe { &mut *(item.get_offset() as *mut usize) };
+
+                let resolver_ptr = item.get_addend() as *const u8;
+                let resolver: fn() -> usize = unsafe { core::mem::transmute(resolver_ptr) };
+
+                *offset = resolver();
+            }
+        }
+    }
+}
+
 fn aero_main() -> ! {
     // NOTE: In this function we only want to initialize essential services, including
     // the task scheduler. Rest of the initializing (including kernel modules) should go

src/aero_kernel/src/unwind.rs

@@ -60,7 +60,7 @@ pub fn prepare_panic() {
 }
 
 pub struct UnwindInfo {
-    kernel_elf: ElfFile<'static>,
+    pub kernel_elf: ElfFile<'static>,
 }
 
 impl UnwindInfo {

src/aero_proc/src/indirect.rs

@@ -0,0 +1,42 @@
+use proc_macro::TokenStream;
+use proc_macro2::{Ident, Span};
+
+pub fn parse(_: TokenStream, item: TokenStream) -> TokenStream {
+    let item = syn::parse_macro_input!(item as syn::ItemFn);
+    let args = item.sig.inputs;
+
+    let name = item.sig.ident.to_string();
+
+    // Underscores at the beginning of the identifier make it reserved, and the more underscores
+    // there are, the more reserveder it is.
+    let resolve_name = Ident::new(&format!("__resolve_{name}"), Span::call_site());
+
+    let inline = format!(
+        r"
+        .global {name}
+
+        .type {name}, @gnu_indirect_function
+        .set {name},{{}}
+        "
+    );
+
+    let name = &item.sig.ident;
+    let resolve_body = &item.block;
+
+    quote::quote! {
+        fn #resolve_name() -> usize {
+            let resolved_function = {
+                #resolve_body
+            };
+
+            resolved_function as usize
+        }
+
+        ::core::arch::global_asm!(#inline, sym #resolve_name);
+
+        extern "C" {
+            fn #name(#args);
+        }
+    }
+    .into()
+}

src/aero_proc/src/lib.rs

@@ -22,6 +22,7 @@ extern crate proc_macro_error;
 
 mod cpu_local;
 mod downcastable;
+mod indirect;
 mod syscall;
 mod test;
 
@@ -64,3 +65,9 @@ pub fn downcastable(attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn cpu_local(attr: TokenStream, item: TokenStream) -> TokenStream {
     cpu_local::parse(attr, item)
 }
+
+#[proc_macro_attribute]
+#[proc_macro_error]
+pub fn indirect(attr: TokenStream, item: TokenStream) -> TokenStream {
+    indirect::parse(attr, item)
+}