Zero-copy passing rust `Vec`s to JavaScript `TypedArray`s and back.

helper.ts

/**
 * Takes a pointer to a wasm `TypedArrayData` and returns an appropriate JavaScript `TypedArray`
 * @param ptr A pointer to a `TypedArrayData` struct
 * @param memory The `WebAssembly.Memory` buffer containing the pointer
 * @param useCapacity If true the returned array will have a length determined by the total allocated capacity instead
 * of the number of elements.
 */
export function pointerToTypedArray(ptr: number, buffer: ArrayBuffer, useCapacity = false) {
    const [structPtr, structLen, structCap, structKind] = new Uint32Array(buffer, ptr, 4);
    const arraySize = useCapacity ? structCap : structLen;

    // tslint:disable:prettier
    switch (structKind) {
        case 0: return new Uint8Array(buffer, structPtr, arraySize);
        case 1: return new Int8Array(buffer, structPtr, arraySize);
        case 2: return new Uint16Array(buffer, structPtr, arraySize);
        case 3: return new Int16Array(buffer, structPtr, arraySize);
        case 4: return new Uint32Array(buffer, structPtr, arraySize);
        case 5: return new Int32Array(buffer, structPtr, arraySize);
        case 6: return new Float32Array(buffer, structPtr, arraySize);
        case 7: return new Float64Array(buffer, structPtr, arraySize);
    }
    // tslint:enable:prettier

    throw new Error(`Invalid TypedArray kind ${structKind}`);
}

js_data.rs

use std::mem;

/// The possible variants of TypedArrays in JavaScript. BigInt64Array and UBigInt64Array will
/// be added eventually as I64 and U64 respectively.
pub enum TypedArrayKind {
    U8 = 0,
    I8 = 1,
    U16 = 2,
    I16 = 3,
    U32 = 4,
    I32 = 5,
    F32 = 6,
    F64 = 7,
}

/// Lets us determine the proper TypedArrayKind for allowed types.
pub trait TypedArrayType {
    fn js_typed_array_kind() -> TypedArrayKind;
}

/// Simple macro to for `TypedArrayType` implementations. Takes a type and an identifier that is
/// a variant of `TypedArrayKind`.
macro_rules! typed_array_type {
    ($t: ty, $e: ident) => {
        impl TypedArrayType for $t {
            fn js_typed_array_kind() -> TypedArrayKind {
                TypedArrayKind::$e
            }
        }
    };
}

typed_array_type!(u8, U8);
typed_array_type!(i8, I8);
typed_array_type!(u16, U16);
typed_array_type!(i16, I16);
typed_array_type!(u32, U32);
typed_array_type!(i32, I32);
typed_array_type!(f32, F32);
typed_array_type!(f64, F64);

/// Contains information about a Vec whose data can be used as a JavaScript TypedArray.
#[repr(C)]
pub struct TypedArrayData {
    ptr: u32,
    len: u32,
    cap: u32,
    kind: u32,
}

impl TypedArrayData {
    ///
    /// Leaks the given data vector and constructs a `TypedArrayData` that contains information
    /// about it. The returned pointer can be passed to JavaScript and used to reconstruct the
    /// vector within the WebAssembly module's memory buffer, allowing for zero-copy passing
    /// of the data from Rust to JavaScript. Combined with `TypedArrayData::to_vec` we can do
    /// bi-directional zero-copy passing of data.
    ///
    /// # Safety
    ///
    /// You are responsible for calling `TypedArrayData::drop` on the returned pointer when the
    /// data is no longer being used.
    ///
    /// # Examples
    ///
    /// A simple function that returns some bytes
    ///
    /// ```rust
    /// #[no_mangle]
    /// pub fn bytes() -> *mut js_data::TypedArrayData {
    ///     let data: Vec<u8> = vec![1, 2, 3];
    ///     js_data::TypedArrayData::new(data)
    /// }
    /// ```
    ///
    /// Called from JavaScript this will return a Number that is the byte offset of our struct
    /// in the WebAssembly memory buffer. We can use this to retrieve the `TypedArrayData`, and
    /// get a TypedArray that points directly to our allocated memory.
    ///
    /// ```javascript
    /// loadWasmInstanceSomehow().then(instance => {
    ///     const pointerToBytes = instance.exports.bytes();
    ///
    ///     // `TypedArrayData` is composed of four consecutive u32's (see `repr(C)`), so we can read
    ///     // it as a `Uint32Array`.
    ///     const struct = new Uint32Array(instance.exports.memory.buffer, pointerToBytes, 4);
    ///
    ///     // Fields from the `TypedArrayData`
    ///     const structPtr = struct[0];
    ///     const structLen = struct[1];
    ///     const structCap = struct[2];
    ///     const structKind = struct[3];
    ///
    ///     // We're expecting u8 data, so make sure that's what we're getting
    ///     if (structKind !== 0 /* TypedArrayKind::U8 */) {
    ///         throw new Error('Expected u8 data, got something else');
    ///     }
    ///
    ///     const bytesLen = new Uint8Array(instance.exports.memory.buffer, structPtr, structLen);
    ///     // or (to get the entire *allocated* vector instead of just the written items)
    ///     const bytesCap = new Uint8Array(instance.exports.memory.buffer, structPtr, structCap);
    /// });
    /// ```
    ///
    /// Or using the `pointerToTypedArray` helper which will auto-detect the proper TypedArray variant
    /// and deal with the `TypedArrayData` struct for us.
    ///
    /// ```javascript
    /// loadWasmInstanceSomehow().then(instance => {
    ///     const pointerToBytes = instance.exports.bytes();
    ///
    ///     const bytesLen = pointerToTypedArray(pointerToBytes, instance.exports.buffer);
    ///     // or (to get the entire *allocated* vector instead of just the written items)
    ///     const bytesCap = pointerToTypedArray(pointerToBytes, instance.exports.buffer, true);
    /// });
    /// ```
    ///
    pub fn new<T: TypedArrayType>(mut data: Vec<T>) -> *mut TypedArrayData {
        let ptr = data.as_mut_ptr() as u32;
        let len = data.len() as u32;
        let cap = data.capacity() as u32;
        let kind = T::js_typed_array_kind() as u32;

        mem::forget(data);

        let boxed = Box::new(TypedArrayData {
            ptr,
            len,
            cap,
            kind,
        });

        Box::into_raw(boxed)
    }

    ///
    /// Reconstructs a vector for any valid underlying data type, and allows it to go out of
    /// scope properly.
    ///
    /// # Safety
    ///
    /// If this function is passed something other than a valid pointer to a `TypedArrayData`
    /// it will almost certainly corrupt memory.
    ///
    /// After a pointer is passed to this the data MUST NOT be used in JavaScript either, as
    /// it is no longer owned by the TypedArrayData.
    ///
    /// # Panics
    ///
    /// If the `TypedArrayData` contains a kind that is not in `TypedArrayKind` we will panic.
    ///
    /// # Examples
    ///
    /// A simple function that will free a `TypedArrayData` pointer passed in from JavaScript
    ///
    /// ```rust
    ///    #[no_mangle]
    ///    pub fn free_typed_array(ptr: *mut js_data::TypedArrayData) {
    ///        js_data::TypedArrayData::drop(ptr);
    ///    }
    /// ```
    ///
    /// And calling it from JavaScript (see example on `TypedArrayData::new` for a definition
    /// of `bytes()`)
    ///
    /// ```javascript
    /// loadWasmInstanceSomehow().then(instance => {
    ///     const pointerToBytes = instance.exports.bytes();
    ///
    ///     instance.exports.free_typed_array(pointerToBytes);
    /// });
    /// ```
    ///
    pub fn drop(ptr: *mut TypedArrayData) {
        let boxed: Box<TypedArrayData> = unsafe { Box::from_raw(ptr) };

        macro_rules! boxed_kind_equals {
            ($e: ident) => {
                boxed.kind == TypedArrayKind::$e as u32
            };
        }

        macro_rules! boxed_to_vec {
            ($t: ty) => {
                unsafe {
                    Vec::from_raw_parts(
                        boxed.ptr as *mut $t,
                        boxed.len as usize,
                        boxed.cap as usize,
                    )
                };
            };
        }

        if boxed_kind_equals!(U8) {
            boxed_to_vec!(u8);
        } else if boxed_kind_equals!(I8) {
            boxed_to_vec!(i8);
        } else if boxed_kind_equals!(U16) {
            boxed_to_vec!(u16);
        } else if boxed_kind_equals!(I16) {
            boxed_to_vec!(i16);
        } else if boxed_kind_equals!(U32) {
            boxed_to_vec!(u32);
        } else if boxed_kind_equals!(I32) {
            boxed_to_vec!(i32);
        } else if boxed_kind_equals!(F32) {
            boxed_to_vec!(f32);
        } else if boxed_kind_equals!(F64) {
            boxed_to_vec!(f64);
        } else {
            panic!("Invalid TypedArray data type {}", boxed.kind);
        };
    }

    ///
    /// Allows us to retrieve a `Vec` that was previously turned into a `TypedArrayData`. This
    /// means we can zero-copy data *back* into rust from JavaScript.
    ///
    /// # Safety
    ///
    /// If this function is passed something other than a valid pointer to a `TypedArrayData`
    /// it will almost certainly corrupt memory.
    ///
    /// After a pointer is passed to this the data can still be changed from JavaScript, but
    /// only as long as the vector remains in scope. This is highly unsafe unless managed
    /// properly, as once the vector goes out of scope in rust the underlying data should be
    /// freed.
    ///
    /// # Panics
    ///
    /// If the passed `TypedArrayData`'s kind doesn't match `T::js_typed_array_kind()` we will
    /// panic.
    ///
    pub fn to_vec<T: TypedArrayType>(ptr: *mut TypedArrayData) -> Vec<T> {
        let boxed: Box<TypedArrayData> = unsafe { Box::from_raw(ptr) };

        if boxed.kind != T::js_typed_array_kind() as u32 {
            panic!("Invalid TypedArray data type {}", boxed.kind);
        };

        unsafe { Vec::from_raw_parts(boxed.ptr as *mut T, boxed.len as usize, boxed.cap as usize) }
    }
}