From e54218be19240a65940c9f6f2165d934bfd84848 Mon Sep 17 00:00:00 2001
From: Yanzhen Yu <yanzhen@smartx.com>
Date: Thu, 6 Aug 2020 21:37:52 +0800
Subject: [PATCH] provide a general visit function to modify rebuilt id node
 map

---
 src/rebuild.ts | 62 +++++++++++++++++++++++++++++++++-----------------
 1 file changed, 41 insertions(+), 21 deletions(-)

diff --git a/src/rebuild.ts b/src/rebuild.ts
index bb26aefa..23025797 100644
--- a/src/rebuild.ts
+++ b/src/rebuild.ts
@@ -62,7 +62,7 @@ export function addHoverClass(cssText: string): string {
   if (!ast.stylesheet) {
     return cssText;
   }
-  ast.stylesheet.rules.forEach(rule => {
+  ast.stylesheet.rules.forEach((rule) => {
     if ('selectors' in rule) {
       (rule.selectors || []).forEach((selector: string) => {
         if (HOVER_SELECTOR.test(selector)) {
@@ -102,7 +102,8 @@ function buildNode(
           continue;
         }
         let value = n.attributes[name];
-        value = typeof value === 'boolean' || typeof value === 'number' ? '' : value;
+        value =
+          typeof value === 'boolean' || typeof value === 'number' ? '' : value;
         // attribute names start with rr_ are internal attributes added by rrweb
         if (!name.startsWith('rr_')) {
           const isTextarea = tagName === 'textarea' && name === 'value';
@@ -128,7 +129,11 @@ function buildNode(
           try {
             if (n.isSVG && name === 'xlink:href') {
               node.setAttributeNS('http://www.w3.org/1999/xlink', name, value);
-            } else if (name == 'onload' || name == 'onclick' || name.substring(0, 7) == 'onmouse') {
+            } else if (
+              name === 'onload' ||
+              name === 'onclick' ||
+              name.substring(0, 7) === 'onmouse'
+            ) {
               // Rename some of the more common atttributes from https://www.w3schools.com/tags/ref_eventattributes.asp
               // as setting them triggers a console.error (which shows up despite the try/catch)
               // Assumption: these attributes are not used to css
@@ -220,25 +225,34 @@ export function buildNodeWithSN(
   return node as INode;
 }
 
-function sideEffects(idNodeMap: idNodeMap) {
-  for(let id in idNodeMap) {
-    const node = idNodeMap[id];
-    const n = node.__sn;
-    if (n.type !== NodeType.Element) {
+function visit(idNodeMap: idNodeMap, onVisit: (node: INode) => void) {
+  function walk(node: INode) {
+    onVisit(node);
+  }
+
+  for (const key in idNodeMap) {
+    if (idNodeMap[key]) {
+      walk(idNodeMap[key]);
+    }
+  }
+}
+
+function handleScroll(node: INode) {
+  const n = node.__sn;
+  if (n.type !== NodeType.Element) {
+    return;
+  }
+  const el = (node as Node) as HTMLElement;
+  for (const name in n.attributes) {
+    if (!(n.attributes.hasOwnProperty(name) && name.startsWith('rr_'))) {
       continue;
     }
-    const el = node as Node as HTMLElement;
-    for(const name in n.attributes) {
-      if (!(n.attributes.hasOwnProperty(name) && name.startsWith('rr_'))) {
-        continue;
-      }
-      const value = n.attributes[name];
-      if (name === 'rr_scrollLeft') {
-        el.scrollLeft = value as number;
-      }
-      if (name === 'rr_scrollTop') {
-        el.scrollTop = value as number;
-      }
+    const value = n.attributes[name];
+    if (name === 'rr_scrollLeft') {
+      el.scrollLeft = value as number;
+    }
+    if (name === 'rr_scrollTop') {
+      el.scrollTop = value as number;
     }
   }
 }
@@ -246,6 +260,7 @@ function sideEffects(idNodeMap: idNodeMap) {
 function rebuild(
   n: serializedNodeWithId,
   doc: Document,
+  onVisit?: (node: INode) => unknown,
   /**
    * This is not a public API yet, just for POC
    */
@@ -253,7 +268,12 @@ function rebuild(
 ): [Node | null, idNodeMap] {
   const idNodeMap: idNodeMap = {};
   const node = buildNodeWithSN(n, doc, idNodeMap, false, HACK_CSS);
-  sideEffects(idNodeMap);
+  visit(idNodeMap, (visitedNode) => {
+    if (onVisit) {
+      onVisit(visitedNode);
+    }
+    handleScroll(visitedNode);
+  });
   return [node, idNodeMap];
 }