fix: document reference on location change

QwikDev · Mar 8, 2022 · fdc7b53 · fdc7b53 · vercel · Mar 8, 2022
1 parent 833ba36
commit fdc7b53
Show file tree

Hide file tree

Showing 7 changed files with 63 additions and 35 deletions.
diff --git a/src/lib/sandbox/main-access-handler.ts b/src/lib/sandbox/main-access-handler.ts
@@ -67,7 +67,16 @@ export const mainAccessHandler = async (
           rtnValue = applyToInstance(worker, instance, applyPath, isLast, task.$groupedGetters$);
 
           if (task.$assignInstanceId$) {
-            setInstanceId(rtnValue, task.$assignInstanceId$);
+            if (typeof task.$assignInstanceId$ === 'string') {
+              setInstanceId(rtnValue, task.$assignInstanceId$);
+            } else {
+              winCtxs[task.$assignInstanceId$.$winId$] = {
+                $winId$: task.$assignInstanceId$.$winId$,
+                $window$: {
+                  document: rtnValue,
+                } as any,
+              };
+            }
           }
 
           if (isPromise(rtnValue)) {

diff --git a/src/lib/sandbox/main-instances.ts b/src/lib/sandbox/main-instances.ts
@@ -1,5 +1,5 @@
 import { CreatedKey, InstanceIdKey, instances, winCtxs, windowIds } from './main-constants';
-import type { InstanceId, WinId } from '../types';
+import { InstanceId, MainWindowContext, WinDocId, WinId } from '../types';
 import { randomId } from '../utils';
 
 export const getAndSetInstanceId = (instance: any, instanceId?: InstanceId) => {
@@ -17,18 +17,32 @@ export const getAndSetInstanceId = (instance: any, instanceId?: InstanceId) => {
 export const getInstance = <T = any>(
   winId: WinId,
   instanceId: InstanceId,
+  win?: MainWindowContext,
+  doc?: Document,
   docId?: string
 ): T | undefined => {
-  if (winId === instanceId && winCtxs[winId] && winCtxs[winId]!.$window$) {
-    return winCtxs[winId]!.$window$ as any;
-  } else {
-    [instanceId, docId] = instanceId.split('.');
-    if (docId) {
-      return instances.get(instanceId)[docId];
-    } else {
-      return instances.get(instanceId);
+  if ((win = winCtxs[winId]) && win.$window$) {
+    if (winId === instanceId) {
+      return win.$window$ as any;
+    }
+
+    doc = win.$window$.document;
+    docId = instanceId.split('.').pop();
+    if (docId === WinDocId.document) {
+      return doc as any;
+    }
+    if (docId === WinDocId.documentElement) {
+      return doc.documentElement as any;
+    }
+    if (docId === WinDocId.head) {
+      return doc.head as any;
+    }
+    if (docId === WinDocId.body) {
+      return doc.body as any;
     }
   }
+
+  return instances.get(instanceId);
 };
 
 export const setInstanceId = (instance: any, instanceId: InstanceId, now?: number) => {

diff --git a/src/lib/sandbox/main-register-window.ts b/src/lib/sandbox/main-register-window.ts
@@ -1,7 +1,6 @@
 import { debug } from '../utils';
 import { logMain, normalizedWinId } from '../log';
-import { MainWindow, PartytownWebWorker, WinDocId, WinId, WorkerMessageType } from '../types';
-import { setInstanceId } from './main-instances';
+import { MainWindow, PartytownWebWorker, WinId, WorkerMessageType } from '../types';
 import { winCtxs, windowIds } from './main-constants';
 
 export const registerWindow = (
@@ -45,8 +44,6 @@ export const registerWindow = (
       onLocationChange();
     };
 
-    setInstanceId(doc, $winId$ + WinDocId.document);
-
     $window$.addEventListener('popstate', onLocationChange);
     $window$.addEventListener('hashchange', onLocationChange);
     doc.addEventListener('visibilitychange', () =>

diff --git a/src/lib/types.ts b/src/lib/types.ts
@@ -22,6 +22,12 @@ export type InstanceId = string;
 
 export type RefId = string;
 
+export interface AssignWinInstanceId {
+  $winId$: WinId;
+}
+
+export type AssignInstanceId = InstanceId | AssignWinInstanceId;
+
 export type MessageFromWorkerToSandbox =
   | [type: WorkerMessageType.MainDataRequestFromWorker]
   | [type: WorkerMessageType.InitializedWebWorker]
@@ -191,9 +197,9 @@ export const enum InterfaceType {
 
 export const enum WinDocId {
   document = 'd',
-  documentElement = 'd.documentElement',
-  head = 'd.head',
-  body = 'd.body',
+  documentElement = 'e',
+  head = 'h',
+  body = 'b',
 }
 
 export interface InitializeScriptData {
@@ -214,7 +220,7 @@ export interface MainAccessTask {
   $instanceId$: InstanceId;
   $applyPath$: ApplyPath;
   $groupedGetters$?: string[];
-  $assignInstanceId$?: InstanceId;
+  $assignInstanceId$?: AssignInstanceId;
   $debug$?: string;
 }
 
@@ -542,7 +548,7 @@ export type CallMethod = (
   applyPath: ApplyPath,
   args: any[],
   callType?: CallType,
-  assignInstanceId?: InstanceId,
+  assignInstanceId?: AssignInstanceId,
   buffer?: ArrayBuffer | ArrayBufferView
 ) => any;
 

diff --git a/src/lib/web-worker/worker-document.ts b/src/lib/web-worker/worker-document.ts
@@ -173,16 +173,18 @@ export const patchDocument = (
         return {
           hasFeature: () => true,
           createHTMLDocument: (title: string) => {
-            const winId = randomId();
-            const docId = winId + WinDocId.document;
-            callMethod(
-              this,
-              ['implementation', 'createHTMLDocument'],
-              [title],
-              CallType.Blocking,
-              docId
+            const $winId$ = randomId();
+            callMethod(this, ['implementation', 'createHTMLDocument'], [title], CallType.Blocking, {
+              $winId$,
+            });
+            const docEnv = createWindow(
+              $winId$,
+              $winId$,
+              env.$location$ + '',
+              'hidden',
+              true,
+              true
             );
-            const docEnv = createWindow(winId, winId, env.$location$ + '', 'hidden', true, true);
             return docEnv.$document$;
           },
         };

diff --git a/src/lib/web-worker/worker-proxy.ts b/src/lib/web-worker/worker-proxy.ts
@@ -1,12 +1,12 @@
 import {
   ApplyPath,
   ApplyPathType,
+  AssignInstanceId,
   CallMethod,
   CallType,
   ConstructGlobal,
   Getter,
   HookOptions,
-  InstanceId,
   MainAccessRequest,
   MainAccessResponse,
   MainAccessTask,
@@ -52,7 +52,7 @@ const queue = (
   instance: WorkerInstance,
   $applyPath$: ApplyPath,
   callType: CallType,
-  $assignInstanceId$?: InstanceId,
+  $assignInstanceId$?: AssignInstanceId,
   $groupedGetters$?: string[],
   buffer?: ArrayBuffer | ArrayBufferView
 ) => {
@@ -194,7 +194,7 @@ export const callMethod: CallMethod = (
   applyPath: ApplyPath,
   args: any[],
   callType?: CallType,
-  assignInstanceId?: InstanceId,
+  assignInstanceId?: AssignInstanceId,
   buffer?: ArrayBuffer | ArrayBufferView,
   rtnValue?: any,
   methodName?: string

diff --git a/src/lib/web-worker/worker-window.ts b/src/lib/web-worker/worker-window.ts
@@ -339,13 +339,13 @@ export const createWindow = (
               // like: <script>globalProp = true</script>
               true,
           }) as any,
-          $document$: $createNode$(NodeName.Document, $winId$ + WinDocId.document) as any,
+          $document$: $createNode$(NodeName.Document, $winId$ + '.' + WinDocId.document) as any,
           $documentElement$: $createNode$(
             NodeName.DocumentElement,
-            $winId$ + WinDocId.documentElement
+            $winId$ + '.' + WinDocId.documentElement
           ) as any,
-          $head$: $createNode$(NodeName.Head, $winId$ + WinDocId.head) as any,
-          $body$: $createNode$(NodeName.Body, $winId$ + WinDocId.body) as any,
+          $head$: $createNode$(NodeName.Head, $winId$ + '.' + WinDocId.head) as any,
+          $body$: $createNode$(NodeName.Body, $winId$ + '.' + WinDocId.body) as any,
           $location$,
           $visibilityState$,
           $isSameOrigin$,