Add support for NIF unload and upgrade

rustler/src/codegen_runtime.rs

@@ -2,6 +2,7 @@
 
 use std::ffi::CString;
 use std::fmt;
+use std::panic::catch_unwind;
 
 use crate::{Encoder, Env, OwnedBinary, Term};
 
@@ -106,10 +107,10 @@ pub unsafe fn handle_nif_init_call(
     let term = Term::new(env, load_info);
 
     if let Some(inner) = function {
-        if inner(env, term) {
-            0
-        } else {
-            1
+        match catch_unwind(|| inner(env, term)) {
+            Ok(true) => 0,
+            Ok(false) => 1,
+            _ => 1,
         }
     } else {
         0
@@ -139,3 +140,36 @@ where
         }
     }
 }
+
+/// # Unsafe
+///
+/// This takes arguments, including raw pointers, that must be correct.
+pub unsafe fn handle_nif_upgrade_call(
+    function: Option<for<'a> fn(Env<'a>, Term<'a>) -> bool>,
+    r_env: NIF_ENV,
+    load_info: NIF_TERM,
+) -> c_int {
+    let env = Env::new(&(), r_env);
+    let term = Term::new(env, load_info);
+
+    if let Some(inner) = function {
+        match catch_unwind(|| inner(env, term)) {
+            Ok(true) => 0,
+            Ok(false) => 1,
+            _ => 1,
+        }
+    } else {
+        0
+    }
+}
+
+/// # Unsafe
+///
+/// This takes arguments, including raw pointers, that must be correct.
+pub unsafe fn handle_nif_unload_call(function: Option<for<'a> fn(Env<'a>)>, r_env: NIF_ENV) {
+    let env = Env::new(&(), r_env);
+
+    if let Some(inner) = function {
+        if catch_unwind(|| inner(env)).is_ok() {}
+    }
+}

rustler_codegen/src/init.rs

@@ -1,63 +1,125 @@
-use proc_macro2::{Span, TokenStream};
+use proc_macro2::TokenStream;
 use quote::quote;
+use std::collections::HashMap;
 use syn::parse::{Parse, ParseStream};
 use syn::punctuated::Punctuated;
 use syn::token::Comma;
-use syn::{Expr, Ident, Result, Token};
+use syn::{Expr, ExprAssign, ExprPath, Result, Token};
 
 #[derive(Debug)]
 pub struct InitMacroInput {
     name: syn::Lit,
     funcs: syn::ExprArray,
-    load: TokenStream,
+    load: Option<Expr>,
+    upgrade: Option<Expr>,
+    unload: Option<Expr>,
 }
 
 impl Parse for InitMacroInput {
     fn parse(input: ParseStream) -> Result<Self> {
         let name = syn::Lit::parse(input)?;
         let _comma = <syn::Token![,]>::parse(input)?;
         let funcs = syn::ExprArray::parse(input)?;
-        let options = parse_expr_assigns(input);
-        let load = extract_option(options, "load");
+        let options = parse_options(input);
 
-        Ok(InitMacroInput { name, funcs, load })
+        let allowed = ["load", "upgrade", "unload"];
+
+        for (key, _) in options.iter() {
+            if !allowed.contains(&key.as_str()) {
+                panic!("Option {} is not supported on init!()", key)
+            }
+        }
+
+        let get = |key| options.get(key).map(Clone::clone);
+
+        let load = get("load");
+        let unload = get("unload");
+        let upgrade = get("upgrade");
+
+        Ok(InitMacroInput {
+            name,
+            funcs,
+            load,
+            upgrade,
+            unload,
+        })
     }
 }
 
-fn parse_expr_assigns(input: ParseStream) -> Vec<syn::ExprAssign> {
-    let mut vec = Vec::new();
+fn parse_options(input: ParseStream) -> HashMap<String, Expr> {
+    let mut res = HashMap::new();
 
     while <Token![,]>::parse(input).is_ok() {
-        match syn::ExprAssign::parse(input) {
-            Ok(expr) => vec.push(expr),
-            Err(err) => panic!("{} (i.e. `load = load`)", err),
+        if input.is_empty() {
+            break;
         }
-    }
-    vec
-}
 
-fn extract_option(args: Vec<syn::ExprAssign>, name: &str) -> TokenStream {
-    for syn::ExprAssign { left, right, .. } in args.into_iter() {
-        if let syn::Expr::Path(syn::ExprPath { path, .. }) = &*left {
-            if let Some(ident) = path.get_ident() {
-                if *ident == name {
-                    let value = *right;
-                    return quote!(Some(#value));
+        match ExprAssign::parse(input) {
+            Ok(ExprAssign { left, right, .. }) => {
+                if let Expr::Path(ExprPath { path, .. }) = &*left {
+                    if let Some(ident) = path.get_ident() {
+                        res.insert(ident.to_string(), *right.clone());
+                    }
                 }
             }
+            Err(err) => panic!("{} (i.e. `load = load`", err),
         }
     }
-
-    let none = Ident::new("None", Span::call_site());
-    quote!(#none)
+    res
 }
 
 impl From<InitMacroInput> for proc_macro2::TokenStream {
     fn from(input: InitMacroInput) -> Self {
         let name = input.name;
         let num_of_funcs = input.funcs.elems.len();
         let funcs = nif_funcs(input.funcs.elems);
-        let load = input.load;
+
+        let nif_load_def = if let Some(load) = input.load {
+            quote! {
+                unsafe extern "C" fn nif_load(
+                    env: rustler::codegen_runtime::NIF_ENV,
+                    _priv_data: *mut *mut rustler::codegen_runtime::c_void,
+                    load_info: rustler::codegen_runtime::NIF_TERM
+                ) -> rustler::codegen_runtime::c_int {
+                    rustler::codegen_runtime::handle_nif_init_call(#load, env, load_info)
+                }
+                Some(nif_load)
+            }
+        } else {
+            quote!(None)
+        };
+
+        let nif_upgrade_def = if let Some(upgrade) = input.upgrade {
+            quote! {
+                unsafe extern "C" fn nif_upgrade(
+                    env: rustler::codegen_runtime::NIF_ENV,
+                    _priv_data: *mut *mut rustler::codegen_runtime::c_void,
+                    _old_priv_data: *mut *mut rustler::codegen_runtime::c_void,
+                    load_info: rustler::codegen_runtime::NIF_TERM,
+                 ) -> rustler::codegen_runtime::c_int {
+                    rustler::codegen_runtime::handle_nif_upgrade_call(#upgrade, env, load_info)
+                };
+
+                Some(nif_upgrade)
+            }
+        } else {
+            quote!(None)
+        };
+
+        let nif_unload_def = if let Some(unload) = input.unload {
+            quote! {
+                unsafe extern "C" fn nif_unload(
+                    env: rustler::codegen_runtime::NIF_ENV,
+                    _priv_data: *mut rustler::codegen_runtime::c_void,
+                 ) {
+                    rustler::codegen_runtime::handle_nif_unload_call(#unload, env)
+                };
+
+                Some(nif_unload)
+            }
+        } else {
+            quote!(None)
+        };
 
         let inner = quote! {
             static mut NIF_ENTRY: Option<rustler::codegen_runtime::DEF_NIF_ENTRY> = None;
@@ -69,22 +131,10 @@ impl From<InitMacroInput> for proc_macro2::TokenStream {
                 name: concat!(#name, "\0").as_ptr() as *const u8,
                 num_of_funcs: #num_of_funcs as rustler::codegen_runtime::c_int,
                 funcs: [#funcs].as_ptr(),
-                load: {
-                    extern "C" fn nif_load(
-                        env: rustler::codegen_runtime::NIF_ENV,
-                        _priv_data: *mut *mut rustler::codegen_runtime::c_void,
-                        load_info: rustler::codegen_runtime::NIF_TERM
-                    ) -> rustler::codegen_runtime::c_int {
-                        unsafe {
-                            // TODO: If an unwrap ever happens, we will unwind right into C! Fix this!
-                            rustler::codegen_runtime::handle_nif_init_call(#load, env, load_info)
-                        }
-                    }
-                    Some(nif_load)
-                },
+                load: { #nif_load_def },
                 reload: None,
-                upgrade: None,
-                unload: None,
+                upgrade: { #nif_upgrade_def },
+                unload: { #nif_unload_def },
                 vm_variant: b"beam.vanilla\0".as_ptr(),
                 options: 0,
                 sizeof_ErlNifResourceTypeInit: rustler::codegen_runtime::get_nif_resource_type_init_size(),

rustler_tests/native/rustler_test/src/lib.rs

@@ -79,10 +79,18 @@ rustler::init!(
         test_nif_attrs::can_rename,
         test_codegen::reserved_keywords::reserved_keywords_type_echo
     ],
-    load = load
+    load = Some(load),
+    upgrade = Some(upgrade),
+    unload = Some(unload),
 );
 
 fn load(env: rustler::Env, _: rustler::Term) -> bool {
     test_resource::on_load(env);
     true
 }
+
+fn upgrade(_env: rustler::Env, _: rustler::Term) -> bool {
+    false
+}
+
+fn unload(_env: rustler::Env) {}