refactor ref assignment

rollup.config.js

@@ -395,7 +395,6 @@ function managlePropsPlugin() {
     $immediateSetters$: '',
     $importScripts$: '',
     $instanceId$: '',
-    $instanceIdByInstance$: '',
     $instances$: '',
     $interfaces$: '',
     $interfaceType$: '',
@@ -406,13 +405,11 @@ function managlePropsPlugin() {
     $memberPath$: '',
     $msgId$: '',
     $newInstanceId$: '',
-    $nextId$: '',
     $nodeName$: '',
     $parentWinId$: '',
     $postMessage$: '',
     $rtnValue$: '',
     $scopePath$: '',
-    $stateProp$: '',
     $tasks$: '',
     $url$: '',
     $window$: '',

src/lib/sandbox/index.ts

@@ -1,3 +1,4 @@
 import { initSandbox } from './init-sandbox';
+import { TOP_WIN_ID } from '../utils';
 
-initSandbox(window);
+initSandbox(window, TOP_WIN_ID);

src/lib/sandbox/init-sandbox.ts

@@ -1,19 +1,17 @@
 import { createWebWorker } from './messenger';
-import { debug, TOP_WIN_ID } from '../utils';
+import { debug } from '../utils';
 import { mainAccessHandler } from './main-access-handler';
 import { MainAccessRequest, MainWindow, MainWindowContext, PlatformInstanceId } from '../types';
 import { readNextScript } from './read-main-scripts';
 import { setInstanceId } from './main-instances';
+import { winCtxs, windows } from './main-constants';
 
-export const initSandbox = async (sandboxWindow: Window) => {
-  let winIds = TOP_WIN_ID;
-  const winCtxs = new Map<number, MainWindowContext>();
-  const windows = new WeakSet<MainWindow>();
+export const initSandbox = async (sandboxWindow: Window, winIds: number) => {
   const mainWindow: MainWindow = sandboxWindow.parent as any;
   const swContainer = sandboxWindow.navigator.serviceWorker;
   const swRegistration = await swContainer.getRegistration();
 
-  const onMessageFromServiceWorker = (ev: MessageEvent<MainAccessRequest>) => {
+  const onMessageFromServiceWorkerToSandbox = (ev: MessageEvent<MainAccessRequest>) => {
     const accessReq = ev.data;
     const accessWinId = accessReq.$winId$;
     const winCtx = winCtxs.get(accessWinId);
@@ -37,9 +35,6 @@ export const initSandbox = async (sandboxWindow: Window) => {
         $parentWinId$: parentWin.partyWinId!,
         $config$: mainWindow.partytown,
         $cleanupInc$: 0,
-        $instanceIdByInstance$: new WeakMap(),
-        $instances$: [],
-        $nextId$: PlatformInstanceId.body + 1,
         $scopePath$: swRegistration!.scope,
         $url$: win.document.baseURI,
         $window$: win,
@@ -56,9 +51,9 @@ export const initSandbox = async (sandboxWindow: Window) => {
       setInstanceId(winCtx, win.localStorage, PlatformInstanceId.localStorage);
       setInstanceId(winCtx, win.sessionStorage, PlatformInstanceId.sessionStorage);
 
-      swContainer.addEventListener('message', onMessageFromServiceWorker);
+      swContainer.addEventListener('message', onMessageFromServiceWorkerToSandbox);
 
-      createWebWorker(winCtxs, winCtx);
+      createWebWorker(winCtx);
 
       win.addEventListener('load', () => readNextScript(winCtx));
     }

src/lib/sandbox/main-access-handler.ts

@@ -80,7 +80,7 @@ export const mainAccessHandler = async (
           rtnValue = await rtnValue;
           accessRsp.$isPromise$ = true;
         }
-        accessRsp.$rtnValue$ = serializeForWorker(winCtx, rtnValue, new Set());
+        accessRsp.$rtnValue$ = serializeForWorker(winCtx, rtnValue);
       } else {
         accessRsp.$errors$.push(`Instance ${instanceId} not found`);
       }

src/lib/sandbox/main-constants.ts

@@ -0,0 +1,8 @@
+import type { MainAccessResponse, MainWindow, MainWindowContext, RefMap } from '../types';
+
+export const forwardMsgResolves = new Map<number, (accessRsp: MainAccessResponse) => void>();
+export const mainInstanceIdByInstance = new WeakMap<any, number>();
+export const mainInstances: [number, any][] = [];
+export const mainInstanceRefs = new WeakMap<any, RefMap>();
+export const winCtxs = new Map<number, MainWindowContext>();
+export const windows = new WeakSet<MainWindow>();

src/lib/sandbox/main-event-forwarding.ts

@@ -11,7 +11,7 @@ export const mainEventForwarding = (winCtx: MainWindowContext, win: MainWindow)
       winCtx.$worker$!.postMessage([
         WorkerMessageType.ForwardEvent,
         forwardConfig,
-        serializeForWorker(winCtx, Array.from(forwardArgs), new Set()),
+        serializeForWorker(winCtx, Array.from(forwardArgs)),
       ]);
 
     for (; i < len(forwardedEvents as any[]); i += 2) {

src/lib/sandbox/main-instances.ts

@@ -1,10 +1,8 @@
-import { len } from '../utils';
+import { len, randomId } from '../utils';
+import { mainInstanceIdByInstance, mainInstances } from './main-constants';
 import { MainWindowContext, NodeName, PlatformInstanceId } from '../types';
 
-export const getInstanceId = (
-  winCtx: MainWindowContext,
-  instance: InstanceType | null | undefined
-) => {
+export const getInstanceId = (instance: InstanceType | null | undefined) => {
   if (instance) {
     const nodeName = (instance as any as Node).nodeName;
     if (nodeName === NodeName.Document) {
@@ -19,8 +17,9 @@ export const getInstanceId = (
     if (nodeName === NodeName.Body) {
       return PlatformInstanceId.body;
     }
-    return winCtx.$instanceIdByInstance$.get(instance);
+    return mainInstanceIdByInstance.get(instance);
   }
+  return -1;
 };
 
 export const getAndSetInstanceId = (
@@ -29,10 +28,9 @@ export const getAndSetInstanceId = (
   instanceId?: number
 ) => {
   if (instance) {
-    instanceId = getInstanceId(winCtx, instance);
+    instanceId = getInstanceId(instance);
     if (typeof instanceId !== 'number') {
-      instanceId = winCtx.$nextId$++;
-      setInstanceId(winCtx, instance, instanceId);
+      setInstanceId(winCtx, instance, (instanceId = randomId()));
     }
     return instanceId;
   }
@@ -57,7 +55,7 @@ export const getInstance = <T = InstanceType | null>(
   if (instanceId === PlatformInstanceId.body) {
     return doc.body as any;
   }
-  instanceItem = winCtx.$instances$.find((i) => i[0] === instanceId);
+  instanceItem = mainInstances.find((i) => i[0] === instanceId);
   return instanceItem ? instanceItem[1] : null;
 };
 
@@ -67,23 +65,22 @@ export const setInstanceId = (
   instanceId: number
 ) => {
   if (instance) {
-    const instances = winCtx.$instances$;
-    instances.push([instanceId, instance]);
-    winCtx.$instanceIdByInstance$.set(instance, instanceId);
+    mainInstances.push([instanceId, instance]);
+    mainInstanceIdByInstance.set(instance, instanceId);
 
     winCtx.$cleanupInc$++;
     if (winCtx.$cleanupInc$ > 99) {
       winCtx.$cleanupInc$ = 0;
       while (true) {
-        let disconnectedNodes = instances.filter(
+        let disconnectedNodes = mainInstances.filter(
           (i) => (i[1] as InstanceNode).nodeType && !(i[1] as InstanceNode).isConnected
         );
         let i: number;
         let l: number;
         if (len(disconnectedNodes) > 99) {
-          for (i = 0, l = len(instances); i < l; i++) {
-            if (!(instances[i][1] as InstanceNode).isConnected) {
-              instances.slice(i, 1);
+          for (i = 0, l = len(mainInstances); i < l; i++) {
+            if (!(mainInstances[i][1] as InstanceNode).isConnected) {
+              mainInstances.slice(i, 1);
               break;
             }
           }

src/lib/sandbox/main-serialization.ts

@@ -4,22 +4,27 @@ import {
   InterfaceType,
   MainWindowContext,
   PlatformInstanceId,
+  RefHandler,
+  RefMap,
   SerializedInstance,
   SerializedTransfer,
   SerializedType,
   WorkerMessageType,
 } from '../types';
+import { mainInstanceRefs } from './main-constants';
 
 export const serializeForWorker = (
   winCtx: MainWindowContext,
   value: any,
-  added: Set<any>,
+  added?: Set<any>,
   type?: string,
   obj?: { [key: string]: SerializedTransfer | undefined },
   key?: string
 ): SerializedTransfer | undefined => {
   if (value !== undefined) {
+    added = added || new Set();
     type = typeof value;
+
     if (type === 'string' || type === 'number' || type === 'boolean' || value == null) {
       return [SerializedType.Primitive, value];
     }
@@ -97,7 +102,7 @@ export const deserializeFromWorker = (
     }
 
     if (serializedType === SerializedType.Ref) {
-      return setRef(winCtx, instanceId, serializedValue);
+      return deserializeRefFromWorker(winCtx, instanceId, serializedValue);
     }
 
     if (serializedType === SerializedType.Array) {
@@ -122,43 +127,37 @@ export const deserializeFromWorker = (
 
 const isNodeList = (cstrName: string) => cstrName === 'HTMLCollection' || cstrName === 'NodeList';
 
-const instanceRefs = new WeakMap<any, RefHandlerMap>();
-
-const setRef = (winCtx: MainWindowContext, instanceId: number, refId: number) => {
+const deserializeRefFromWorker = (winCtx: MainWindowContext, instanceId: number, refId: number) => {
   let instance = getInstance(winCtx, instanceId);
-  let refHandlerMap: RefHandlerMap | undefined;
-  let refHandler: RefHandler | undefined;
+  let mainRefHandlerMap: RefMap | undefined;
+  let mainRefHandler: RefHandler | undefined;
 
   if (instance) {
-    refHandlerMap = instanceRefs.get(instance);
-    if (!refHandlerMap) {
-      refHandlerMap = new Map<number, RefHandler>();
+    mainRefHandlerMap = mainInstanceRefs.get(instance);
+    if (!mainRefHandlerMap) {
+      mainInstanceRefs.set(instance, (mainRefHandlerMap = {}));
     }
 
-    refHandler = refHandlerMap.get(refId);
-    if (!refHandler) {
-      refHandler = createRefHandler(winCtx, refId);
-      refHandlerMap.set(refId, refHandler);
+    mainRefHandler = mainRefHandlerMap[refId];
+    if (!mainRefHandler) {
+      mainRefHandler = createMainRefHandler(winCtx, refId);
+      mainRefHandlerMap[refId] = mainRefHandler;
     }
 
-    return refHandler;
+    return mainRefHandler;
   }
 
   return noop;
 };
 
-const createRefHandler = (winCtx: MainWindowContext, refId: number): RefHandler =>
+const createMainRefHandler = (winCtx: MainWindowContext, refId: number): RefHandler =>
   function (this: any, ...args: any[]) {
-    const serializedTarget = serializeForWorker(winCtx, this, new Set());
-    const serializedArgs = serializeForWorker(winCtx, args, new Set());
+    const serializedTarget = serializeForWorker(winCtx, this);
+    const serializedArgs = serializeForWorker(winCtx, args);
     winCtx.$worker$!.postMessage([
       WorkerMessageType.RefHandlerCallback,
       refId,
       serializedTarget,
       serializedArgs,
     ]);
   };
-
-type RefHandler = (...args: any[]) => void;
-
-type RefHandlerMap = Map<number, RefHandler>;

src/lib/sandbox/messenger.ts

@@ -1,4 +1,5 @@
 import { debug } from '../utils';
+import { forwardMsgResolves, winCtxs } from './main-constants';
 import { getAndSetInstanceId } from './main-instances';
 import {
   InitWebWorkerData,
@@ -13,11 +14,7 @@ import { initializedWorkerScript, readNextScript } from './read-main-scripts';
 import { readMainInterfaces } from './read-interfaces';
 import WebWorkerBlob from '@web-worker-blob';
 
-const onMessageFromWebWorker = (
-  winCtxs: Map<number, MainWindowContext>,
-  winCtx: MainWindowContext,
-  msg: MessageFromWorkerToSandbox
-) => {
+const onMessageFromWebWorker = (winCtx: MainWindowContext, msg: MessageFromWorkerToSandbox) => {
   const msgType = msg[0];
   const doc = winCtx.$window$.document;
 
@@ -65,8 +62,6 @@ const onMessageFromWebWorker = (
   }
 };
 
-const forwardMsgResolves = new Map<number, (accessRsp: MainAccessResponse) => void>();
-
 export const forwardToWinAccessHandler = (
   worker: PartytownWebWorker,
   accessReq: MainAccessRequest
@@ -76,10 +71,7 @@ export const forwardToWinAccessHandler = (
     worker.postMessage([WorkerMessageType.ForwardMainDataRequest, accessReq]);
   });
 
-export const createWebWorker = (
-  winCtxs: Map<number, MainWindowContext>,
-  winCtx: MainWindowContext
-) => {
+export const createWebWorker = (winCtx: MainWindowContext) => {
   winCtx.$worker$ = new Worker(
     debug
       ? './partytown-ww.debug.js'
@@ -91,5 +83,5 @@ export const createWebWorker = (
     { name: `Partytown (${winCtx.$winId$}) 🎉` }
   );
 
-  winCtx.$worker$.onmessage = (ev) => onMessageFromWebWorker(winCtxs, winCtx, ev.data);
+  winCtx.$worker$.onmessage = (ev) => onMessageFromWebWorker(winCtx, ev.data);
 };

src/lib/service-worker/sw-message.ts

@@ -1,6 +1,6 @@
-import type { MainAccessRequest, MainAccessResponse } from '../types';
-import { debug } from '../utils';
 import { CacheControl, ContentType, response } from './response';
+import { debug } from '../utils';
+import type { MainAccessRequest, MainAccessResponse } from '../types';
 
 const resolves = new Map<number, MessageResolve>();
 
@@ -55,7 +55,7 @@ export const httpRequestFromWebWorker = (self: ServiceWorkerGlobalScope, req: Re
     const accessReq: MainAccessRequest = await req.clone().json();
     const responseData = await sendMessageToSandboxFromServiceWorker(self, accessReq);
 
-    resolve(response(JSON.stringify(responseData), ContentType.JSON, CacheControl.Immutable));
+    resolve(response(JSON.stringify(responseData), ContentType.JSON, CacheControl.NoStore));
   });
 
 type MessageResolve = [(data?: any) => void, any];