cargo fmt, adds nif scheduling type and macro

Author: sunny-g

rustler/Cargo.toml

@@ -7,20 +7,6 @@ authors = ["Hansihe <me@hansihe.com>"]
 license = "MIT/Apache-2.0"
 edition = "2018"
 
-[features]
-default = ["derive"]
-derive = ["rustler_codegen"]
-alternative_nif_init_name = []
-
-[dependencies]
-rustler_codegen = { path = "../rustler_codegen", optional = true }
-erl_nif_sys = { path = "../erl_nif_sys" }
-lazy_static = "1.2"
-
-[build-dependencies]
-lazy_static = "1.2"
-which = "2"
-
 [package.metadata.release]
 
 [[package.metadata.release.pre-release-replacements]]
@@ -37,3 +23,17 @@ replace = "def rustler_version, do: \"{{version}}\""
 file = "../rustler_codegen/Cargo.toml"
 search = "# rustler_codegen version\nversion = \"[^\"]+\""
 replace = "# rustler_codegen version\nversion = \"{{version}}\""
+
+[features]
+default = ["derive"]
+derive = ["rustler_codegen"]
+alternative_nif_init_name = []
+
+[dependencies]
+erl_nif_sys = { path = "../erl_nif_sys" }
+lazy_static = "1.2"
+rustler_codegen = { path = "../rustler_codegen", optional = true }
+
+[build-dependencies]
+lazy_static = "1.2"
+which = "2"

rustler/src/codegen_runtime.rs

@@ -21,6 +21,7 @@ unsafe impl<T> NifReturnable for T
 where
     T: crate::Encoder,
 {
+    #[inline]
     unsafe fn as_returned(self, env: Env) -> NifReturned {
         NifReturned::Term(self.encode(env).as_c_arg())
     }
@@ -31,7 +32,7 @@ where
 {
     unsafe fn as_returned(self, env: Env) -> NifReturned {
         match self {
-            Ok(inner) => NifReturned::Term(inner.encode(env).as_c_arg()),
+            Ok(inner) => inner.as_returned(env),
             Err(inner) => inner.as_returned(env),
         }
     }

rustler/src/export.rs

@@ -21,8 +21,31 @@ macro_rules! rustler_export_nifs {
         $crate::rustler_export_nifs!($name, [$( $exported_nif ),*], $on_load);
     };
     ($name:expr, [$( $exported_nif:tt ),*], $on_load:expr) => {
+        type ExternNifFuncMap = std::collections::HashMap<&'static str, ExternNifFunc>;
+        type ExternNifFunc = unsafe extern "C" fn(env: $crate::codegen_runtime::NIF_ENV,
+            argc: $crate::codegen_runtime::c_int,
+            argv: *const $crate::codegen_runtime::NIF_TERM)
+            -> $crate::codegen_runtime::NIF_TERM;
+
         static mut NIF_ENTRY: Option<$crate::codegen_runtime::DEF_NIF_ENTRY> = None;
 
+        const FUN_ENTRIES: &'static [$crate::codegen_runtime::DEF_NIF_FUNC] = &[
+            $($crate::rustler_export_nifs!(internal_item_init, $exported_nif)),*
+        ];
+
+        lazy_static::lazy_static!{
+            pub static ref EXTERN_NIF_MAP: ExternNifFuncMap = {
+                let mut map: ExternNifFuncMap = std::collections::HashMap::new();
+                for nif_func_entry in FUN_ENTRIES {
+                    let name = unsafe {
+                        std::ffi::CStr::from_ptr(nif_func_entry.name as *const i8).to_str().unwrap()
+                    };
+                    map.insert(name, nif_func_entry.function);
+                }
+                map
+            };
+        }
+
         $crate::rustler_export_nifs!(internal_platform_init, ({
             // TODO: If an unwrap ever happens, we will unwind right into C! Fix this!
 
@@ -36,10 +59,6 @@ macro_rules! rustler_export_nifs {
                 }
             }
 
-            const FUN_ENTRIES: &'static [$crate::codegen_runtime::DEF_NIF_FUNC] = &[
-                $($crate::rustler_export_nifs!(internal_item_init, $exported_nif)),*
-            ];
-
             let entry = $crate::codegen_runtime::DEF_NIF_ENTRY {
                 major: $crate::codegen_runtime::NIF_MAJOR_VERSION,
                 minor: $crate::codegen_runtime::NIF_MINOR_VERSION,

rustler/src/schedule.rs

@@ -1,5 +1,9 @@
-use crate::wrapper::ErlNifTaskFlags;
-use crate::Env;
+use crate::{
+    codegen_runtime::{NifReturnable, NifReturned},
+    wrapper::{ErlNifTaskFlags, NIF_ENV, NIF_TERM},
+    Env,
+};
+use std::ffi::CString;
 
 pub enum SchedulerFlags {
     Normal = ErlNifTaskFlags::ERL_NIF_NORMAL_JOB as isize,
@@ -11,3 +15,89 @@ pub fn consume_timeslice(env: Env, percent: i32) -> bool {
     let success = unsafe { erl_nif_sys::enif_consume_timeslice(env.as_c_arg(), percent) };
     success == 1
 }
+
+/// An `Either`-like type representing either the final return value of a NIF, or the parameters required to schedule a future NIF call.
+///
+/// See `schedule_nif!` for a convenience wrapper around this type.
+pub enum NifScheduleResult<T> {
+    Final(T),
+    Continue {
+        fun_name: &'static str,
+        fun: unsafe extern "C" fn(NIF_ENV, i32, *const NIF_TERM) -> NIF_TERM,
+        args: Vec<T>,
+        flags: SchedulerFlags,
+    },
+}
+
+impl<T> From<T> for NifScheduleResult<T>
+where
+    T: crate::Encoder,
+{
+    #[inline]
+    fn from(t: T) -> Self {
+        NifScheduleResult::Final(t)
+    }
+}
+
+unsafe impl<T> NifReturnable for Result<NifScheduleResult<T>, crate::error::Error>
+where
+    T: crate::Encoder,
+{
+    unsafe fn as_returned(self, env: Env) -> NifReturned {
+        match self {
+            Err(inner) => inner.as_returned(env),
+            Ok(inner) => match inner {
+                NifScheduleResult::Final(inner) => inner.as_returned(env),
+                NifScheduleResult::Continue {
+                    fun_name,
+                    fun,
+                    flags,
+                    args,
+                } => NifReturned::Reschedule {
+                    fun_name: CString::new(fun_name).unwrap(),
+                    fun,
+                    flags,
+                    args: args.iter().map(|arg| arg.encode(env).as_c_arg()).collect(),
+                },
+            },
+        }
+    }
+}
+
+/// Convenience macro for scheduling a NIF call.
+///
+/// NOTE: can only be used when the `static ref EXTERN_NIF_MAP` created by `rustler_export_nifs!` macro is in scope.
+///
+/// ## Example:
+/// ```
+/// rustler_export_nifs!(
+///     "Elixir.RustlerTest",
+///     ("factorial", 2, test_schedule::factorial),
+///     None,
+/// )
+///
+/// pub fn factorial<'a>(_env: Env<'a>, args: &[Term<'a>]) -> NifResult<NifScheduleResult<u32>> {
+///     let input: u32 = args[0].decode()?;
+///     let result: u32 = args[1].decode::<Option<u32>>()?.unwrap_or(1);
+///     if input == 0 {
+///         return Ok(result.into());
+///     }
+
+///     return schedule_nif!("factorial", vec![input - 1, result * input]);
+/// }
+/// ```
+#[macro_export]
+macro_rules! schedule_nif {
+    ($fun_name:expr, $args:expr) => {
+        schedule_nif!($fun_name, $args, $crate::schedule::SchedulerFlags::Normal)
+    };
+    ($fun_name:expr, $args:expr, $flags:expr) => {{
+        let func = EXTERN_NIF_MAP.get($fun_name).unwrap();
+        Ok($crate::schedule::NifScheduleResult::Continue {
+            fun_name: $fun_name,
+            fun: *func,
+            args: $args,
+            flags: $flags,
+        })
+    }};
+}

rustler/src/term.rs

@@ -1,7 +1,7 @@
 use crate::types::binary::OwnedBinary;
 use crate::wrapper::env::term_to_binary;
 use crate::wrapper::NIF_TERM;
-use crate::{Decoder, Env, NifResult, Binary};
+use crate::{Binary, Decoder, Env, NifResult};
 use std::cmp::Ordering;
 use std::fmt::{self, Debug};
 

rustler/src/thread.rs

@@ -1,8 +1,8 @@
 use crate::env::OwnedEnv;
-use std::panic;
-use std::thread;
 use crate::types::atom::Atom;
 use crate::{Encoder, Env, Term};
+use std::panic;
+use std::thread;
 
 /// A `JobSpawner` is a value that can run Rust code on non-Erlang system threads.
 /// Abstracts away details of thread management for `spawn()`.

rustler_codegen/src/ex_struct.rs

@@ -83,7 +83,11 @@ pub fn gen_decoder(
             let ident = field.ident.as_ref().unwrap();
             let ident_str = ident.to_string();
             let atom_fun = Ident::new(&format!("atom_{}", ident_str), Span::call_site());
-            let error_message = format!("Could not decode field :{} on %{}{{}}", ident.to_string(), struct_name.to_string());
+            let error_message = format!(
+                "Could not decode field :{} on %{}{{}}",
+                ident.to_string(),
+                struct_name.to_string()
+            );
             quote! {
                 #ident: match ::rustler::Decoder::decode(term.map_get(#atom_fun().encode(env))?) {
                     Err(_) => return Err(::rustler::Error::RaiseTerm(Box::new(#error_message))),

rustler_tests/lib/rustler_test.ex

@@ -63,4 +63,6 @@ defmodule RustlerTest do
   def dirty_cpu(), do: err()
 
   def sum_range(_), do: err()
+
+  def scheduled_fac(_, _ \\ nil), do: err()
 end

rustler_tests/native/rustler_test/src/lib.rs

@@ -11,6 +11,7 @@ mod test_map;
 mod test_primitives;
 mod test_range;
 mod test_resource;
+mod test_schedule;
 mod test_term;
 mod test_thread;
 
@@ -89,6 +90,7 @@ rustler::rustler_export_nifs!(
         ),
         ("dirty_io", 0, test_dirty::dirty_io, SchedulerFlags::DirtyIo),
         ("sum_range", 1, test_range::sum_range),
+        ("scheduled_fac", 2, test_schedule::scheduled_fac),
     ],
     Some(on_load)
 );

rustler_tests/native/rustler_test/src/test_schedule.rs

@@ -0,0 +1,14 @@
+use super::EXTERN_NIF_MAP;
+use rustler::schedule::NifScheduleResult;
+use rustler::schedule_nif;
+use rustler::{Env, NifResult, Term};
+
+pub fn scheduled_fac<'a>(_env: Env<'a>, args: &[Term<'a>]) -> NifResult<NifScheduleResult<u32>> {
+    let input: u32 = args[0].decode()?;
+    let result: u32 = args[1].decode::<Option<u32>>()?.unwrap_or(1);
+    if input == 0 {
+        return Ok(result.into());
+    }
+
+    return schedule_nif!("scheduled_fac", vec![input - 1, result * input]);
+}