Implement NIF discovery using linkme to statically collect NIF implem…

…entations
rusterlium · Jun 7, 2024 · cacb96f · cacb96f
1 parent 868c01b
commit cacb96f
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 45 deletions.
diff --git a/rustler/Cargo.toml b/rustler/Cargo.toml
@@ -22,7 +22,7 @@ nif_version_2_17 = ["nif_version_2_16", "rustler_sys/nif_version_2_17"]
 serde = ["dep:serde"]
 
 [dependencies]
-inventory = "0.3"
+linkme = "0.3"
 rustler_codegen = { path = "../rustler_codegen", version = "0.33.0", optional = true}
 rustler_sys = { path = "../rustler_sys", version = "~2.4.1" }
 num-bigint = { version = "0.4", optional = true }

diff --git a/rustler/src/codegen_runtime.rs b/rustler/src/codegen_runtime.rs
@@ -5,8 +5,9 @@ use std::fmt;
 
 use crate::{Encoder, Env, OwnedBinary, Term};
 
-// Re-export of inventory
-pub use inventory;
+// Re-export of linkme
+pub use crate::nif::RUSTLER_NIFS as NIFS;
+pub use linkme;
 
 // Names used by the `rustler::init!` macro or other generated code.
 pub use crate::wrapper::exception::raise_exception;

diff --git a/rustler/src/nif.rs b/rustler/src/nif.rs
@@ -1,5 +1,6 @@
 use crate::codegen_runtime::{c_char, c_int, c_uint, DEF_NIF_FUNC, NIF_ENV, NIF_TERM};
 
+#[repr(C)]
 pub struct Nif {
     pub name: *const c_char,
     pub arity: c_uint,
@@ -22,4 +23,6 @@ impl Nif {
 
 unsafe impl Sync for Nif {}
 
-inventory::collect!(Nif);
+#[no_mangle]
+#[linkme::distributed_slice]
+pub static RUSTLER_NIFS: [Nif];
diff --git a/rustler_codegen/Cargo.toml b/rustler_codegen/Cargo.toml
@@ -17,7 +17,6 @@ syn = { version = "2.0", features = ["full", "extra-traits"] }
 quote = "1.0"
 heck = "0.5"
 proc-macro2 = "1.0"
-inventory = "0.3"
 
 [dev-dependencies]
 trybuild = "1.0"
diff --git a/rustler_codegen/src/init.rs b/rustler_codegen/src/init.rs
@@ -86,8 +86,9 @@ impl From<InitMacroInput> for proc_macro2::TokenStream {
         let inner = quote! {
             static mut NIF_ENTRY: Option<rustler::codegen_runtime::DEF_NIF_ENTRY> = None;
             let nif_funcs: Box<[_]> =
-                rustler::codegen_runtime::inventory::iter::<rustler::Nif>()
-                .map(rustler::Nif::get_def)
+                rustler::codegen_runtime::NIFS
+                .iter()
+                .map(|n| n.get_def())
                 .collect();
 
             let entry = rustler::codegen_runtime::DEF_NIF_ENTRY {

diff --git a/rustler_codegen/src/nif.rs b/rustler_codegen/src/nif.rs
@@ -55,47 +55,48 @@ pub fn transcoder_decorator(nif_attributes: NifAttributes, fun: syn::ItemFn) ->
     }
 
     quote! {
-        rustler::codegen_runtime::inventory::submit!(
-            rustler::Nif {
-                name: concat!(#erl_func_name, "\0").as_ptr()
-                    as *const rustler::codegen_runtime::c_char,
-                arity: #arity,
-                flags: #flags as rustler::codegen_runtime::c_uint,
-                raw_func: {
-                    unsafe extern "C" fn nif_func(
-                        nif_env: rustler::codegen_runtime::NIF_ENV,
-                        argc: rustler::codegen_runtime::c_int,
-                        argv: *const rustler::codegen_runtime::NIF_TERM
-                    ) -> rustler::codegen_runtime::NIF_TERM {
-                        let lifetime = ();
-                        let env = rustler::Env::new(&lifetime, nif_env);
-
-                        let terms = std::slice::from_raw_parts(argv, argc as usize)
-                            .iter()
-                            .map(|term| rustler::Term::new(env, *term))
-                            .collect::<Vec<rustler::Term>>();
-
-                        fn wrapper<'a>(
-                            env: rustler::Env<'a>,
-                            args: &[rustler::Term<'a>]
-                        ) -> rustler::codegen_runtime::NifReturned {
-                            let result: std::thread::Result<_> =
-                                std::panic::catch_unwind(move || {
-                                    #decoded_terms
-                                    #function
-                                    Ok(#name(#argument_names))
-                                });
-
-                            rustler::codegen_runtime::handle_nif_result(
-                                result, env
-                            )
-                        }
-                        wrapper(env, &terms).apply(env)
+        #[allow(non_upper_case_globals)]
+        #[rustler::codegen_runtime::linkme::distributed_slice(rustler::codegen_runtime::NIFS)]
+        #[linkme(crate = rustler::codegen_runtime::linkme)]
+        static #name: rustler::Nif = rustler::Nif {
+            name: concat!(#erl_func_name, "\0").as_ptr()
+                as *const rustler::codegen_runtime::c_char,
+            arity: #arity,
+            flags: #flags as rustler::codegen_runtime::c_uint,
+            raw_func: {
+                unsafe extern "C" fn nif_func(
+                    nif_env: rustler::codegen_runtime::NIF_ENV,
+                    argc: rustler::codegen_runtime::c_int,
+                    argv: *const rustler::codegen_runtime::NIF_TERM
+                ) -> rustler::codegen_runtime::NIF_TERM {
+                    let lifetime = ();
+                    let env = rustler::Env::new(&lifetime, nif_env);
+
+                    let terms = std::slice::from_raw_parts(argv, argc as usize)
+                        .iter()
+                        .map(|term| rustler::Term::new(env, *term))
+                        .collect::<Vec<rustler::Term>>();
+
+                    fn wrapper<'a>(
+                        env: rustler::Env<'a>,
+                        args: &[rustler::Term<'a>]
+                    ) -> rustler::codegen_runtime::NifReturned {
+                        let result: std::thread::Result<_> =
+                            std::panic::catch_unwind(move || {
+                                #decoded_terms
+                                #function
+                                Ok(#name(#argument_names))
+                            });
+
+                        rustler::codegen_runtime::handle_nif_result(
+                            result, env
+                        )
                     }
-                    nif_func
+                    wrapper(env, &terms).apply(env)
                 }
+                nif_func
             }
-        );
+        };
     }
 }