NIF Codegen Rework

Signed-off-by: Sonny Scroggin <[email protected]>

README.md

@@ -35,28 +35,12 @@ NOTE: If you have previously used Rustler, you need to run `mix archive.uninstal
 This is the code for a minimal NIF that adds two numbers and returns the result.
 
 ```rust
-use rustler::{Encoder, Env, Error, Term};
-
-mod atoms {
-    rustler::atoms! {
-        ok,
-    }
+#[rustler::nif]
+fn add(a: i64, b: i64) -> i64 {
+    a + b
 }
 
-rustler::init!(
-    "Elixir.Math",
-    [
-        ("add", 2, add)
-    ],
-    None
-);
-
-fn add<'a>(env: Env<'a>, args: &[Term<'a>]) -> Result<Term<'a>, Error> {
-    let a: i64 = args[0].decode()?;
-    let b: i64 = args[1].decode()?;
-
-    Ok((atoms::ok(), a + b).encode(env))
-}
+rustler::init!("Elixir.Math", [add]);
 ```
 
 #### Supported nif_version

rustler/Cargo.toml

@@ -13,9 +13,9 @@ derive = ["rustler_codegen"]
 alternative_nif_init_name = []
 
 [dependencies]
+lazy_static = "1.4"
 rustler_codegen = { path = "../rustler_codegen", version = "0.21.0", optional = true }
 rustler_sys = { path = "../rustler_sys", version = "~2.0" }
-lazy_static = "1.4"
 
 [build-dependencies]
 lazy_static = "1.4"

rustler/src/codegen_runtime.rs

@@ -1,10 +1,11 @@
 //! Functions used by runtime generated code. Should not be used.
 
 use std::ffi::CString;
+use std::fmt;
 
 use crate::{Env, Term};
 
-// Names used by the `init!` macro or other generated code.
+// Names used by the `rustler::init!` macro or other generated code.
 pub use crate::wrapper::exception::raise_exception;
 pub use crate::wrapper::{
     c_int, c_void, get_nif_resource_type_init_size, DEF_NIF_ENTRY, DEF_NIF_FUNC,
@@ -17,6 +18,7 @@ pub use rustler_sys::{TWinDynNifCallbacks, WIN_DYN_NIF_CALLBACKS};
 pub unsafe trait NifReturnable {
     unsafe fn as_returned(self, env: Env) -> NifReturned;
 }
+
 unsafe impl<T> NifReturnable for T
 where
     T: crate::Encoder,
@@ -25,6 +27,7 @@ where
         NifReturned::Term(self.encode(env).as_c_arg())
     }
 }
+
 unsafe impl<T> NifReturnable for Result<T, crate::error::Error>
 where
     T: crate::Encoder,
@@ -48,6 +51,7 @@ pub enum NifReturned {
         args: Vec<NIF_TERM>,
     },
 }
+
 impl NifReturned {
     pub unsafe fn apply(self, env: Env) -> NIF_TERM {
         match self {
@@ -73,6 +77,17 @@ impl NifReturned {
     }
 }
 
+impl fmt::Debug for NifReturned {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
+        match self {
+            NifReturned::BadArg => write!(fmt, "{{error, badarg}}"),
+            NifReturned::Term(ref s) => write!(fmt, "{{ok, {}}}", s),
+            NifReturned::Raise(ref s) => write!(fmt, "throw({})", s),
+            NifReturned::Reschedule { .. } => write!(fmt, "reschedule()"),
+        }
+    }
+}
+
 /// # Unsafe
 ///
 /// This takes arguments, including raw pointers, that must be correct.
@@ -94,3 +109,27 @@ pub unsafe fn handle_nif_init_call(
         0
     }
 }
+
+pub fn handle_nif_result<T>(
+    result: std::thread::Result<Result<T, crate::error::Error>>,
+    env: Env,
+) -> NifReturned
+where
+    T: NifReturnable,
+{
+    unsafe {
+        match result {
+            Ok(res) => match res {
+                Ok(res) => NifReturnable::as_returned(res, env),
+                Err(err) => NifReturnable::as_returned(err, env),
+            },
+            Err(err) => match err.downcast::<NifReturned>() {
+                Ok(ty) => NifReturned::Term(ty.apply(env)),
+                Err(_) => {
+                    let term = crate::types::atom::nif_panicked().as_c_arg();
+                    NifReturned::Raise(term)
+                }
+            },
+        }
+    }
+}

rustler/src/env.rs

@@ -23,7 +23,7 @@ pub struct Env<'a> {
 
 /// Two environments are equal if they're the same `NIF_ENV` value.
 ///
-/// A `Env<'a>` is equal to a `Env<'b>` iff `'a` and `'b` are the same lifetime.
+/// A `Env<'a>` is equal to a `Env<'b>` if and only if `'a` and `'b` are the same lifetime.
 impl<'a, 'b> PartialEq<Env<'b>> for Env<'a> {
     fn eq(&self, other: &Env<'b>) -> bool {
         self.env == other.env

rustler/src/export.rs

@@ -15,23 +15,16 @@
 /// The third argument is an `Option<fn(env: &Env, load_info: Term) -> bool>`. If this is
 /// `Some`, the function will execute when the NIF is first loaded by the BEAM.
 #[macro_export]
-#[deprecated(since = "0.22.0", note = "Please use `init!` instead.")]
+#[deprecated(since = "0.22.0", note = "Please use `rustler::init!` instead.")]
 macro_rules! rustler_export_nifs {
-    ($name:expr, [$( $exported_nif:tt ),+,], $on_load:expr) => {
-        $crate::init!($name, [$( $exported_nif ),*], $on_load);
-    };
-}
-
-#[macro_export]
-macro_rules! init {
     // Strip trailing comma.
     ($name:expr, [$( $exported_nif:tt ),+,], $on_load:expr) => {
-        $crate::init!($name, [$( $exported_nif ),*], $on_load);
+        $crate::rustler_export_nifs!($name, [$( $exported_nif ),*], $on_load);
     };
     ($name:expr, [$( $exported_nif:tt ),*], $on_load:expr) => {
         static mut NIF_ENTRY: Option<$crate::codegen_runtime::DEF_NIF_ENTRY> = None;
 
-        $crate::init!(internal_platform_init, ({
+        $crate::rustler_export_nifs!(internal_platform_init, ({
             // TODO: If an unwrap ever happens, we will unwind right into C! Fix this!
 
             extern "C" fn nif_load(
@@ -45,7 +38,7 @@ macro_rules! init {
             }
 
             const FUN_ENTRIES: &'static [$crate::codegen_runtime::DEF_NIF_FUNC] = &[
-                $($crate::init!(internal_item_init, $exported_nif)),*
+                $($crate::rustler_export_nifs!(internal_item_init, $exported_nif)),*
             ];
 
             let entry = $crate::codegen_runtime::DEF_NIF_ENTRY {
@@ -69,7 +62,7 @@ macro_rules! init {
     };
 
     (internal_item_init, ($nif_name:expr, $nif_arity:expr, $nif_fun:path)) => {
-        $crate::init!(internal_item_init, ($nif_name, $nif_arity, $nif_fun, $crate::schedule::SchedulerFlags::Normal))
+        $crate::rustler_export_nifs!(internal_item_init, ($nif_name, $nif_arity, $nif_fun, $crate::schedule::SchedulerFlags::Normal))
     };
     (internal_item_init, ($nif_name:expr, $nif_arity:expr, $nif_fun:path, $nif_flag:expr)) => {
         $crate::codegen_runtime::DEF_NIF_FUNC {
@@ -82,12 +75,9 @@ macro_rules! init {
                     argv: *const $crate::codegen_runtime::NIF_TERM)
                     -> $crate::codegen_runtime::NIF_TERM {
                         unsafe {
-                            $crate::init!(
+                            $crate::rustler_export_nifs!(
                                 internal_handle_nif_call, ($nif_fun, $nif_arity, env, argc, argv))
                         }
-                    //unsafe {
-                    //    $crate::codegen_runtime::handle_nif_call($nif_fun, $nif_arity, env, argc, argv)
-                    //}
                 }
                 nif_func
             },

rustler/src/lib.rs

@@ -25,7 +25,8 @@
 #[macro_use(enif_snprintf)]
 extern crate rustler_sys;
 
-mod wrapper;
+#[doc(hidden)]
+pub mod wrapper;
 
 #[doc(hidden)]
 pub mod codegen_runtime;
@@ -62,7 +63,13 @@ pub use crate::error::Error;
 pub mod r#return;
 pub use crate::r#return::Return;
 
+#[doc(hidden)]
+mod nif;
+pub use nif::Nif;
+
 pub type NifResult<T> = Result<T, Error>;
 
 #[cfg(feature = "derive")]
-pub use rustler_codegen::{NifMap, NifRecord, NifStruct, NifTuple, NifUnitEnum, NifUntaggedEnum};
+pub use rustler_codegen::{
+    init, nif, NifMap, NifRecord, NifStruct, NifTuple, NifUnitEnum, NifUntaggedEnum,
+};

rustler/src/nif.rs

@@ -0,0 +1,13 @@
+use crate::codegen_runtime::{c_int, DEF_NIF_FUNC, NIF_ENV, NIF_TERM};
+
+pub trait Nif {
+    const NAME: *const u8;
+    const ARITY: u32;
+    const FLAGS: u32;
+    const FUNC: DEF_NIF_FUNC;
+    const RAW_FUNC: unsafe extern "C" fn(
+        nif_env: NIF_ENV,
+        argc: c_int,
+        argv: *const NIF_TERM,
+    ) -> NIF_TERM;
+}

rustler/src/return.rs

@@ -6,6 +6,7 @@ pub enum Return<'a> {
     Term(Term<'a>),
     Error(Error),
 }
+
 unsafe impl<'b> NifReturnable for Return<'b> {
     unsafe fn as_returned(self, env: Env) -> NifReturned {
         match self {

rustler/src/types/atom.rs

@@ -269,6 +269,9 @@ macro_rules! rustler_atoms {
 }
 
 atoms! {
+    /// The `nif_panicked` atom.
+    nif_panicked,
+
     /// The `nil` atom.
     nil,
 

rustler_codegen/Cargo.toml

@@ -12,7 +12,7 @@ name = "rustler_codegen"
 proc_macro = true
 
 [dependencies]
-syn = { version = "1.0.5", features = ["derive"] }
+syn = { version = "1.0.5", features = ["full", "extra-traits"] }
 quote = "1.0"
 heck = "0.3"
 proc-macro2 = "1.0"

rustler_codegen/src/init.rs

@@ -0,0 +1,112 @@
+use proc_macro2::{Span, TokenStream};
+use quote::quote;
+use syn::parse::{Parse, ParseStream};
+use syn::punctuated::Punctuated;
+use syn::token::Comma;
+use syn::{Expr, Ident, Result, Token};
+
+#[derive(Debug)]
+pub struct InitMacroInput {
+    module: syn::Lit,
+    funcs: syn::ExprArray,
+    on_load: Option<TokenStream>,
+}
+
+impl Parse for InitMacroInput {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let module = syn::Lit::parse(input)?;
+        let _comma = <syn::Token![,]>::parse(input)?;
+        let funcs = syn::ExprArray::parse(input)?;
+        let on_load = if input.peek(Token![,]) {
+            let _comma = <Token![,]>::parse(input)?;
+            Some(TokenStream::parse(input)?)
+        } else {
+            let none = Ident::new("None", Span::call_site());
+            Some(quote!(#none))
+        };
+
+        Ok(InitMacroInput {
+            module,
+            funcs,
+            on_load,
+        })
+    }
+}
+
+impl Into<proc_macro2::TokenStream> for InitMacroInput {
+    fn into(self) -> proc_macro2::TokenStream {
+        let name = self.module;
+        let num_of_funcs = self.funcs.elems.len();
+        let funcs = nif_funcs(self.funcs.elems);
+        let on_load = self.on_load;
+
+        let inner = quote! {
+            static mut NIF_ENTRY: Option<rustler::codegen_runtime::DEF_NIF_ENTRY> = None;
+            use rustler::Nif;
+
+            let entry = rustler::codegen_runtime::DEF_NIF_ENTRY {
+                major: rustler::codegen_runtime::NIF_MAJOR_VERSION,
+                minor: rustler::codegen_runtime::NIF_MINOR_VERSION,
+                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(#on_load, env, load_info)
+                        }
+                    }
+                    Some(nif_load)
+                },
+                reload: None,
+                upgrade: None,
+                unload: None,
+                vm_variant: b"beam.vanilla\0".as_ptr(),
+                options: 0,
+                sizeof_ErlNifResourceTypeInit: rustler::codegen_runtime::get_nif_resource_type_init_size(),
+            };
+
+            unsafe {
+                NIF_ENTRY = Some(entry);
+                NIF_ENTRY.as_ref().unwrap()
+            }
+        };
+
+        quote! {
+            #[cfg(unix)]
+            #[no_mangle]
+            pub extern "C" fn nif_init() -> *const rustler::codegen_runtime::DEF_NIF_ENTRY {
+                #inner
+            }
+
+            #[cfg(windows)]
+            #[no_mangle]
+            pub extern "C" fn nif_init(callbacks: *mut rustler::codegen_runtime::TWinDynNifCallbacks) -> *const rustler::codegen_runtime::DEF_NIF_ENTRY {
+                unsafe {
+                    rustler::codegen_runtime::WIN_DYN_NIF_CALLBACKS = Some(*callbacks);
+                }
+
+                #inner
+            }
+        }
+    }
+}
+
+fn nif_funcs(funcs: Punctuated<Expr, Comma>) -> TokenStream {
+    let mut tokens = TokenStream::new();
+
+    for func in funcs.iter() {
+        if let Expr::Path(_) = *func {
+            tokens.extend(quote!(#func::FUNC,));
+        } else {
+            panic!("Expected an expression, found: {}", stringify!(func));
+        }
+    }
+
+    tokens
+}