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

use std::cell::Cell;
use std::ffi::c_void;
use std::ops::{Deref, DerefMut};
use std::rc::Rc;

use lazy_static::lazy_static;

use crate::{
  bindgen_runtime::{PersistedSingleThreadHashMap, ToNapiValue},
  check_status, Env, Error, Result, Status,
};

type RefInformation = (
  *mut c_void,
  crate::sys::napi_ref,
  *const Cell<*mut dyn FnOnce()>,
);

lazy_static! {
  pub(crate) static ref REFERENCE_MAP: PersistedSingleThreadHashMap<*mut c_void, RefInformation> =
    Default::default();
}

/// ### Experimental feature
///
/// Create a `reference` from `Class` instance.
/// Unref the `Reference` when the `Reference` is dropped.
pub struct Reference<T: 'static> {
  raw: *mut T,
  napi_ref: crate::sys::napi_ref,
  env: *mut c_void,
  finalize_callbacks: Rc<Cell<*mut dyn FnOnce()>>,
}

unsafe impl<T: Send> Send for Reference<T> {}
unsafe impl<T: Sync> Sync for Reference<T> {}

impl<T> Drop for Reference<T> {
  fn drop(&mut self) {
    let rc_strong_count = Rc::strong_count(&self.finalize_callbacks);
    let mut ref_count = 0;
    // If Rc strong count == 1, then the referenced object is dropped on GC
    // It would happen when the process is exiting
    // In general, the `drop` of the `Reference` would happen first
    if rc_strong_count > 1 {
      let status = unsafe {
        crate::sys::napi_reference_unref(
          self.env as crate::sys::napi_env,
          self.napi_ref,
          &mut ref_count,
        )
      };
      debug_assert!(
        status == crate::sys::Status::napi_ok,
        "Reference unref failed, status code: {}",
        crate::Status::from(status)
      );
    };
  }
}

impl<T: 'static> Reference<T> {
  #[doc(hidden)]
  pub fn add_ref(t: *mut c_void, value: RefInformation) {
    REFERENCE_MAP.borrow_mut(|map| {
      map.insert(t, value);
    });
  }

  #[doc(hidden)]
  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.borrow_mut(|map| map.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,
        napi_ref,
        env: env as *mut c_void,
        finalize_callbacks,
      })
    } else {
      Err(Error::new(
        Status::InvalidArg,
        format!("Class for Type {:?} not found", t),
      ))
    }
  }
}

impl<T: 'static> ToNapiValue for Reference<T> {
  unsafe fn to_napi_value(env: crate::sys::napi_env, val: Self) -> Result<crate::sys::napi_value> {
    let mut result = std::ptr::null_mut();
    check_status!(
      unsafe { crate::sys::napi_get_reference_value(env, val.napi_ref, &mut result) },
      "Failed to get reference value"
    )?;
    Ok(result)
  }
}

impl<T: 'static> Reference<T> {
  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: s_ptr,
      owner: self,
    })
  }
}

impl<T: 'static> Deref for Reference<T> {
  type Target = T;

  fn deref(&self) -> &Self::Target {
    unsafe { Box::leak(Box::from_raw(self.raw)) }
  }
}

impl<T: 'static> DerefMut for Reference<T> {
  fn deref_mut(&mut self) -> &mut Self::Target {
    unsafe { Box::leak(Box::from_raw(self.raw)) }
  }
}

/// ### Experimental feature
///
/// Create a `SharedReference` from an existed `Reference`.
pub struct SharedReference<T: 'static, S: 'static> {
  raw: *mut S,
  owner: Reference<T>,
}

unsafe impl<T: Send, S: Send> Send for SharedReference<T, S> {}
unsafe impl<T: Sync, S: Sync> Sync for SharedReference<T, S> {}

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)?,
    })
  }

  pub fn clone_owner(&self, env: Env) -> Result<Reference<T>> {
    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,
      owner: self.owner,
    })
  }
}

impl<T: 'static, S: 'static> ToNapiValue for SharedReference<T, S> {
  unsafe fn to_napi_value(env: crate::sys::napi_env, val: Self) -> Result<crate::sys::napi_value> {
    let mut result = std::ptr::null_mut();
    check_status!(
      unsafe { crate::sys::napi_get_reference_value(env, val.owner.napi_ref, &mut result) },
      "Failed to get reference value"
    )?;
    Ok(result)
  }
}

impl<T: 'static, S: 'static> Deref for SharedReference<T, S> {
  type Target = S;

  fn deref(&self) -> &Self::Target {
    unsafe { Box::leak(Box::from_raw(self.raw)) }
  }
}

impl<T: 'static, S: 'static> DerefMut for SharedReference<T, S> {
  fn deref_mut(&mut self) -> &mut Self::Target {
    unsafe { Box::leak(Box::from_raw(self.raw)) }
  }
}