Address pr comments

Author: dhaval-valotia

packages/rrdom/src/index.ts

@@ -368,7 +368,7 @@ export class Mirror implements IMirror<RRNode> {
 
   // removes the node from idNodeMap
   // doesn't remove the node from nodeMetaMap
-  removeNodeFromMap(n: RRNode) {
+  removeNodeFromMap(n: RRNode, _permanent?: boolean) {
     const id = this.getId(n);
     this.idNodeMap.delete(id);
 

packages/rrweb-snapshot/src/utils.ts

@@ -203,26 +203,15 @@ export class Mirror implements IMirror<Node> {
   }
 
   // removes the node from idNodeMap
-  // doesn't remove the node from nodeMetaMap
-  removeNodeFromMap(n: Node) {
+  // if permanent is true, also removes from nodeMetaMap
+  removeNodeFromMap(n: Node, permanent = false) {
     const id = this.getId(n);
     this.idNodeMap.delete(id);
+    if (permanent) this.nodeMetaMap.delete(n);
 
     if (n.childNodes) {
       n.childNodes.forEach((childNode) =>
-        this.removeNodeFromMap(childNode as unknown as Node),
-      );
-    }
-  }
-
-  removeNodeFromMapPermanently(n: Node) {
-    const id = this.getId(n);
-    this.idNodeMap.delete(id);
-    this.nodeMetaMap.delete(n);
-
-    if (n.childNodes) {
-      n.childNodes.forEach((childNode) =>
-        this.removeNodeFromMapPermanently(childNode as unknown as Node),
+        this.removeNodeFromMap(childNode as unknown as Node, permanent),
       );
     }
   }

packages/rrweb-snapshot/test/utils.test.ts

@@ -283,7 +283,7 @@ describe('utils', () => {
   });
 
   describe('Mirror', () => {
-    describe('removeNodeFromMapPermanently', () => {
+    describe('removeNodeFromMap with permanent flag', () => {
       it('should remove node from both idNodeMap and nodeMetaMap', () => {
         const mirror = new Mirror();
         const div = document.createElement('div');
@@ -300,7 +300,7 @@ describe('utils', () => {
         expect(mirror.hasNode(div)).toBe(true);
         expect(mirror.getNode(1)).toBe(div);
 
-        mirror.removeNodeFromMapPermanently(div);
+        mirror.removeNodeFromMap(div, true);
 
         expect(mirror.getId(div)).toBe(-1);
         expect(mirror.hasNode(div)).toBe(false);
@@ -358,7 +358,7 @@ describe('utils', () => {
         expect(mirror.hasNode(child2)).toBe(true);
         expect(mirror.hasNode(grandchild)).toBe(true);
 
-        mirror.removeNodeFromMapPermanently(parent);
+        mirror.removeNodeFromMap(parent, true);
 
         expect(mirror.hasNode(parent)).toBe(false);
         expect(mirror.hasNode(child1)).toBe(false);
@@ -370,7 +370,7 @@ describe('utils', () => {
         expect(mirror.getId(grandchild)).toBe(-1);
       });
 
-      it('should differ from removeNodeFromMap by also clearing nodeMetaMap', () => {
+      it('should differ from non-permanent removeNodeFromMap by also clearing nodeMetaMap', () => {
         const mirror1 = new Mirror();
         const div1 = document.createElement('div');
         const meta1 = {
@@ -399,9 +399,9 @@ describe('utils', () => {
         } as serializedNodeWithId;
 
         mirror2.add(div2, meta2);
-        mirror2.removeNodeFromMapPermanently(div2);
+        mirror2.removeNodeFromMap(div2, true);
 
-        // removeNodeFromMapPermanently clears both maps
+        // removeNodeFromMap with permanent=true clears both maps
         expect(mirror2.getNode(1)).toBeNull();
         expect(mirror2.hasNode(div2)).toBe(false); // Also removed from nodeMetaMap
       });

packages/rrweb/src/record/iframe-manager.ts

@@ -78,6 +78,16 @@ export class IframeManager {
   }
 
   public removeIframe(iframeEl: HTMLIFrameElement): void {
+    const storedDoc = this.iframeContentDocumentMap.get(iframeEl);
+
+    this.stylesheetManager.cleanupStylesheetsForRemovedNode(
+      iframeEl,
+      storedDoc,
+    );
+    if (storedDoc) {
+      this.mirror.removeNodeFromMap(storedDoc, true);
+    }
+
     this.iframes.delete(iframeEl);
     this.iframeContentDocumentMap.delete(iframeEl);
 

packages/rrweb/src/record/mutation.ts

@@ -370,18 +370,8 @@ export default class MutationBuffer {
       const removedNode = this.mapRemoves.shift()!;
 
       if (removedNode.nodeName === 'IFRAME') {
-        const iframeEl = removedNode as HTMLIFrameElement;
         try {
-          const storedDoc =
-            this.iframeManager.getIframeContentDocument(iframeEl);
-          this.stylesheetManager.cleanupStylesheetsForRemovedNode(
-            removedNode,
-            storedDoc,
-          );
-          if (storedDoc) {
-            this.mirror.removeNodeFromMapPermanently(storedDoc);
-          }
-          this.iframeManager.removeIframe(iframeEl);
+          this.iframeManager.removeIframe(removedNode as HTMLIFrameElement);
         } catch (e) {
           // Ignore errors during iframe cleanup
         }

packages/types/src/index.ts

@@ -828,9 +828,7 @@ export interface IMirror<TNode> {
 
   getMeta(n: TNode): serializedNodeWithId | null;
 
-  removeNodeFromMap(n: TNode): void;
-
-  removeNodeFromMapPermanently?(n: TNode): void;
+  removeNodeFromMap(n: TNode, permanent?: boolean): void;
 
   has(id: number): boolean;