add transferables support

Author: anulman

src/index.ts

@@ -1,6 +1,7 @@
 import guest from "./guest";
 import host from "./host";
+import { withTransferable } from "./rpc";
 
-export { host, guest };
+export { host, guest, withTransferable };
 
 export * from "./types";

src/rpc.ts

@@ -10,6 +10,9 @@ import {
 } from "./helpers";
 import { actions, events, IRPCRequestPayload, IRPCResolvePayload, ISchema } from "./types";
 
+/** Private symbol to which we will assign transferable objects */
+const SYM_TRANSFERABLES = Symbol();
+
 /**
  * for each function in the schema
  * 1. subscribe to an event that the remote can call
@@ -47,24 +50,19 @@ export function registerLocalMethods(
       };
 
       // run function and return the results to the remote
+      const cb = get(schema, methodName);
+      const transferables = cb[SYM_TRANSFERABLES];
+
       try {
-        const result = await get(schema, methodName)(...args);
-
-        if (!result) {
-          // if the result is falsy (null, undefined, "", etc), set it directly
-          payload.result = result;
-        } else {
-          // otherwise parse a stringified version of it
-          payload.result = JSON.parse(JSON.stringify(result));
-        }
+        payload.result = await cb(...args);
       } catch (error) {
         payload.action = actions.RPC_REJECT;
         payload.error = JSON.parse(JSON.stringify(error, Object.getOwnPropertyNames(error)));
       }
 
-      if (guest) guest.postMessage(payload);
-      else if (isWorker()) (self as any).postMessage(payload);
-      else event.source.postMessage(payload, event.origin);
+      if (guest) guest.postMessage(payload, transferables);
+      else if (isWorker()) (self as any).postMessage(payload, transferables);
+      else event.source.postMessage(payload, event.origin, transferables);
     }
 
     // subscribe to the call event
@@ -99,7 +97,7 @@ export function createRPC(
   listeners: Array<() => void> = [],
   guest?: Worker
 ) {
-  return (...args: any) => {
+  return (...args: any[]) => {
     return new Promise((resolve, reject) => {
       const callID = generateId();
 
@@ -120,7 +118,7 @@ export function createRPC(
       // send the RPC request with arguments
       const payload = {
         action: actions.RPC_REQUEST,
-        args: JSON.parse(JSON.stringify(args)),
+        args,
         callID,
         callName: _callName,
         connectionID: _connectionID,
@@ -134,9 +132,12 @@ export function createRPC(
         listeners.push(() => removeEventListener(self, events.MESSAGE, handleResponse));
       }
 
-      if (guest) guest.postMessage(payload);
-      else if (isWorker() || isNodeEnv()) (self as any).postMessage(payload);
-      else (event.source || event.target).postMessage(payload, event.origin);
+      // @ts-expect-error: we know this is an array of transferables (if it exists)
+      const transferables = args[SYM_TRANSFERABLES] ?? (args.length === 1 ? args[0][SYM_TRANSFERABLES] : []);
+
+      if (guest) guest.postMessage(payload, transferables);
+      else if (isWorker() || isNodeEnv()) (self as any).postMessage(payload, transferables);
+      else (event.source || event.target).postMessage(payload, event.origin, transferables);
     });
   };
 }
@@ -171,3 +172,43 @@ export function registerRemoteMethods(
     unregisterRemote: () => listeners.forEach((unregister) => unregister()),
   };
 }
+
+/**
+ * This function is used by API schema declarations and remote function calls alike to
+ * indicate which variables should be declared as transferable over `postMessage` calls.
+ *
+ * @param cb a function that takes a transfer function as an argument and returns an object
+ *           (in the loose, `typeof foo === "object"` sense)
+ * @return result the callback's return value, with an extra array of transferable objects
+ *                assigned to rimless' private symbol `SYM_TRANSFERABLES`
+ *
+ * The `transfer(...)` function can be called with one or more objects to transfer. When
+ * called with one object, it returns that object. When called with zero or multiple objects,
+ * it returns an array of the objects. Calling `transfer` will only modify the callback result,
+ * not the original object itself (or objects themselves).
+ *
+ * @example
+ * host.connect({
+ *   foo: withTransferable((transfer) => {
+ *     const foo = new ArrayBuffer(8);
+ *     const bar = new ArrayBuffer(8);
+ *     transfer(foo, bar);
+ *
+ *     return { foo, bar };
+ *   }),
+ * });
+ *
+ * @example
+ * host.remote.foo(withTransferable((transfer) => transfer(new ArrayBuffer(8))));
+ */
+export const withTransferable = <T, V extends object>(cb: (transfer: (transferable: T) => void) => V) => {
+  const transferables: T[] = [];
+  const transfer = (...toTransfer: [T, ...T[]]) => {
+    transferables.push(...toTransfer);
+    return toTransfer.length === 1 ? toTransfer[0] : toTransfer;
+  };
+
+  const result = cb(transfer);
+
+  return Object.assign(result, { [SYM_TRANSFERABLES]: transferables });
+};