From 070c75c46d1ccf841a27c0fa095fc47e66c718cc Mon Sep 17 00:00:00 2001
From: Chris Villa <chrisvilla@me.com>
Date: Wed, 28 Feb 2024 10:16:10 +0100
Subject: [PATCH] refactor: revert use-style-marshal and rely on user-land
 style syncing

---
 .../use-style-marshal/use-style-marshal.ts    | 119 +++++++-----------
 1 file changed, 45 insertions(+), 74 deletions(-)

diff --git a/src/view/use-style-marshal/use-style-marshal.ts b/src/view/use-style-marshal/use-style-marshal.ts
index 672773bc3..7b2e64703 100644
--- a/src/view/use-style-marshal/use-style-marshal.ts
+++ b/src/view/use-style-marshal/use-style-marshal.ts
@@ -1,3 +1,4 @@
+import { useRef, MutableRefObject } from 'react';
 import memoizeOne from 'memoize-one';
 import { useMemo, useCallback } from 'use-memo-one';
 import { invariant } from '../../invariant';
@@ -7,11 +8,10 @@ import getStyles from './get-styles';
 import type { Styles } from './get-styles';
 import { prefix } from '../data-attributes';
 import useLayoutEffect from '../use-isomorphic-layout-effect';
-import { querySelectorAll } from '../../query-selector-all';
-import querySelectorAllIframe from '../iframe/query-selector-all-iframe';
 
-const getHead = (doc: Document): HTMLHeadElement | null => {
-  const head: HTMLHeadElement | null = doc.querySelector('head');
+const getHead = (): HTMLHeadElement => {
+  const head: HTMLHeadElement | null = document.querySelector('head');
+  invariant(head, 'Cannot find the head to append a style to');
   return head;
 };
 
@@ -24,97 +24,64 @@ const createStyleEl = (nonce?: string): HTMLStyleElement => {
   return el;
 };
 
-const alwaysDataAttr = `${prefix}-always`;
-const dynamicDataAttr = `${prefix}-dynamic`;
-
 export default function useStyleMarshal(contextId: ContextId, nonce?: string) {
   const styles: Styles = useMemo(() => getStyles(contextId), [contextId]);
+  const alwaysRef = useRef<HTMLStyleElement | null>(null);
+  const dynamicRef = useRef<HTMLStyleElement | null>(null);
 
   // eslint-disable-next-line react-hooks/exhaustive-deps
   const setDynamicStyle = useCallback(
     // Using memoizeOne to prevent frequent updates to textContext
     memoizeOne((proposed: string) => {
-      const selector = `[${dynamicDataAttr}="${contextId}"]`;
-
-      querySelectorAllIframe(selector).forEach((el) => {
-        invariant(el, 'Cannot set dynamic style element if it is not set');
-        el.textContent = proposed;
-      });
+      const el: HTMLStyleElement | null = dynamicRef.current;
+      invariant(el, 'Cannot set dynamic style element if it is not set');
+      el.textContent = proposed;
     }),
-    [contextId],
+    [],
   );
 
-  const setAlwaysStyle = useCallback(
-    (proposed: string) => {
-      const selector = `[${alwaysDataAttr}="${contextId}"]`;
-
-      querySelectorAllIframe(selector).forEach((el) => {
-        invariant(el, 'Cannot set dynamic style element if it is not set');
-        el.textContent = proposed;
-      });
-    },
-    [contextId],
-  );
+  const setAlwaysStyle = useCallback((proposed: string) => {
+    const el: HTMLStyleElement | null = alwaysRef.current;
+    invariant(el, 'Cannot set dynamic style element if it is not set');
+    el.textContent = proposed;
+  }, []);
 
   // using layout effect as programatic dragging might start straight away (such as for cypress)
   useLayoutEffect(() => {
-    const alwaysSelector = `[${alwaysDataAttr}="${contextId}"]`;
-    const dynamicSelector = `[${dynamicDataAttr}="${contextId}"]`;
-
-    const heads = [
-      getHead(document),
-      ...(
-        querySelectorAll(document, `[${prefix}-iframe]`) as HTMLIFrameElement[]
-      ).map((iframe) => getHead(iframe.contentWindow!.document)),
-    ];
-
-    // Create initial style elements
-    heads.forEach((head) => {
-      if (!head) return;
-
-      const alwaysElements = querySelectorAll(
-        head.ownerDocument,
-        alwaysSelector,
-      );
-      const dynamicElements = querySelectorAll(
-        head.ownerDocument,
-        dynamicSelector,
-      );
-
-      if (
-        alwaysElements.length >= heads.length ||
-        dynamicElements.length >= heads.length
-      ) {
-        return;
-      }
+    invariant(
+      !alwaysRef.current && !dynamicRef.current,
+      'style elements already mounted',
+    );
 
-      const always: HTMLStyleElement = createStyleEl(nonce);
-      const dynamic: HTMLStyleElement = createStyleEl(nonce);
+    const always: HTMLStyleElement = createStyleEl(nonce);
+    const dynamic: HTMLStyleElement = createStyleEl(nonce);
 
-      // for easy identification
-      always.setAttribute(alwaysDataAttr, contextId);
-      dynamic.setAttribute(dynamicDataAttr, contextId);
+    // store their refs
+    alwaysRef.current = always;
+    dynamicRef.current = dynamic;
 
-      head.appendChild(always);
-      head.appendChild(dynamic);
+    // for easy identification
+    always.setAttribute(`${prefix}-always`, contextId);
+    dynamic.setAttribute(`${prefix}-dynamic`, contextId);
 
-      // set initial style
-      setAlwaysStyle(styles.always);
-      setDynamicStyle(styles.resting);
-    });
+    // add style tags to head
+    getHead().appendChild(always);
+    getHead().appendChild(dynamic);
 
-    return () => {
-      const remove = (selector: string) => {
-        const elements = querySelectorAllIframe(selector);
+    // set initial style
+    setAlwaysStyle(styles.always);
+    setDynamicStyle(styles.resting);
 
-        elements.forEach((el) => {
-          invariant(el, 'Cannot unmount element as it is not set');
-          el.ownerDocument.head.removeChild(el);
-        });
+    return () => {
+      const remove = (ref: MutableRefObject<HTMLStyleElement | null>) => {
+        const current: HTMLStyleElement | null = ref.current;
+        invariant(current, 'Cannot unmount ref as it is not set');
+        getHead().removeChild(current);
+        ref.current = null;
       };
 
-      remove(alwaysSelector);
-      remove(dynamicSelector);
+      remove(alwaysRef);
+      remove(dynamicRef);
     };
   }, [
     nonce,
@@ -140,6 +107,10 @@ export default function useStyleMarshal(contextId: ContextId, nonce?: string) {
     [setDynamicStyle, styles.dropAnimating, styles.userCancel],
   );
   const resting = useCallback(() => {
+    // Can be called defensively
+    if (!dynamicRef.current) {
+      return;
+    }
     setDynamicStyle(styles.resting);
   }, [setDynamicStyle, styles.resting]);