call afterAppend hook in a consistent traversal order

packages/rrdom/src/diff.ts

@@ -1,4 +1,8 @@
-import { NodeType as RRNodeType, Mirror as NodeMirror } from 'rrweb-snapshot';
+import {
+  NodeType as RRNodeType,
+  Mirror as NodeMirror,
+  elementNode,
+} from 'rrweb-snapshot';
 import type {
   canvasMutationData,
   canvasEventWithTime,
@@ -88,6 +92,9 @@ export type ReplayerHandler = {
   afterAppend?(node: Node, id: number): void;
 };
 
+// A set contains newly appended nodes. It's used to make sure the afterAppend callback can iterate newly appended nodes in the same traversal order as that in the `rrweb-snapshot` package.
+let createdNodeSet: WeakSet<Node> | null = null;
+
 /**
  * Make the old tree to have the same structure and properties as the new tree with the diff algorithm.
  * @param oldTree - The old tree to be modified.
@@ -102,19 +109,12 @@ export function diff(
   rrnodeMirror: Mirror = (newTree as RRDocument).mirror ||
     (newTree.ownerDocument as RRDocument).mirror,
 ) {
-  // If the Mirror data has some flaws, the diff function may throw errors. We check the node consistency here to make it robust.
-  if (!sameNodeType(oldTree, newTree)) {
-    const calibratedOldTree = createOrGetNode(
-      newTree,
-      replayer.mirror,
-      rrnodeMirror,
-    );
-    oldTree.parentNode?.replaceChild(calibratedOldTree, oldTree);
-    oldTree = calibratedOldTree;
-    replayer.afterAppend?.(oldTree, replayer.mirror.getId(oldTree));
-  }
-
-  diffBeforeUpdatingChildren(oldTree, newTree, replayer, rrnodeMirror);
+  oldTree = diffBeforeUpdatingChildren(
+    oldTree,
+    newTree,
+    replayer,
+    rrnodeMirror,
+  );
 
   const oldChildren = oldTree.childNodes;
   const newChildren = newTree.childNodes;
@@ -140,6 +140,22 @@ function diffBeforeUpdatingChildren(
   replayer: ReplayerHandler,
   rrnodeMirror: Mirror,
 ) {
+  if (replayer.afterAppend && !createdNodeSet) {
+    createdNodeSet = new WeakSet();
+    setTimeout(() => {
+      createdNodeSet = null;
+    }, 0);
+  }
+  // If the Mirror data has some flaws, the diff function may throw errors. We check the node consistency here to make it robust.
+  if (!sameNodeType(oldTree, newTree)) {
+    const calibratedOldTree = createOrGetNode(
+      newTree,
+      replayer.mirror,
+      rrnodeMirror,
+    );
+    oldTree.parentNode?.replaceChild(calibratedOldTree, oldTree);
+    oldTree = calibratedOldTree;
+  }
   switch (newTree.RRNodeType) {
     case RRNodeType.Document: {
       /**
@@ -154,7 +170,7 @@ function diffBeforeUpdatingChildren(
           (oldTree as Document).close();
           (oldTree as Document).open();
           replayer.mirror.add(oldTree, newMeta);
-          replayer.afterAppend?.(oldTree, replayer.mirror.getId(oldTree));
+          createdNodeSet?.add(oldTree);
         }
       }
       break;
@@ -196,6 +212,7 @@ function diffBeforeUpdatingChildren(
       break;
     }
   }
+  return oldTree;
 }
 
 /**
@@ -291,6 +308,10 @@ function diffAfterUpdatingChildren(
       break;
     }
   }
+  if (createdNodeSet?.has(oldTree)) {
+    createdNodeSet.delete(oldTree);
+    replayer.afterAppend?.(oldTree, replayer.mirror.getId(oldTree));
+  }
 }
 
 function diffProps(
@@ -303,8 +324,8 @@ function diffProps(
 
   for (const name in newAttributes) {
     const newValue = newAttributes[name];
-    const sn = rrnodeMirror.getMeta(newTree);
-    if (sn && 'isSVG' in sn && sn.isSVG && NAMESPACES[name])
+    const sn = rrnodeMirror.getMeta(newTree) as elementNode | null;
+    if (sn?.isSVG && NAMESPACES[name])
       oldTree.setAttributeNS(NAMESPACES[name], name, newValue);
     else if (newTree.tagName === 'CANVAS' && name === 'rr_dataURL') {
       const image = document.createElement('img');
@@ -441,7 +462,6 @@ function diffChildren(
         try {
           parentNode.insertBefore(newNode, oldStartNode || null);
           diff(newNode, newStartNode, replayer, rrnodeMirror);
-          replayer.afterAppend?.(newNode, replayer.mirror.getId(newNode));
         } catch (e) {
           console.warn(e);
         }
@@ -465,7 +485,6 @@ function diffChildren(
       try {
         parentNode.insertBefore(newNode, referenceNode);
         diff(newNode, newChildren[newStartIndex], replayer, rrnodeMirror);
-        replayer.afterAppend?.(newNode, replayer.mirror.getId(newNode));
       } catch (e) {
         console.warn(e);
       }
@@ -526,6 +545,11 @@ export function createOrGetNode(
   }
 
   if (sn) domMirror.add(node, { ...sn });
+  try {
+    createdNodeSet?.add(node);
+  } catch (e) {
+    // Just for safety concern.
+  }
   return node;
 }
 

packages/rrdom/test/diff.test.ts

@@ -33,7 +33,6 @@ import type {
 } from '@rrweb/types';
 import { EventType, IncrementalSource } from '@rrweb/types';
 import { compileTSCode } from './utils';
-import { printRRDom } from '../src/index';
 
 const elementSn = {
   type: RRNodeType.Element,
@@ -114,7 +113,9 @@ describe('diff algorithm for rrdom', () => {
       applyInput: () => {},
       applyScroll: () => {},
       applyStyleSheetMutation: () => {},
+      afterAppend: () => {},
     };
+    document.write('<!DOCTYPE html><html><head></head><body></body></html>');
   });
 
   describe('diff single node', () => {
@@ -133,7 +134,7 @@ describe('diff algorithm for rrdom', () => {
         x: 0,
         y: 0,
       };
-      replayer.applyScroll = jest.fn();
+      const applyScrollFn = jest.spyOn(replayer, 'applyScroll');
       diff(document, rrNode, replayer);
       expect(document.childNodes.length).toEqual(1);
       expect(document.childNodes[0]).toBeInstanceOf(DocumentType);
@@ -142,7 +143,24 @@ describe('diff algorithm for rrdom', () => {
         '-//W3C//DTD XHTML 1.0 Transitional//EN',
       );
       expect(document.doctype?.systemId).toEqual('');
-      expect(replayer.applyScroll).toBeCalledTimes(1);
+      expect(applyScrollFn).toHaveBeenCalledTimes(1);
+      applyScrollFn.mockRestore();
+    });
+
+    it('should apply scroll data on an element', () => {
+      const element = document.createElement('div');
+      const rrDocument = new RRDocument();
+      const rrNode = rrDocument.createElement('div');
+      rrNode.scrollData = {
+        source: IncrementalSource.Scroll,
+        id: 0,
+        x: 0,
+        y: 0,
+      };
+      const applyScrollFn = jest.spyOn(replayer, 'applyScroll');
+      diff(element, rrNode, replayer);
+      expect(applyScrollFn).toHaveBeenCalledTimes(1);
+      applyScrollFn.mockRestore();
     });
 
     it('should apply input data on an input element', () => {
@@ -1442,6 +1460,159 @@ describe('diff algorithm for rrdom', () => {
     });
   });
 
+  describe('afterAppend callback', () => {
+    it('should call afterAppend callback', () => {
+      const afterAppendFn = jest.spyOn(replayer, 'afterAppend');
+      const node = createTree(
+        {
+          tagName: 'div',
+          id: 1,
+        },
+        undefined,
+        mirror,
+      ) as Node;
+
+      const rrdom = new RRDocument();
+      const rrNode = createTree(
+        {
+          tagName: 'div',
+          id: 1,
+          children: [
+            {
+              tagName: 'span',
+              id: 2,
+            },
+          ],
+        },
+        rrdom,
+      ) as RRNode;
+      diff(node, rrNode, replayer);
+      expect(afterAppendFn).toHaveBeenCalledTimes(1);
+      expect(afterAppendFn).toHaveBeenCalledWith(node.childNodes[0], 2);
+      afterAppendFn.mockRestore();
+    });
+
+    it('should diff without afterAppend callback', () => {
+      replayer.afterAppend = undefined;
+      const rrdom = buildFromDom(document);
+      document.open();
+      diff(document, rrdom, replayer);
+      replayer.afterAppend = () => {};
+    });
+
+    it('should call afterAppend callback in the post traversal order', () => {
+      const afterAppendFn = jest.spyOn(replayer, 'afterAppend');
+      document.open();
+
+      const rrdom = new RRDocument();
+      rrdom.mirror.add(rrdom, getDefaultSN(rrdom, 1));
+      const rrNode = createTree(
+        {
+          tagName: 'html',
+          id: 1,
+          children: [
+            {
+              tagName: 'head',
+              id: 2,
+            },
+            {
+              tagName: 'body',
+              id: 3,
+              children: [
+                {
+                  tagName: 'span',
+                  id: 4,
+                  children: [
+                    {
+                      tagName: 'li',
+                      id: 5,
+                    },
+                    {
+                      tagName: 'li',
+                      id: 6,
+                    },
+                  ],
+                },
+                {
+                  tagName: 'p',
+                  id: 7,
+                },
+                {
+                  tagName: 'p',
+                  id: 8,
+                  children: [
+                    {
+                      tagName: 'li',
+                      id: 9,
+                    },
+                    {
+                      tagName: 'li',
+                      id: 10,
+                    },
+                  ],
+                },
+              ],
+            },
+          ],
+        },
+        rrdom,
+      ) as RRNode;
+      diff(document, rrNode, replayer);
+
+      expect(afterAppendFn).toHaveBeenCalledTimes(10);
+      // the correct traversal order
+      [2, 5, 6, 4, 7, 9, 10, 8, 3, 1].forEach((id, index) => {
+        expect((mirror.getNode(id) as HTMLElement).tagName).toEqual(
+          (rrdom.mirror.getNode(id) as IRRElement).tagName,
+        );
+        expect(afterAppendFn).toHaveBeenNthCalledWith(
+          index + 1,
+          mirror.getNode(id),
+          id,
+        );
+      });
+    });
+
+    it('should only call afterAppend for newly created nodes', () => {
+      const afterAppendFn = jest.spyOn(replayer, 'afterAppend');
+      const rrdom = buildFromDom(document, replayer.mirror) as RRDocument;
+
+      // Append 3 nodes to rrdom.
+      const rrNode = createTree(
+        {
+          tagName: 'span',
+          id: 1,
+          children: [
+            {
+              tagName: 'li',
+              id: 2,
+            },
+            {
+              tagName: 'li',
+              id: 3,
+            },
+          ],
+        },
+        rrdom,
+      ) as RRNode;
+      rrdom.body?.appendChild(rrNode);
+      diff(document, rrdom, replayer);
+      expect(afterAppendFn).toHaveBeenCalledTimes(3);
+      // Should only call afterAppend for 3 newly appended nodes.
+      [2, 3, 1].forEach((id, index) => {
+        expect((mirror.getNode(id) as HTMLElement).tagName).toEqual(
+          (rrdom.mirror.getNode(id) as IRRElement).tagName,
+        );
+        expect(afterAppendFn).toHaveBeenNthCalledWith(
+          index + 1,
+          mirror.getNode(id),
+          id,
+        );
+      });
+      afterAppendFn.mockClear();
+    });
+  });
+
   describe('create or get a Node', () => {
     it('create a real HTML element from RRElement', () => {
       const rrDocument = new RRDocument();