yielding api

rustler/src/codegen_runtime.rs

@@ -1,6 +1,5 @@
 //! Functions used by runtime generated code. Should not be used.
 
-use std::ffi::CString;
 use std::fmt;
 
 use crate::{Encoder, Env, OwnedBinary, Term};
@@ -51,9 +50,9 @@ pub enum NifReturned {
     Raise(NIF_TERM),
     BadArg,
     Reschedule {
-        fun_name: CString,
+        fun_name: *const c_char,
         flags: crate::schedule::SchedulerFlags,
-        fun: unsafe extern "C" fn(NIF_ENV, i32, *const NIF_TERM) -> NIF_TERM,
+        fun: crate::nif::RawFunc,
         args: Vec<NIF_TERM>,
     },
 }
@@ -73,7 +72,7 @@ impl NifReturned {
                 args,
             } => rustler_sys::enif_schedule_nif(
                 env.as_c_arg(),
-                fun_name.as_ptr() as *const c_char,
+                fun_name,
                 flags as i32,
                 fun,
                 args.len() as i32,

rustler/src/lib.rs

@@ -62,6 +62,9 @@ pub use crate::error::Error;
 pub mod r#return;
 pub use crate::r#return::Return;
 
+pub mod r#yield;
+pub use crate::r#yield::{Branch, Yield};
+
 #[doc(hidden)]
 mod nif;
 pub use nif::Nif;

rustler/src/nif.rs

@@ -1,13 +1,12 @@
 use crate::codegen_runtime::{c_char, c_int, c_uint, DEF_NIF_FUNC, NIF_ENV, NIF_TERM};
 
+pub(crate) type RawFunc =
+    unsafe extern "C" fn(nif_env: NIF_ENV, argc: c_int, argv: *const NIF_TERM) -> NIF_TERM;
+
 pub trait Nif {
     const NAME: *const c_char;
     const ARITY: c_uint;
     const FLAGS: c_uint;
     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;
+    const RAW_FUNC: RawFunc;
 }

rustler/src/yield.rs

@@ -0,0 +1,77 @@
+use crate::{
+    codegen_runtime::{NifReturnable, NifReturned},
+    nif::RawFunc,
+    schedule::SchedulerFlags,
+    Env, Nif, Term,
+};
+
+/// If a function needs to break up long-running work, it can return
+/// a `Yield`. The runtime will resume execution by calling the provided
+/// NIF.
+///
+/// ```rust
+/// use rustler::{Encoder, Env, Yield};
+///
+/// #[rustler::nif]
+/// fn start_expensive_work(env: Env) -> Yield {
+///     let args = vec![42_i32.encode(env)];
+///     Yield::to(perform_expensive_work, args)
+/// }
+///
+/// #[rustler::nif]
+/// fn perform_expensive_work(env: Env, value: i32) {
+///     assert_eq!(value, 42);
+/// }
+/// ```
+pub struct Yield<'a> {
+    fun_name: *const i8,
+    fun: RawFunc,
+    flags: SchedulerFlags,
+    args: Vec<Term<'a>>,
+}
+
+impl<'a> Yield<'a> {
+    /// Yield back to the runtime with the provided nif.
+    pub fn to<N: Nif>(nif: N, args: Vec<Term<'a>>) -> Self {
+        Self::to_with_flags(nif, SchedulerFlags::Normal, args)
+    }
+
+    /// Yield back to the runtime with the provided function and scheduler flags.
+    pub fn to_with_flags<N: Nif>(nif: N, flags: SchedulerFlags, args: Vec<Term<'a>>) -> Self {
+        let _ = nif;
+        Self {
+            fun_name: N::NAME,
+            fun: N::RAW_FUNC,
+            flags,
+            args,
+        }
+    }
+}
+
+unsafe impl<'a> NifReturnable for Yield<'a> {
+    unsafe fn into_returned(self, _env: Env) -> NifReturned {
+        NifReturned::Reschedule {
+            fun_name: self.fun_name,
+            fun: self.fun,
+            flags: self.flags,
+            args: self.args.into_iter().map(|term| term.as_c_arg()).collect(),
+        }
+    }
+}
+
+pub enum Branch<'a, T> {
+    Yield(Yield<'a>),
+    Stop(T),
+}
+
+unsafe impl<'a, T> NifReturnable for Branch<'a, T>
+where
+    T: NifReturnable,
+{
+    unsafe fn into_returned(self, env: Env) -> NifReturned {
+        match self {
+            Self::Yield(inner) => inner.into_returned(env),
+            Self::Stop(inner) => inner.into_returned(env),
+        }
+    }
+}

rustler_tests/lib/rustler_test.ex

@@ -111,6 +111,7 @@ defmodule RustlerTest do
 
   def dirty_io(), do: err()
   def dirty_cpu(), do: err()
+  def yields(), do: err()
 
   def sum_range(_), do: err()
 

rustler_tests/native/rustler_test/src/lib.rs

@@ -82,6 +82,7 @@ rustler::init!(
         test_codegen::tuplestruct_record_echo,
         test_dirty::dirty_cpu,
         test_dirty::dirty_io,
+        test_dirty::yields,
         test_range::sum_range,
         test_error::bad_arg_error,
         test_error::atom_str_error,

rustler_tests/native/rustler_test/src/test_dirty.rs

@@ -1,4 +1,4 @@
-use rustler::Atom;
+use rustler::{Atom, Encoder, Env, Yield};
 use std::time::Duration;
 
 mod atoms {
@@ -22,3 +22,15 @@ pub fn dirty_io() -> Atom {
 
     atoms::ok()
 }
+
+#[rustler::nif]
+fn yield_resume(input: i32) -> Atom {
+    assert_eq!(input, 42);
+    atoms::ok()
+}
+
+#[rustler::nif]
+pub fn yields(env: Env) -> Yield {
+    let term = 42.encode(env);
+    Yield::to(yield_resume, vec![term])
+}

rustler_tests/test/dirty_test.exs

@@ -8,4 +8,8 @@ defmodule RustlerTest.DirtyTest do
   test "dirty cpu" do
     RustlerTest.dirty_cpu()
   end
+
+  test "yields" do
+    RustlerTest.yields()
+  end
 end