feat(napi): redesign the Reference API (#1118)

- Reference now is inject by #[napi] macro
- Class Reference and underlaying data now have the same lifetime

crates/backend/src/codegen/fn.rs

@@ -100,10 +100,16 @@ impl NapiFn {
     if let Some(parent) = &self.parent {
       match self.fn_self {
         Some(FnSelf::Ref) => {
-          arg_conversions.push(quote! { let this = cb.unwrap_borrow::<#parent>()?; });
+          arg_conversions.push(quote! {
+            let this_ptr = unsafe { cb.unwrap_raw::<#parent>()? };
+            let this: &#parent = Box::leak(Box::from_raw(this_ptr));
+          });
         }
         Some(FnSelf::MutRef) => {
-          arg_conversions.push(quote! { let this = cb.unwrap_borrow_mut::<#parent>()?; });
+          arg_conversions.push(quote! {
+            let this_ptr = unsafe { cb.unwrap_raw::<#parent>()? };
+            let this: &mut #parent = Box::leak(Box::from_raw(this_ptr));
+          });
         }
         _ => {}
       };
@@ -120,6 +126,30 @@ impl NapiFn {
             args.push(quote! { napi::bindgen_prelude::Env::from(env) });
             skipped_arg_count += 1;
           } else {
+            if self.parent.is_some() {
+              if let syn::Type::Path(path) = path.ty.as_ref() {
+                if let Some(p) = path.path.segments.first() {
+                  if p.ident == "Reference" {
+                    if let syn::PathArguments::AngleBracketed(
+                      syn::AngleBracketedGenericArguments { args: angle_bracketed_args, .. },
+                    ) = &p.arguments
+                    {
+                      if let Some(syn::GenericArgument::Type(syn::Type::Path(path))) = angle_bracketed_args.first() {
+                          if let Some(p) = path.path.segments.first() {
+                            if p.ident == *self.parent.as_ref().unwrap() {
+                              args.push(
+                                quote! { napi::bindgen_prelude::Reference::from_value_ptr(this_ptr as *mut std::ffi::c_void, env)? },
+                              );
+                              skipped_arg_count += 1;
+                              return;
+                            }
+                          }
+                      }
+                    }
+                  }
+                }
+              }
+            }
             arg_conversions.push(self.gen_ty_arg_conversion(&ident, i, path));
             args.push(quote! { #ident });
           }

crates/backend/src/codegen/struct.rs

@@ -240,6 +240,8 @@ impl NapiStruct {
             )?;
             let wrapped_value = Box::into_raw(Box::new(val)) as *mut std::ffi::c_void;
             let mut object_ref = std::ptr::null_mut();
+            let initial_finalize: Box<dyn FnOnce()> = Box::new(|| {});
+            let finalize_callbacks_ptr = std::rc::Rc::into_raw(std::rc::Rc::new(std::cell::Cell::new(Box::into_raw(initial_finalize))));
             napi::check_status!(
               napi::sys::napi_wrap(
                 env,
@@ -252,7 +254,7 @@ impl NapiStruct {
               "Failed to wrap native object of class `{}`",
               #js_name_str
             )?;
-            napi::bindgen_prelude::Reference::<#name>::add_ref(std::any::TypeId::of::<#name>(), (wrapped_value, env, object_ref));
+            napi::bindgen_prelude::Reference::<#name>::add_ref(wrapped_value, (wrapped_value, object_ref, finalize_callbacks_ptr));
             napi::bindgen_prelude::___CALL_FROM_FACTORY.with(|inner| inner.store(false, std::sync::atomic::Ordering::Relaxed));
             Ok(result)
           } else {
@@ -262,12 +264,6 @@ impl NapiStruct {
           }
         }
       }
-
-      impl #name {
-        pub fn create_reference(&self) -> napi::bindgen_prelude::Result<napi::bindgen_prelude::Reference<#name>> {
-          napi::bindgen_prelude::Reference::<#name>::from_typeid(std::any::TypeId::of::<#name>())
-        }
-      }
     }
   }
 

crates/backend/src/typegen/fn.rs

@@ -116,7 +116,8 @@ impl NapiFn {
         .iter()
         .filter_map(|arg| match arg {
           crate::NapiFnArgKind::PatType(path) => {
-            if path.ty.to_token_stream().to_string() == "Env" {
+            let ty_string = path.ty.to_token_stream().to_string();
+            if ty_string == "Env" || ty_string.replace(' ', "").starts_with("Reference<") {
               return None;
             }
             let mut path = path.clone();

crates/napi/src/bindgen_runtime/callback_info.rs

@@ -1,5 +1,7 @@
+use std::cell::Cell;
 use std::ffi::c_void;
 use std::ptr;
+use std::rc::Rc;
 use std::sync::atomic::{AtomicBool, Ordering};
 
 use crate::{bindgen_prelude::*, check_status, sys, Result};
@@ -69,6 +71,8 @@ impl<const N: usize> CallbackInfo<N> {
     let this = self.this();
     let value_ref = Box::into_raw(obj) as *mut c_void;
     let mut object_ref = ptr::null_mut();
+    let initial_finalize: Box<dyn FnOnce()> = Box::new(|| {});
+    let finalize_callbacks_ptr = Rc::into_raw(Rc::new(Cell::new(Box::into_raw(initial_finalize))));
     unsafe {
       check_status!(
         sys::napi_wrap(
@@ -84,15 +88,11 @@ impl<const N: usize> CallbackInfo<N> {
       )?;
     };
 
-    Reference::<T>::add_ref(
-      std::any::TypeId::of::<T>(),
-      (value_ref, self.env, object_ref),
-    );
+    Reference::<T>::add_ref(value_ref, (value_ref, object_ref, finalize_callbacks_ptr));
     Ok(this)
   }
 
   pub fn factory<T: 'static>(&self, js_name: &str, obj: T) -> Result<sys::napi_value> {
-    let obj = Box::new(obj);
     let this = self.this();
     let mut instance = ptr::null_mut();
     unsafe {
@@ -106,7 +106,10 @@ impl<const N: usize> CallbackInfo<N> {
         sys::napi_throw(self.env, exception);
         return Ok(ptr::null_mut());
       }
-
+      let obj = Box::new(obj);
+      let initial_finalize: Box<dyn FnOnce()> = Box::new(|| {});
+      let finalize_callbacks_ptr =
+        Rc::into_raw(Rc::new(Cell::new(Box::into_raw(initial_finalize))));
       let mut object_ref = ptr::null_mut();
       let value_ref = Box::into_raw(obj) as *mut c_void;
       check_status!(
@@ -122,10 +125,7 @@ impl<const N: usize> CallbackInfo<N> {
         js_name,
       )?;
 
-      Reference::<T>::add_ref(
-        std::any::TypeId::of::<T>(),
-        (value_ref, self.env, object_ref),
-      );
+      Reference::<T>::add_ref(value_ref, (value_ref, object_ref, finalize_callbacks_ptr));
     };
 
     Ok(instance)
@@ -134,6 +134,23 @@ impl<const N: usize> CallbackInfo<N> {
   pub fn unwrap_borrow_mut<T>(&mut self) -> Result<&'static mut T>
   where
     T: FromNapiMutRef + TypeName,
+  {
+    unsafe { self.unwrap_raw::<T>() }.map(|raw| Box::leak(unsafe { Box::from_raw(raw) }))
+  }
+
+  pub fn unwrap_borrow<T>(&mut self) -> Result<&'static T>
+  where
+    T: FromNapiRef + TypeName,
+  {
+    unsafe { self.unwrap_raw::<T>() }
+      .map(|raw| Box::leak(unsafe { Box::from_raw(raw) }) as &'static T)
+  }
+
+  #[doc(hidden)]
+  #[inline]
+  pub unsafe fn unwrap_raw<T>(&mut self) -> Result<*mut T>
+  where
+    T: TypeName,
   {
     let mut wrapped_val: *mut c_void = std::ptr::null_mut();
 
@@ -144,24 +161,7 @@ impl<const N: usize> CallbackInfo<N> {
         T::type_name(),
       )?;
 
-      Ok(&mut *(wrapped_val as *mut T))
-    }
-  }
-
-  pub fn unwrap_borrow<T>(&mut self) -> Result<&'static T>
-  where
-    T: FromNapiRef + TypeName,
-  {
-    let mut wrapped_val: *mut c_void = std::ptr::null_mut();
-
-    unsafe {
-      check_status!(
-        sys::napi_unwrap(self.env, self.this, &mut wrapped_val),
-        "Failed to unwrap shared reference of `{}` type from napi value",
-        T::type_name(),
-      )?;
-
-      Ok(&*(wrapped_val as *const T))
+      Ok(wrapped_val as *mut T)
     }
   }
 }

crates/napi/src/bindgen_runtime/js_values/value_ref.rs

@@ -1,15 +1,19 @@
-use std::any::TypeId;
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 use std::ffi::c_void;
 use std::ops::{Deref, DerefMut};
+use std::rc::Rc;
 
-use crate::{check_status, Error, Result, Status};
+use crate::{check_status, Env, Error, Result, Status};
 
-type RefInformation = (*mut c_void, crate::sys::napi_env, crate::sys::napi_ref);
+type RefInformation = (
+  *mut c_void,
+  crate::sys::napi_ref,
+  *const Cell<*mut dyn FnOnce()>,
+);
 
 thread_local! {
-  static REFERENCE_MAP: RefCell<HashMap<TypeId, RefInformation>> = Default::default();
+  pub(crate) static REFERENCE_MAP: RefCell<HashMap<*mut c_void, RefInformation>> = Default::default();
 }
 
 /// ### Experimental feature
@@ -19,7 +23,8 @@ thread_local! {
 pub struct Reference<T> {
   raw: *mut T,
   napi_ref: crate::sys::napi_ref,
-  env: crate::sys::napi_env,
+  env: *mut c_void,
+  finalize_callbacks: Rc<Cell<*mut dyn FnOnce()>>,
 }
 
 unsafe impl<T: Send> Send for Reference<T> {}
@@ -27,7 +32,9 @@ unsafe impl<T: Sync> Sync for Reference<T> {}
 
 impl<T> Drop for Reference<T> {
   fn drop(&mut self) {
-    let status = unsafe { crate::sys::napi_reference_unref(self.env, self.napi_ref, &mut 0) };
+    let status = unsafe {
+      crate::sys::napi_reference_unref(self.env as crate::sys::napi_env, self.napi_ref, &mut 0)
+    };
     debug_assert!(
       status == crate::sys::Status::napi_ok,
       "Reference unref failed"
@@ -37,25 +44,30 @@ impl<T> Drop for Reference<T> {
 
 impl<T> Reference<T> {
   #[doc(hidden)]
-  pub fn add_ref(t: TypeId, value: RefInformation) {
+  pub fn add_ref(t: *mut c_void, value: RefInformation) {
     REFERENCE_MAP.with(|map| {
       map.borrow_mut().insert(t, value);
     });
   }
 
   #[doc(hidden)]
-  pub fn from_typeid(t: TypeId) -> Result<Self> {
-    if let Some((wrapped_value, env, napi_ref)) =
+  pub unsafe fn from_value_ptr(t: *mut c_void, env: crate::sys::napi_env) -> Result<Self> {
+    if let Some((wrapped_value, napi_ref, finalize_callbacks_ptr)) =
       REFERENCE_MAP.with(|map| map.borrow().get(&t).cloned())
     {
       check_status!(
         unsafe { crate::sys::napi_reference_ref(env, napi_ref, &mut 0) },
         "Failed to ref napi reference"
       )?;
+      let finalize_callbacks_raw = unsafe { Rc::from_raw(finalize_callbacks_ptr) };
+      let finalize_callbacks = finalize_callbacks_raw.clone();
+      // Leak the raw finalize callbacks
+      Rc::into_raw(finalize_callbacks_raw);
       Ok(Self {
         raw: wrapped_value as *mut T,
-        env,
         napi_ref,
+        env: env as *mut c_void,
+        finalize_callbacks,
       })
     } else {
       Err(Error::new(
@@ -67,13 +79,36 @@ impl<T> Reference<T> {
 }
 
 impl<T: 'static> Reference<T> {
-  pub fn share_with<S, F: FnOnce(&'static mut T) -> Result<S>>(
+  pub fn clone(&self, env: Env) -> Result<Self> {
+    let mut ref_count = 0;
+    check_status!(
+      unsafe { crate::sys::napi_reference_ref(env.0, self.napi_ref, &mut ref_count) },
+      "Failed to ref napi reference"
+    )?;
+    Ok(Self {
+      raw: self.raw,
+      napi_ref: self.napi_ref,
+      env: env.0 as *mut c_void,
+      finalize_callbacks: self.finalize_callbacks.clone(),
+    })
+  }
+
+  /// Safety to share because caller can provide `Env`
+  pub fn share_with<S: 'static, F: FnOnce(&'static mut T) -> Result<S>>(
     self,
+    #[allow(unused_variables)] env: Env,
     f: F,
   ) -> Result<SharedReference<T, S>> {
     let s = f(Box::leak(unsafe { Box::from_raw(self.raw) }))?;
+    let s_ptr = Box::into_raw(Box::new(s));
+    let prev_drop_fn = unsafe { Box::from_raw(self.finalize_callbacks.get()) };
+    let drop_fn = Box::new(move || {
+      unsafe { Box::from_raw(s_ptr) };
+      prev_drop_fn();
+    });
+    self.finalize_callbacks.set(Box::into_raw(drop_fn));
     Ok(SharedReference {
-      raw: Box::into_raw(Box::new(s)),
+      raw: s_ptr,
       owner: self,
     })
   }
@@ -93,22 +128,6 @@ impl<T> DerefMut for Reference<T> {
   }
 }
 
-impl<T> Clone for Reference<T> {
-  fn clone(&self) -> Self {
-    let mut ref_count = 0;
-    let status = unsafe { crate::sys::napi_reference_ref(self.env, self.napi_ref, &mut ref_count) };
-    debug_assert!(
-      status == crate::sys::Status::napi_ok,
-      "Reference ref failed"
-    );
-    Self {
-      raw: self.raw,
-      napi_ref: self.napi_ref,
-      env: self.env,
-    }
-  }
-}
-
 /// ### Experimental feature
 ///
 /// Create a `SharedReference` from an existed `Reference`.
@@ -120,14 +139,30 @@ pub struct SharedReference<T, S> {
 unsafe impl<T: Send, S: Send> Send for SharedReference<T, S> {}
 unsafe impl<T: Sync, S: Sync> Sync for SharedReference<T, S> {}
 
-impl<T, S: 'static> SharedReference<T, S> {
-  pub fn share_with<U, F: FnOnce(&'static mut S) -> Result<U>>(
+impl<T: 'static, S: 'static> SharedReference<T, S> {
+  pub fn clone(&self, env: Env) -> Result<Self> {
+    Ok(SharedReference {
+      raw: self.raw,
+      owner: self.owner.clone(env)?,
+    })
+  }
+
+  /// Safety to share because caller can provide `Env`
+  pub fn share_with<U: 'static, F: FnOnce(&'static mut S) -> Result<U>>(
     self,
+    #[allow(unused_variables)] env: Env,
     f: F,
   ) -> Result<SharedReference<T, U>> {
     let s = f(Box::leak(unsafe { Box::from_raw(self.raw) }))?;
+    let raw = Box::into_raw(Box::new(s));
+    let prev_drop_fn = unsafe { Box::from_raw(self.owner.finalize_callbacks.get()) };
+    let drop_fn = Box::new(move || {
+      unsafe { Box::from_raw(raw) };
+      prev_drop_fn();
+    });
+    self.owner.finalize_callbacks.set(Box::into_raw(drop_fn));
     Ok(SharedReference {
-      raw: Box::into_raw(Box::new(s)),
+      raw,
       owner: self.owner,
     })
   }
@@ -146,18 +181,3 @@ impl<T, S> DerefMut for SharedReference<T, S> {
     unsafe { Box::leak(Box::from_raw(self.raw)) }
   }
 }
-
-impl<T, S> Clone for SharedReference<T, S> {
-  fn clone(&self) -> Self {
-    let status =
-      unsafe { crate::sys::napi_reference_ref(self.owner.env, self.owner.napi_ref, &mut 0) };
-    debug_assert!(
-      status == crate::sys::Status::napi_ok,
-      "Reference ref failed"
-    );
-    Self {
-      raw: self.raw,
-      owner: self.owner.clone(),
-    }
-  }
-}

crates/napi/src/bindgen_runtime/mod.rs

@@ -1,5 +1,6 @@
 use std::ffi::c_void;
 use std::mem;
+use std::rc::Rc;
 
 pub use callback_info::*;
 pub use ctor::ctor;
@@ -8,6 +9,7 @@ pub use js_values::*;
 pub use module_register::*;
 
 use super::sys;
+use crate::Status;
 
 mod callback_info;
 mod env;
@@ -19,11 +21,24 @@ mod module_register;
 ///
 /// called when node wrapper objects destroyed
 pub unsafe extern "C" fn raw_finalize_unchecked<T>(
-  _env: sys::napi_env,
+  env: sys::napi_env,
   finalize_data: *mut c_void,
   _finalize_hint: *mut c_void,
 ) {
   unsafe { Box::from_raw(finalize_data as *mut T) };
+  if let Some((_, ref_val, finalize_callbacks_ptr)) =
+    REFERENCE_MAP.with(|reference_map| reference_map.borrow_mut().remove(&finalize_data))
+  {
+    let finalize_callbacks_rc = unsafe { Rc::from_raw(finalize_callbacks_ptr) };
+    let finalize = unsafe { Box::from_raw(finalize_callbacks_rc.get()) };
+    finalize();
+    let delete_reference_status = unsafe { sys::napi_delete_reference(env, ref_val) };
+    debug_assert!(
+      delete_reference_status == sys::Status::napi_ok,
+      "Delete reference in finalize callback failed {}",
+      Status::from(delete_reference_status)
+    );
+  }
 }
 
 /// # Safety

examples/napi/src/reference.rs

@@ -35,11 +35,9 @@ impl JsRepo {
   }
 
   #[napi]
-  pub fn remote(&self) -> Result<JsRemote> {
+  pub fn remote(&self, reference: Reference<JsRepo>, env: Env) -> Result<JsRemote> {
     Ok(JsRemote {
-      inner: self
-        .create_reference()?
-        .share_with(|repo| Ok(repo.inner.remote()))?,
+      inner: reference.share_with(env, |repo| Ok(repo.inner.remote()))?,
     })
   }
 }