Merge pull request #1224 from napi-rs/clone-ref

feat(napi): clone reference for TypedArray/Buffer

crates/backend/src/codegen/struct.rs

@@ -10,6 +10,19 @@ use crate::{
 };
 
 static NAPI_IMPL_ID: AtomicU32 = AtomicU32::new(0);
+const TYPED_ARRAY_TYPE: &[&str] = &[
+  "Int8Array",
+  "Uint8Array",
+  "Uint8ClampedArray",
+  "Int16Array",
+  "Uint16Array",
+  "Int32Array",
+  "Uint32Array",
+  "Float32Array",
+  "Float64Array",
+  "BigInt64Array",
+  "BigUint64Array",
+];
 
 // Generate trait implementations for given Struct.
 fn gen_napi_value_map_impl(name: &Ident, to_napi_val_impl: TokenStream) -> TokenStream {
@@ -584,6 +597,30 @@ impl NapiStruct {
       let setter_name = Ident::new(&format!("set_{}", field_name), Span::call_site());
 
       if field.getter {
+        let default_to_napi_value_convert = quote! {
+          let val = obj.#field_ident.to_owned();
+          unsafe { <#ty as napi::bindgen_prelude::ToNapiValue>::to_napi_value(env, val) }
+        };
+        let to_napi_value_convert = if let syn::Type::Path(syn::TypePath {
+          path: syn::Path { segments, .. },
+          ..
+        }) = ty
+        {
+          if let Some(syn::PathSegment { ident, .. }) = segments.last() {
+            if TYPED_ARRAY_TYPE.iter().any(|name| ident == name) || ident == "Buffer" {
+              quote! {
+                let val = &mut obj.#field_ident;
+                unsafe { <&mut #ty as napi::bindgen_prelude::ToNapiValue>::to_napi_value(env, val) }
+              }
+            } else {
+              default_to_napi_value_convert
+            }
+          } else {
+            default_to_napi_value_convert
+          }
+        } else {
+          default_to_napi_value_convert
+        };
         getters_setters.push((
           field.js_name.clone(),
           quote! {
@@ -592,10 +629,9 @@ impl NapiStruct {
               cb: napi::bindgen_prelude::sys::napi_callback_info
             ) -> napi::bindgen_prelude::sys::napi_value {
               napi::bindgen_prelude::CallbackInfo::<0>::new(env, cb, Some(0))
-                .and_then(|mut cb| unsafe { cb.unwrap_borrow::<#struct_name>() })
+                .and_then(|mut cb| unsafe { cb.unwrap_borrow_mut::<#struct_name>() })
                 .and_then(|obj| {
-                  let val = obj.#field_ident.to_owned();
+                  #to_napi_value_convert
-                  unsafe { <#ty as napi::bindgen_prelude::ToNapiValue>::to_napi_value(env, val) }
                 })
                 .unwrap_or_else(|e| {
                   unsafe { napi::bindgen_prelude::JsError::from(e).throw_into(env) };

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

@@ -2,31 +2,87 @@ use std::ffi::c_void;
 use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::ptr;
+use std::sync::Arc;
 
 pub use crate::js_values::TypedArrayType;
 use crate::{check_status, sys, Error, Result, Status};
 
 use super::{FromNapiValue, ToNapiValue, TypeName, ValidateNapiValue};
 
+trait Finalizer {
+  type RustType;
+
+  fn finalizer(&mut self) -> Box<dyn FnOnce(*mut Self::RustType, usize)>;
+
+  fn data_managed_type(&self) -> &DataManagedType;
+
+  fn len(&self) -> &usize;
+
+  fn ref_count(&self) -> usize;
+}
+
 macro_rules! impl_typed_array {
   ($name:ident, $rust_type:ident, $typed_array_type:expr) => {
     pub struct $name {
+      /// Used for `clone_reference` Owned | External TypedArray
+      data_reference: Option<Arc<()>>,
       data: *mut $rust_type,
       length: usize,
       data_managed_type: DataManagedType,
       byte_offset: usize,
+      raw: Option<(crate::sys::napi_ref, crate::sys::napi_env)>,
       finalizer_notify: Box<dyn FnOnce(*mut $rust_type, usize)>,
     }
 
+    unsafe impl Send for $name {}
+
+    impl Finalizer for $name {
+      type RustType = $rust_type;
+
+      fn finalizer(&mut self) -> Box<dyn FnOnce(*mut Self::RustType, usize)> {
+        std::mem::replace(&mut self.finalizer_notify, Box::new($name::noop_finalize))
+      }
+
+      fn data_managed_type(&self) -> &DataManagedType {
+        &self.data_managed_type
+      }
+
+      fn len(&self) -> &usize {
+        &self.length
+      }
+
+      fn ref_count(&self) -> usize {
+        if let Some(inner) = &self.data_reference {
+          Arc::strong_count(inner)
+        } else {
+          0
+        }
+      }
+    }
+
+    impl Drop for $name {
+      fn drop(&mut self) {
+        if let Some((ref_, env)) = self.raw {
+          crate::check_status_or_throw!(
+            env,
+            unsafe { sys::napi_reference_unref(env, ref_, &mut 0) },
+            "Failed to delete Buffer reference in drop"
+          );
+        }
+      }
+    }
+
     impl $name {
       fn noop_finalize(_data: *mut $rust_type, _length: usize) {}
 
       pub fn new(mut data: Vec<$rust_type>) -> Self {
         let ret = $name {
+          data_reference: Some(Arc::new(())),
           data: data.as_mut_ptr(),
           length: data.len(),
           data_managed_type: DataManagedType::Owned,
           byte_offset: 0,
+          raw: None,
           finalizer_notify: Box::new(Self::noop_finalize),
         };
         mem::forget(data);
@@ -39,10 +95,12 @@ macro_rules! impl_typed_array {
       {
         let mut data_copied = data.as_ref().to_vec();
         let ret = $name {
+          data_reference: Some(Arc::new(())),
           data: data_copied.as_mut_ptr(),
           length: data.as_ref().len(),
           data_managed_type: DataManagedType::Owned,
           finalizer_notify: Box::new(Self::noop_finalize),
+          raw: None,
           byte_offset: 0,
         };
         mem::forget(data_copied);
@@ -57,13 +115,40 @@ macro_rules! impl_typed_array {
         F: 'static + FnOnce(*mut $rust_type, usize),
       {
         $name {
+          data_reference: Some(Arc::new(())),
           data,
           length,
           data_managed_type: DataManagedType::External,
           finalizer_notify: Box::new(notify),
+          raw: None,
           byte_offset: 0,
         }
       }
+
+      /// Clone reference, the inner data is not copied nor moved
+      pub fn clone(&mut self, env: crate::Env) -> crate::Result<$name> {
+        let raw = if let Some((ref_, _)) = self.raw {
+          crate::check_status!(
+            unsafe { sys::napi_reference_ref(env.0, ref_, &mut 0) },
+            "Failed to ref Buffer reference in TypedArray::clone"
+          )?;
+          Some((ref_, env.0))
+        } else {
+          None
+        };
+        Ok(Self {
+          data_reference: match self.data_managed_type {
+            DataManagedType::Owned | DataManagedType::External => self.data_reference.clone(),
+            _ => None,
+          },
+          data: self.data,
+          length: self.length,
+          data_managed_type: self.data_managed_type,
+          finalizer_notify: Box::new(Self::noop_finalize),
+          raw,
+          byte_offset: self.byte_offset,
+        })
+      }
     }
 
     impl Deref for $name {
@@ -110,7 +195,7 @@ macro_rules! impl_typed_array {
       unsafe fn validate(
         env: sys::napi_env,
         napi_val: sys::napi_value,
-      ) -> Result<$crate::sys::napi_value> {
+      ) -> Result<crate::sys::napi_value> {
         let mut is_typed_array = false;
         check_status!(
           unsafe { sys::napi_is_typedarray(env, napi_val, &mut is_typed_array) },
@@ -133,6 +218,11 @@ macro_rules! impl_typed_array {
         let mut data = ptr::null_mut();
         let mut array_buffer = ptr::null_mut();
         let mut byte_offset = 0;
+        let mut ref_ = ptr::null_mut();
+        check_status!(
+          unsafe { sys::napi_create_reference(env, napi_val, 1, &mut ref_) },
+          "Failed to create reference from Buffer"
+        )?;
         check_status!(
           unsafe {
             sys::napi_get_typedarray_info(
@@ -154,9 +244,11 @@ macro_rules! impl_typed_array {
           ));
         }
         Ok($name {
+          data_reference: None,
           data: data as *mut $rust_type,
           length,
           byte_offset,
+          raw: Some((ref_, env)),
           data_managed_type: DataManagedType::Vm,
           finalizer_notify: Box::new(Self::noop_finalize),
         })
@@ -164,15 +256,22 @@ macro_rules! impl_typed_array {
     }
 
     impl ToNapiValue for $name {
-      unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+      unsafe fn to_napi_value(env: sys::napi_env, mut val: Self) -> Result<sys::napi_value> {
+        if let Some((ref_, _)) = val.raw {
+          let mut napi_value = std::ptr::null_mut();
+          check_status!(
+            unsafe { sys::napi_get_reference_value(env, ref_, &mut napi_value) },
+            "Failed to delete reference from Buffer"
+          )?;
+          val.raw = None;
+          return Ok(napi_value);
+        }
         let mut arraybuffer_value = ptr::null_mut();
         let ratio = mem::size_of::<$rust_type>() / mem::size_of::<u8>();
         let length = val.length * ratio;
-        let hint_ptr = Box::into_raw(Box::new((
+        let val_data = val.data;
-          val.data_managed_type,
+        let val_length = val.length;
-          val.length,
+        let hint_ptr = Box::into_raw(Box::new(val));
-          val.finalizer_notify,
-        )));
         check_status!(
           if length == 0 {
             // Rust uses 0x1 as the data pointer for empty buffers,
@@ -185,9 +284,9 @@ macro_rules! impl_typed_array {
             unsafe {
               sys::napi_create_external_arraybuffer(
                 env,
-                val.data as *mut c_void,
+                val_data as *mut c_void,
                 length,
-                Some(finalizer::<$rust_type>),
+                Some(finalizer::<$rust_type, $name>),
                 hint_ptr as *mut c_void,
                 &mut arraybuffer_value,
               )
@@ -201,7 +300,7 @@ macro_rules! impl_typed_array {
             sys::napi_create_typedarray(
               env,
               $typed_array_type as i32,
-              val.length,
+              val_length,
               arraybuffer_value,
               0,
               &mut napi_val,
@@ -212,31 +311,64 @@ macro_rules! impl_typed_array {
         Ok(napi_val)
       }
     }
+
+    impl ToNapiValue for &mut $name {
+      unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+        if let Some((ref_, _)) = val.raw {
+          let mut napi_value = std::ptr::null_mut();
+          check_status!(
+            unsafe { sys::napi_get_reference_value(env, ref_, &mut napi_value) },
+            "Failed to delete reference from Buffer"
+          )?;
+          Ok(napi_value)
+        } else {
+          let cloned_value = $name {
+            data_reference: match val.data_managed_type {
+              DataManagedType::Owned | DataManagedType::External => val.data_reference.clone(),
+              _ => None,
+            },
+            data: val.data,
+            length: val.length,
+            data_managed_type: val.data_managed_type,
+            finalizer_notify: Box::new($name::noop_finalize),
+            raw: None,
+            byte_offset: val.byte_offset,
+          };
+          unsafe { ToNapiValue::to_napi_value(env, cloned_value) }
+        }
+      }
+    }
   };
 }
 
-unsafe extern "C" fn finalizer<T>(
+unsafe extern "C" fn finalizer<Data, T: Finalizer<RustType = Data>>(
   _env: sys::napi_env,
   finalize_data: *mut c_void,
   finalize_hint: *mut c_void,
 ) {
-  let (data_managed_type, length, finalizer_notify) = unsafe {
+  let mut data = unsafe { *Box::from_raw(finalize_hint as *mut T) };
-    *Box::from_raw(finalize_hint as *mut (DataManagedType, usize, Box<dyn FnOnce(*mut T, usize)>))
+  let data_managed_type = *data.data_managed_type();
-  };
+  let length = *data.len();
+  let finalizer_notify = data.finalizer();
   match data_managed_type {
     DataManagedType::Vm => {
       // do nothing
     }
     DataManagedType::Owned => {
+      if data.ref_count() == 0 {
         let length = length;
-      unsafe { Vec::from_raw_parts(finalize_data as *mut T, length, length) };
+        unsafe { Vec::from_raw_parts(finalize_data as *mut Data, length, length) };
+      }
     }
     DataManagedType::External => {
-      (finalizer_notify)(finalize_data as *mut T, length);
+      if data.ref_count() == 0 {
+        (finalizer_notify)(finalize_data as *mut Data, length);
+      }
     }
   }
 }
 
+#[derive(PartialEq, Eq, Clone, Copy)]
 enum DataManagedType {
   /// Vm managed data, passed in from JavaScript
   Vm,

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

@@ -4,6 +4,7 @@ use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::ptr;
 use std::slice;
+use std::sync::Arc;
 #[cfg(debug_assertions)]
 use std::sync::Mutex;
 
@@ -19,6 +20,7 @@ thread_local! {
 /// So it is safe to use it in `async fn`, the `&[u8]` under the hood will not be dropped until the `drop` called.
 /// Clone will create a new `Reference` to the same underlying `JavaScript Buffer`.
 pub struct Buffer {
+  data_reference: Option<Arc<()>>,
   inner: &'static mut [u8],
   capacity: usize,
   raw: Option<(sys::napi_ref, sys::napi_env)>,
@@ -40,22 +42,21 @@ unsafe impl Send for Buffer {}
 
 impl Buffer {
   pub fn clone(&mut self, env: &Env) -> Result<Self> {
-    if let Some((ref_, _)) = self.raw {
+    let raw = if let Some((ref_, _)) = self.raw {
       check_status!(
         unsafe { sys::napi_reference_ref(env.0, ref_, &mut 0) },
         "Failed to ref Buffer reference in Buffer::clone"
       )?;
+      Some((ref_, env.0))
+    } else {
+      None
+    };
     Ok(Self {
+      data_reference: self.data_reference.clone(),
       inner: unsafe { slice::from_raw_parts_mut(self.inner.as_mut_ptr(), self.inner.len()) },
       capacity: self.capacity,
-        raw: Some((ref_, env.0)),
+      raw,
     })
-    } else {
-      Err(Error::new(
-        Status::InvalidArg,
-        "clone only available when the buffer is created from a JavaScript Buffer".to_owned(),
-      ))
-    }
   }
 }
 
@@ -76,6 +77,7 @@ impl From<Vec<u8>> for Buffer {
     let capacity = data.capacity();
     mem::forget(data);
     Buffer {
+      data_reference: Some(Arc::new(())),
       inner: unsafe { slice::from_raw_parts_mut(inner_ptr, len) },
       capacity,
       raw: None,
@@ -146,6 +148,7 @@ impl FromNapiValue for Buffer {
     )?;
 
     Ok(Self {
+      data_reference: None,
       inner: unsafe { slice::from_raw_parts_mut(buf as *mut _, len) },
       capacity: len,
       raw: Some((ref_, env)),
@@ -196,6 +199,32 @@ impl ToNapiValue for Buffer {
   }
 }
 
+impl ToNapiValue for &mut Buffer {
+  unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+    // From Node.js value, not from `Vec<u8>`
+    if let Some((ref_, _)) = val.raw {
+      let mut buf = ptr::null_mut();
+      check_status!(
+        unsafe { sys::napi_get_reference_value(env, ref_, &mut buf) },
+        "Failed to get Buffer value from reference"
+      )?;
+      check_status!(
+        unsafe { sys::napi_delete_reference(env, ref_) },
+        "Failed to delete Buffer reference"
+      )?;
+      Ok(buf)
+    } else {
+      let buf = Buffer {
+        data_reference: val.data_reference.clone(),
+        inner: unsafe { slice::from_raw_parts_mut(val.inner.as_mut_ptr(), val.capacity) },
+        capacity: val.capacity,
+        raw: None,
+      };
+      unsafe { ToNapiValue::to_napi_value(env, buf) }
+    }
+  }
+}
+
 impl ValidateNapiValue for Buffer {
   unsafe fn validate(env: sys::napi_env, napi_val: sys::napi_value) -> Result<sys::napi_value> {
     let mut is_buffer = false;

examples/napi/__test__/typegen.spec.ts.md

@@ -187,6 +187,7 @@ Generated by [AVA](https://avajs.dev).
     export function mutateTypedArray(input: Float32Array): void␊
     export function derefUint8Array(a: Uint8Array, b: Uint8ClampedArray): number␊
     export function bufferPassThrough(buf: Buffer): Promise<Buffer>␊
+    export function arrayBufferPassThrough(buf: Uint8Array): Promise<Uint8Array>␊
     export function asyncReduceBuffer(buf: Buffer): Promise<number>␊
     /**␊
      * \`constructor\` option for \`struct\` requires all fields to be public,␊
@@ -240,8 +241,10 @@ Generated by [AVA](https://avajs.dev).
     export class Blake2BKey { }␊
     export class Context {␊
       maybeNeed?: boolean␊
+      buffer: Uint8Array␊
       constructor()␊
       static withData(data: string): Context␊
+      static withBuffer(buf: Uint8Array): Context␊
       method(): string␊
     }␊
     export class AnimalWithDefaultConstructor {␊

examples/napi/__test__/values.spec.ts

@@ -90,6 +90,7 @@ import {
   derefUint8Array,
   chronoDateAdd1Minute,
   bufferPassThrough,
+  arrayBufferPassThrough,
   JsRepo,
   CssStyleSheet,
   asyncReduceBuffer,
@@ -208,6 +209,14 @@ test('class constructor return Result', (t) => {
   t.is(c.method(), 'not empty')
 })
 
+test('class default field is TypedArray', (t) => {
+  const c = new Context()
+  t.deepEqual(c.buffer, new Uint8Array([0, 1, 2, 3]))
+  const fixture = new Uint8Array([0, 1, 2, 3, 4, 5, 6])
+  const c2 = Context.withBuffer(fixture)
+  t.is(c2.buffer, fixture)
+})
+
 test('class Factory return Result', (t) => {
   const c = Context.withData('not empty')
   t.is(c.method(), 'not empty')
@@ -479,6 +488,12 @@ test('buffer passthrough', async (t) => {
   t.deepEqual(ret, fixture)
 })
 
+test('arraybuffer passthrough', async (t) => {
+  const fixture = new Uint8Array([1, 2, 3, 4, 5])
+  const ret = await arrayBufferPassThrough(fixture)
+  t.deepEqual(ret, fixture)
+})
+
 test('async reduce buffer', async (t) => {
   const input = [1, 2, 3, 4, 5, 6]
   const fixture = Buffer.from(input)

examples/napi/index.d.ts

@@ -177,6 +177,7 @@ export function createExternalTypedArray(): Uint32Array
 export function mutateTypedArray(input: Float32Array): void
 export function derefUint8Array(a: Uint8Array, b: Uint8ClampedArray): number
 export function bufferPassThrough(buf: Buffer): Promise<Buffer>
+export function arrayBufferPassThrough(buf: Uint8Array): Promise<Uint8Array>
 export function asyncReduceBuffer(buf: Buffer): Promise<number>
 /**
  * `constructor` option for `struct` requires all fields to be public,
@@ -230,8 +231,10 @@ export type Blake2bKey = Blake2BKey
 export class Blake2BKey { }
 export class Context {
   maybeNeed?: boolean
+  buffer: Uint8Array
   constructor()
   static withData(data: string): Context
+  static withBuffer(buf: Uint8Array): Context
   method(): string
 }
 export class AnimalWithDefaultConstructor {

examples/napi/src/class.rs

@@ -1,5 +1,5 @@
 use napi::{
-  bindgen_prelude::{Buffer, ClassInstance, This},
+  bindgen_prelude::{Buffer, ClassInstance, This, Uint8Array},
   Env, Result,
 };
 
@@ -159,6 +159,7 @@ impl Blake2bKey {
 pub struct Context {
   data: String,
   pub maybe_need: Option<bool>,
+  pub buffer: Uint8Array,
 }
 
 // Test for return `napi::Result` and `Result`
@@ -169,6 +170,7 @@ impl Context {
     Ok(Self {
       data: "not empty".into(),
       maybe_need: None,
+      buffer: Uint8Array::new(vec![0, 1, 2, 3]),
     })
   }
 
@@ -177,9 +179,19 @@ impl Context {
     Ok(Self {
       data,
       maybe_need: Some(true),
+      buffer: Uint8Array::new(vec![0, 1, 2, 3]),
     })
   }
 
+  #[napi(factory)]
+  pub fn with_buffer(buf: Uint8Array) -> Self {
+    Self {
+      data: "not empty".into(),
+      maybe_need: None,
+      buffer: buf,
+    }
+  }
+
   #[napi]
   pub fn method(&self) -> String {
     self.data.clone()

examples/napi/src/typed_array.rs

@@ -44,6 +44,11 @@ async fn buffer_pass_through(buf: Buffer) -> Result<Buffer> {
   Ok(buf)
 }
 
+#[napi]
+async fn array_buffer_pass_through(buf: Uint8Array) -> Result<Uint8Array> {
+  Ok(buf)
+}
+
 struct AsyncBuffer {
   buf: Buffer,
 }