diff --git a/rustler/src/codegen_runtime.rs b/rustler/src/codegen_runtime.rs index 725f42c3..e1f4d3bf 100644 --- a/rustler/src/codegen_runtime.rs +++ b/rustler/src/codegen_runtime.rs @@ -8,6 +8,9 @@ use crate::{Encoder, Env, OwnedBinary, Term}; // Re-export of inventory pub use inventory; +// Re-export of resource registration +pub use crate::resource::Registration as ResourceRegistration; + // Names used by the `rustler::init!` macro or other generated code. pub use crate::wrapper::exception::raise_exception; pub use crate::wrapper::{ diff --git a/rustler/src/lib.rs b/rustler/src/lib.rs index a7aea28b..005c732e 100644 --- a/rustler/src/lib.rs +++ b/rustler/src/lib.rs @@ -72,8 +72,8 @@ pub use nif::Nif; pub type NifResult = Result; pub use rustler_codegen::{ - init, nif, NifException, NifMap, NifRecord, NifStruct, NifTaggedEnum, NifTuple, NifUnitEnum, - NifUntaggedEnum, + init, nif, resource_impl, NifException, NifMap, NifRecord, NifStruct, NifTaggedEnum, NifTuple, + NifUnitEnum, NifUntaggedEnum, }; #[cfg(feature = "serde")] diff --git a/rustler/src/resource/mod.rs b/rustler/src/resource/mod.rs index 5fa9fa29..fc078168 100644 --- a/rustler/src/resource/mod.rs +++ b/rustler/src/resource/mod.rs @@ -15,6 +15,7 @@ mod util; pub use arc::ResourceArc; pub use error::ResourceInitError; pub use monitor::Monitor; +pub use registration::Registration; pub use traits::Resource; use traits::ResourceExt; diff --git a/rustler/src/resource/registration.rs b/rustler/src/resource/registration.rs index c06db5e4..65054a4f 100644 --- a/rustler/src/resource/registration.rs +++ b/rustler/src/resource/registration.rs @@ -13,12 +13,16 @@ use std::mem::MaybeUninit; use std::ptr; #[derive(Debug)] -struct Registration { +pub struct Registration { get_type_id: fn() -> TypeId, get_type_name: fn() -> &'static str, init: ErlNifResourceTypeInit, } +unsafe impl Sync for Registration {} + +inventory::collect!(Registration); + impl<'a> Env<'a> { /// Register a resource type, see `Registration::register`. pub fn register(&self) -> Result<(), ResourceInitError> { @@ -34,6 +38,15 @@ impl<'a> Env<'a> { /// `std::mem::needs_drop`). All other callbacks are only registered if `IMPLEMENTS_...` is set to /// `true`. impl Registration { + /// Register all resource types that have been submitted to the inventory. + pub fn register_all_collected(env: Env) -> Result<(), ResourceInitError> { + for reg in inventory::iter::() { + reg.register(env)?; + } + + Ok(()) + } + /// Generate a new (pending) resource type registration. pub const fn new() -> Self { Self { diff --git a/rustler_codegen/src/init.rs b/rustler_codegen/src/init.rs index 6d11ca5e..15e8effa 100644 --- a/rustler_codegen/src/init.rs +++ b/rustler_codegen/src/init.rs @@ -106,6 +106,11 @@ impl From for proc_macro2::TokenStream { let mut env = rustler::Env::new_init_env(&env, env); // TODO: If an unwrap ever happens, we will unwind right into C! Fix this! let load_info = rustler::Term::new(env, load_info); + + if !rustler::codegen_runtime::ResourceRegistration::register_all_collected(env).is_ok() { + return 1; + } + #load.map_or(0, |inner| { rustler::codegen_runtime::handle_nif_init_call( inner, env, load_info diff --git a/rustler_codegen/src/lib.rs b/rustler_codegen/src/lib.rs index 8bace2f9..411fbed2 100644 --- a/rustler_codegen/src/lib.rs +++ b/rustler_codegen/src/lib.rs @@ -9,6 +9,7 @@ mod init; mod map; mod nif; mod record; +mod resource_impl; mod tagged_enum; mod tuple; mod unit_enum; @@ -400,3 +401,10 @@ pub fn nif_untagged_enum(input: TokenStream) -> TokenStream { let ast = syn::parse(input).unwrap(); untagged_enum::transcoder_decorator(&ast).into() } + +#[proc_macro_attribute] +pub fn resource_impl(_attr: TokenStream, item: TokenStream) -> TokenStream { + let input = syn::parse_macro_input!(item as syn::ItemImpl); + + resource_impl::transcoder_decorator(input).into() +} diff --git a/rustler_codegen/src/resource_impl.rs b/rustler_codegen/src/resource_impl.rs new file mode 100644 index 00000000..80ebe44f --- /dev/null +++ b/rustler_codegen/src/resource_impl.rs @@ -0,0 +1,42 @@ +use proc_macro2::{Span, TokenStream}; +use quote::quote; +use std::collections::HashSet; + +pub fn transcoder_decorator(mut input: syn::ItemImpl) -> TokenStream { + // Should be `Resource` but will fail somewhere else anyway if it isn't. + // let (_, _trait_path, _) = input.trait_.unwrap(); + let type_path = match *input.self_ty { + syn::Type::Path(ref type_path) => type_path.clone(), + _ => panic!("Can only implement trait on concrete types"), + }; + + let mut to_add: HashSet = HashSet::new(); + let mut already_has: HashSet = HashSet::new(); + + for item in input.items.iter() { + if let syn::ImplItem::Fn(f) = item { + to_add.insert( + format!("IMPLEMENTS_{}", f.sig.ident.to_string().to_uppercase()).to_string(), + ); + } + + if let syn::ImplItem::Const(f) = item { + already_has.insert(f.ident.to_string()); + } + } + + for add in to_add.difference(&already_has) { + let ident = syn::Ident::new(add, Span::call_site()); + let impl_item: syn::ImplItem = syn::parse_quote!(const #ident: bool = true;); + + input.items.push(impl_item); + } + + quote!( + #input + + rustler::codegen_runtime::inventory::submit!( + rustler::codegen_runtime::ResourceRegistration::new::<#type_path>() + ); + ) +} diff --git a/rustler_tests/native/rustler_test/src/test_resource.rs b/rustler_tests/native/rustler_test/src/test_resource.rs index 2ede9954..95f6771a 100644 --- a/rustler_tests/native/rustler_test/src/test_resource.rs +++ b/rustler_tests/native/rustler_test/src/test_resource.rs @@ -14,9 +14,8 @@ pub struct TestMonitorResource { inner: Mutex, } +#[rustler::resource_impl(register = true)] impl Resource for TestMonitorResource { - const IMPLEMENTS_DOWN: bool = true; - fn down<'a>(&'a self, _env: Env<'a>, _pid: LocalPid, mon: Monitor) { let mut inner = self.inner.lock().unwrap(); assert!(Some(mon) == inner.mon); @@ -32,6 +31,7 @@ pub struct ImmutableResource { b: u32, } +#[rustler::resource_impl(register = false)] impl Resource for ImmutableResource {} pub struct WithBinaries { @@ -39,11 +39,12 @@ pub struct WithBinaries { b: Vec, } +impl Resource for WithBinaries {} + pub fn on_load(env: Env) -> bool { rustler::resource!(TestResource, env) + && env.register::().is_ok() && env.register::().is_ok() - && env.register::().is_ok() - && rustler::resource!(WithBinaries, env) } #[rustler::nif]