fix: inserted styles lost when moving elements (#233)

fix code for nodejs tests
change fix direction to avoid issues with duplicate styles
format issues
swap waitForTimeout for waitForRAF in test that flaked
Add unit tests for new functions
Fix broken test causes by file formatting removing spaced

---------

Co-authored-by: jaj1014 <ajax@pendo.io>
Co-authored-by: jaj1014 <jaj1014@users.noreply.github.com>

packages/rrdom/src/diff.ts

@@ -390,7 +390,7 @@ function diffChildren(
       nodeMatching(oldStartNode, newEndNode, replayer.mirror, rrnodeMirror)
     ) {
       try {
-        oldTree.insertBefore(oldStartNode, oldEndNode.nextSibling);
+        handleInsertBefore(oldTree, oldStartNode, oldEndNode.nextSibling);
       } catch (e) {
         console.warn(e);
       }
@@ -401,7 +401,7 @@ function diffChildren(
       nodeMatching(oldEndNode, newStartNode, replayer.mirror, rrnodeMirror)
     ) {
       try {
-        oldTree.insertBefore(oldEndNode, oldStartNode);
+        handleInsertBefore(oldTree, oldEndNode, oldStartNode);
       } catch (e) {
         console.warn(e);
       }
@@ -426,7 +426,7 @@ function diffChildren(
         nodeMatching(nodeToMove, newStartNode, replayer.mirror, rrnodeMirror)
       ) {
         try {
-          oldTree.insertBefore(nodeToMove, oldStartNode);
+          handleInsertBefore(oldTree, nodeToMove, oldStartNode);
         } catch (e) {
           console.warn(e);
         }
@@ -460,7 +460,7 @@ function diffChildren(
         }
 
         try {
-          oldTree.insertBefore(newNode, oldStartNode || null);
+          handleInsertBefore(oldTree, newNode, oldStartNode || null);
         } catch (e) {
           console.warn(e);
         }
@@ -482,7 +482,7 @@ function diffChildren(
         rrnodeMirror,
       );
       try {
-        oldTree.insertBefore(newNode, referenceNode);
+        handleInsertBefore(oldTree, newNode, referenceNode);
       } catch (e) {
         console.warn(e);
       }
@@ -590,3 +590,64 @@ export function nodeMatching(
   if (node1Id === -1 || node1Id !== node2Id) return false;
   return sameNodeType(node1, node2);
 }
+
+/**
+ * Copies CSSRules and their position from HTML style element which don't exist in it's innerText
+ */
+function getInsertedStylesFromElement(
+  styleElement: HTMLStyleElement,
+): Array<{ index: number; cssRuleText: string }> | undefined {
+  const elementCssRules = styleElement.sheet?.cssRules;
+  if (!elementCssRules || !elementCssRules.length) return;
+  // style sheet w/ innerText styles to diff with actual and get only inserted styles
+  const tempStyleSheet = new CSSStyleSheet();
+  tempStyleSheet.replaceSync(styleElement.innerText);
+
+  const innerTextStylesMap: { [key: string]: CSSRule } = {};
+
+  for (let i = 0; i < tempStyleSheet.cssRules.length; i++) {
+    innerTextStylesMap[tempStyleSheet.cssRules[i].cssText] =
+      tempStyleSheet.cssRules[i];
+  }
+
+  const insertedStylesStyleSheet = [];
+
+  for (let i = 0; i < elementCssRules?.length; i++) {
+    const cssRuleText = elementCssRules[i].cssText;
+
+    if (!innerTextStylesMap[cssRuleText]) {
+      insertedStylesStyleSheet.push({
+        index: i,
+        cssRuleText,
+      });
+    }
+  }
+
+  return insertedStylesStyleSheet;
+}
+
+/**
+ * Conditionally copy insertedStyles for STYLE nodes and apply after calling insertBefore'
+ * For non-STYLE nodes, just insertBefore
+ */
+export function handleInsertBefore(
+  oldTree: Node,
+  nodeToMove: Node,
+  insertBeforeNode: Node | null,
+): void {
+  let insertedStyles;
+
+  if (nodeToMove.nodeName === 'STYLE') {
+    insertedStyles = getInsertedStylesFromElement(
+      nodeToMove as HTMLStyleElement,
+    );
+  }
+
+  oldTree.insertBefore(nodeToMove, insertBeforeNode);
+
+  if (insertedStyles && insertedStyles.length) {
+    insertedStyles.forEach(({ cssRuleText, index }) => {
+      (nodeToMove as HTMLStyleElement).sheet?.insertRule(cssRuleText, index);
+    });
+  }
+}

packages/rrdom/test/diff.test.ts

@@ -23,6 +23,7 @@ import {
   ReplayerHandler,
   nodeMatching,
   sameNodeType,
+  handleInsertBefore,
 } from '../src/diff';
 import type { IRRElement, IRRNode } from '../src/document';
 import { Replayer } from '@sentry-internal/rrweb';
@@ -1469,7 +1470,7 @@ describe('diff algorithm for rrdom', () => {
           const rrHtmlEl = rrDocument.createElement('html');
           rrDocument.mirror.add(rrHtmlEl, rrdom.getDefaultSN(rrHtmlEl, ${htmlElId}));
           rrIframeEl.contentDocument.appendChild(rrHtmlEl);
-          
+
           const replayer = {
             mirror: rrdom.createMirror(),
             applyCanvas: () => {},
@@ -1478,7 +1479,7 @@ describe('diff algorithm for rrdom', () => {
             applyStyleSheetMutation: () => {},
           };
           rrdom.diff(iframeEl, rrIframeEl, replayer);
-          
+
           iframeEl.contentDocument.documentElement.className =
             '${className.toLowerCase()}';
           iframeEl.contentDocument.childNodes.length === 2 &&
@@ -2000,4 +2001,69 @@ describe('diff algorithm for rrdom', () => {
       expect(nodeMatching(node1, node2, NodeMirror, rrdomMirror)).toBeFalsy();
     });
   });
+
+  describe('test handleInsertBefore function', () => {
+    it('should insert nodeToMove before insertBeforeNode in oldTree for non-style elements', () => {
+      const oldTree = document.createElement('div');
+      const nodeToMove = document.createElement('div');
+      const insertBeforeNode = document.createElement('div');
+      oldTree.appendChild(insertBeforeNode);
+
+      expect(oldTree.children.length).toEqual(1);
+
+      handleInsertBefore(oldTree, nodeToMove, insertBeforeNode);
+
+      expect(oldTree.children.length).toEqual(2);
+      expect(oldTree.children[0]).toEqual(nodeToMove);
+    });
+
+    it('should not drop inserted styles when moving a style element with inserted styles', async () => {
+      function MockCSSStyleSheet() {
+        this.replaceSync = jest.fn();
+        this.cssRules = [{ cssText: baseStyle }];
+      }
+
+      jest
+        .spyOn(window, 'CSSStyleSheet')
+        .mockImplementationOnce(MockCSSStyleSheet as any);
+
+      const baseStyle = 'body {margin: 0;}';
+      const insertedStyle = 'div {display: flex;}';
+
+      document.write('<html></html>');
+
+      const insertBeforeNode = document.createElement('style');
+      document.documentElement.appendChild(insertBeforeNode);
+
+      const nodeToMove = document.createElement('style');
+      nodeToMove.appendChild(document.createTextNode(baseStyle));
+      document.documentElement.appendChild(nodeToMove);
+      nodeToMove.sheet?.insertRule(insertedStyle);
+
+      // validate dom prior to moving element
+      expect(document.documentElement.children.length).toEqual(4);
+      expect(document.documentElement.children[2]).toEqual(insertBeforeNode);
+      expect(document.documentElement.children[3]).toEqual(nodeToMove);
+      expect(nodeToMove.sheet?.cssRules.length).toEqual(2);
+      expect(nodeToMove.sheet?.cssRules[0].cssText).toEqual(insertedStyle);
+      expect(nodeToMove.sheet?.cssRules[1].cssText).toEqual(baseStyle);
+
+      // move the node
+      handleInsertBefore(
+        document.documentElement,
+        nodeToMove,
+        insertBeforeNode,
+      );
+
+      // nodeToMove was inserted before
+      expect(document.documentElement.children.length).toEqual(4);
+      expect(document.documentElement.children[2]).toEqual(nodeToMove);
+      expect(document.documentElement.children[3]).toEqual(insertBeforeNode);
+      // styles persisted on the moved element
+      // w/ document.documentElement.insertBefore(nodeToMove, insertBeforeNode) insertedStyle wouldn't be copied
+      expect(nodeToMove.sheet?.cssRules.length).toEqual(2);
+      expect(nodeToMove.sheet?.cssRules[0].cssText).toEqual(insertedStyle);
+      expect(nodeToMove.sheet?.cssRules[1].cssText).toEqual(baseStyle);
+    });
+  });
 });