feat(napi): support TypedArray input and output

crates/napi/src/bindgen_runtime/js_values.rs

@@ -3,6 +3,7 @@ use std::ptr;
 use crate::{check_status, sys, Error, JsUnknown, NapiRaw, NapiValue, Result, Status, ValueType};
 
 mod array;
+mod arraybuffer;
 #[cfg(feature = "napi6")]
 mod bigint;
 mod boolean;
@@ -23,6 +24,7 @@ mod symbol;
 mod task;
 
 pub use array::*;
+pub use arraybuffer::*;
 #[cfg(feature = "napi6")]
 pub use bigint::*;
 pub use buffer::*;

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

@@ -0,0 +1,228 @@
+use std::ffi::c_void;
+use std::mem;
+use std::ops::{Deref, DerefMut};
+use std::ptr;
+
+pub use crate::js_values::TypedArrayType;
+use crate::{check_status, sys, Error, Result, Status};
+
+use super::{FromNapiValue, ToNapiValue, TypeName};
+
+macro_rules! impl_typed_array {
+  ($name:ident, $rust_type:ident, $typed_array_type:expr) => {
+    pub struct $name {
+      data: *mut $rust_type,
+      length: usize,
+      data_managed_type: DataManagedType,
+      byte_offset: usize,
+      finalizer_notify: Box<dyn FnOnce(*mut $rust_type, usize)>,
+    }
+
+    impl $name {
+      fn noop_finalize(_data: *mut $rust_type, _length: usize) {}
+
+      pub fn new(mut data: Vec<$rust_type>) -> Self {
+        let ret = $name {
+          data: data.as_mut_ptr(),
+          length: data.len(),
+          data_managed_type: DataManagedType::Owned,
+          byte_offset: 0,
+          finalizer_notify: Box::new(Self::noop_finalize),
+        };
+        mem::forget(data);
+        ret
+      }
+
+      pub fn with_data_copied<D>(data: D) -> Self
+      where
+        D: AsRef<[$rust_type]>,
+      {
+        let mut data_copied = data.as_ref().to_vec();
+        let ret = $name {
+          data: data_copied.as_mut_ptr(),
+          length: data.as_ref().len(),
+          data_managed_type: DataManagedType::Owned,
+          finalizer_notify: Box::new(Self::noop_finalize),
+          byte_offset: 0,
+        };
+        mem::forget(data_copied);
+        ret
+      }
+
+      /// # Safety
+      ///
+      /// The caller will be notified when the data is deallocated by vm
+      pub unsafe fn with_external_data<F>(data: *mut $rust_type, length: usize, notify: F) -> Self
+      where
+        F: 'static + FnOnce(*mut $rust_type, usize),
+      {
+        $name {
+          data,
+          length,
+          data_managed_type: DataManagedType::External,
+          finalizer_notify: Box::new(notify),
+          byte_offset: 0,
+        }
+      }
+    }
+
+    impl Deref for $name {
+      type Target = [$rust_type];
+
+      fn deref(&self) -> &Self::Target {
+        unsafe { std::slice::from_raw_parts(self.data, self.length) }
+      }
+    }
+
+    impl DerefMut for $name {
+      fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { std::slice::from_raw_parts_mut(self.data, self.length) }
+      }
+    }
+
+    impl AsRef<[$rust_type]> for $name {
+      fn as_ref(&self) -> &[$rust_type] {
+        unsafe { std::slice::from_raw_parts(self.data, self.length) }
+          .split_at(self.byte_offset)
+          .1
+      }
+    }
+
+    impl AsMut<[$rust_type]> for $name {
+      fn as_mut(&mut self) -> &mut [$rust_type] {
+        unsafe { std::slice::from_raw_parts_mut(self.data, self.length) }
+          .split_at_mut(self.byte_offset)
+          .1
+      }
+    }
+
+    impl TypeName for $name {
+      fn type_name() -> &'static str {
+        "TypedArray"
+      }
+
+      fn value_type() -> crate::ValueType {
+        crate::ValueType::Object
+      }
+    }
+
+    impl FromNapiValue for $name {
+      unsafe fn from_napi_value(
+        env: napi_sys::napi_env,
+        napi_val: napi_sys::napi_value,
+      ) -> Result<Self> {
+        let mut typed_array_type = 0;
+        let mut length = 0;
+        let mut data = ptr::null_mut();
+        let mut array_buffer = ptr::null_mut();
+        let mut byte_offset = 0;
+        check_status!(
+          sys::napi_get_typedarray_info(
+            env,
+            napi_val,
+            &mut typed_array_type,
+            &mut length,
+            &mut data,
+            &mut array_buffer,
+            &mut byte_offset
+          ),
+          "Get TypedArray info failed"
+        )?;
+        if typed_array_type != $typed_array_type as i32 {
+          return Err(Error::new(
+            Status::InvalidArg,
+            format!("Expected $name, got {}", typed_array_type),
+          ));
+        }
+        Ok($name {
+          data: data as *mut $rust_type,
+          length,
+          byte_offset,
+          data_managed_type: DataManagedType::Vm,
+          finalizer_notify: Box::new(Self::noop_finalize),
+        })
+      }
+    }
+
+    impl ToNapiValue for $name {
+      unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::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((
+          val.data_managed_type,
+          val.length,
+          val.finalizer_notify,
+        )));
+        check_status!(
+          sys::napi_create_external_arraybuffer(
+            env,
+            val.data as *mut c_void,
+            length,
+            Some(finalizer::<$rust_type>),
+            hint_ptr as *mut c_void,
+            &mut arraybuffer_value
+          ),
+          "Create external arraybuffer failed"
+        )?;
+        let mut napi_val = ptr::null_mut();
+        check_status!(
+          sys::napi_create_typedarray(
+            env,
+            $typed_array_type as i32,
+            val.length,
+            arraybuffer_value,
+            0,
+            &mut napi_val,
+          ),
+          "Create TypedArray failed"
+        )?;
+        Ok(napi_val)
+      }
+    }
+  };
+}
+
+unsafe extern "C" fn finalizer<T>(
+  _env: sys::napi_env,
+  finalize_data: *mut c_void,
+  finalize_hint: *mut c_void,
+) {
+  let (data_managed_type, length, finalizer_notify) =
+    *Box::from_raw(finalize_hint as *mut (DataManagedType, usize, Box<dyn FnOnce(*mut T, usize)>));
+  match data_managed_type {
+    DataManagedType::Vm => {
+      // do nothing
+    }
+    DataManagedType::Owned => {
+      let length = length;
+      Vec::from_raw_parts(finalize_data as *mut T, length, length);
+    }
+    DataManagedType::External => {
+      (finalizer_notify)(finalize_data as *mut T, length);
+    }
+  }
+}
+
+enum DataManagedType {
+  /// Vm managed data, passed in from JavaScript
+  Vm,
+  /// Rust owned data, which need to be deallocated in the finalizer
+  Owned,
+  /// External data, which need to be notice to the owner in finalizer
+  External,
+}
+
+impl_typed_array!(Int8Array, i8, TypedArrayType::Int8);
+impl_typed_array!(Uint8Array, u8, TypedArrayType::Uint8);
+impl_typed_array!(Uint8ClampedArray, u8, TypedArrayType::Uint8Clamped);
+impl_typed_array!(Int16Array, i16, TypedArrayType::Int16);
+impl_typed_array!(Uint16Array, u16, TypedArrayType::Uint16);
+impl_typed_array!(Int32Array, i32, TypedArrayType::Int32);
+impl_typed_array!(Uint32Array, u32, TypedArrayType::Uint32);
+impl_typed_array!(Float32Array, f32, TypedArrayType::Float32);
+impl_typed_array!(Float64Array, f64, TypedArrayType::Float64);
+#[cfg(feature = "napi6")]
+impl_typed_array!(BigInt64Array, i64, TypedArrayType::BigInt64);
+#[cfg(feature = "napi6")]
+impl_typed_array!(BigUint64Array, u64, TypedArrayType::BigUint64);

crates/napi/src/js_values/arraybuffer.rs

@@ -39,8 +39,8 @@ impl TypeName for JsTypedArray {
 pub struct JsTypedArrayValue {
   pub arraybuffer: JsArrayBuffer,
   data: *mut c_void,
-  pub byte_offset: u64,
-  pub length: u64,
+  pub byte_offset: usize,
+  pub length: usize,
   pub typedarray_type: TypedArrayType,
 }
 
@@ -234,19 +234,19 @@ impl JsTypedArray {
   /// ***Warning***: Use caution while using this API since the underlying data buffer is managed by the VM.
   pub fn into_value(self) -> Result<JsTypedArrayValue> {
     let mut typedarray_type = 0;
-    let mut len = 0u64;
+    let mut len = 0;
     let mut data = ptr::null_mut();
     let mut arraybuffer_value = ptr::null_mut();
-    let mut byte_offset = 0u64;
+    let mut byte_offset = 0;
     check_status!(unsafe {
       sys::napi_get_typedarray_info(
         self.0.env,
         self.0.value,
         &mut typedarray_type,
-        &mut len as *mut u64 as *mut _,
+        &mut len,
         &mut data,
         &mut arraybuffer_value,
-        &mut byte_offset as *mut u64 as *mut usize,
+        &mut byte_offset,
       )
     })?;
 
@@ -260,34 +260,51 @@ impl JsTypedArray {
   }
 }
 
+impl JsTypedArrayValue {
+  #[inline]
+  fn is_valid_as_ref(&self, dest_type: TypedArrayType) {
+    if self.typedarray_type == TypedArrayType::Uint8Clamped && dest_type != TypedArrayType::Uint8 {
+      return;
+    }
+    if self.typedarray_type != dest_type {
+      panic!(
+        "invalid typedarray type: expected {:?}, got {:?}",
+        dest_type, self.typedarray_type
+      );
+    }
+  }
+}
+
 macro_rules! impl_as_ref {
-  ($ref_type:ident) => {
+  ($ref_type:ident, $expect_type:expr) => {
     impl AsRef<[$ref_type]> for JsTypedArrayValue {
       fn as_ref(&self) -> &[$ref_type] {
-        unsafe { slice::from_raw_parts(self.data as *const $ref_type, self.length as usize) }
+        self.is_valid_as_ref($expect_type);
+        unsafe { slice::from_raw_parts(self.data as *const $ref_type, self.length) }
       }
     }
 
     impl AsMut<[$ref_type]> for JsTypedArrayValue {
       fn as_mut(&mut self) -> &mut [$ref_type] {
-        unsafe { slice::from_raw_parts_mut(self.data as *mut $ref_type, self.length as usize) }
+        self.is_valid_as_ref($expect_type);
+        unsafe { slice::from_raw_parts_mut(self.data as *mut $ref_type, self.length) }
       }
     }
   };
 }
 
-impl_as_ref!(u8);
-impl_as_ref!(i8);
-impl_as_ref!(u16);
-impl_as_ref!(i16);
-impl_as_ref!(u32);
-impl_as_ref!(i32);
-impl_as_ref!(f32);
-impl_as_ref!(f64);
+impl_as_ref!(u8, TypedArrayType::Uint8);
+impl_as_ref!(i8, TypedArrayType::Int8);
+impl_as_ref!(u16, TypedArrayType::Uint16);
+impl_as_ref!(i16, TypedArrayType::Int16);
+impl_as_ref!(u32, TypedArrayType::Uint32);
+impl_as_ref!(i32, TypedArrayType::Int32);
+impl_as_ref!(f32, TypedArrayType::Float32);
+impl_as_ref!(f64, TypedArrayType::Float64);
 #[cfg(feature = "napi6")]
-impl_as_ref!(i64);
+impl_as_ref!(i64, TypedArrayType::BigInt64);
 #[cfg(feature = "napi6")]
-impl_as_ref!(u64);
+impl_as_ref!(u64, TypedArrayType::BigUint64);
 
 impl JsDataView {
   pub fn into_value(self) -> Result<JsDataViewValue> {