From cfc8798b53c4293324e9692a40a2bdd7811465f9 Mon Sep 17 00:00:00 2001
From: Yanzhen Yu <yanzhen@smartx.com>
Date: Wed, 1 Apr 2026 12:00:00 +0800
Subject: [PATCH] return id node map when snapshot

---
 index.d.ts          |  4 ++--
 src/snapshot.ts     | 29 ++++++++++++++++++++++++-----
 src/types.ts        |  8 ++++++++
 test/integration.ts |  6 +++---
 4 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/index.d.ts b/index.d.ts
index bb15b0be..fa9eae71 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -1,5 +1,5 @@
-import { serializedNodeWithId } from './src/types';
+import { serializedNodeWithId, idNodeMap } from './src/types';
 export * from './src/types';
 
-export function snapshot(n: Document): serializedNodeWithId | null;
+export function snapshot(n: Document): [serializedNodeWithId | null, idNodeMap];
 export function rebuild(n: serializedNodeWithId): Node | null;
diff --git a/src/snapshot.ts b/src/snapshot.ts
index 913210a4..15642b7f 100644
--- a/src/snapshot.ts
+++ b/src/snapshot.ts
@@ -3,6 +3,8 @@ import {
   serializedNodeWithId,
   NodeType,
   attributes,
+  INode,
+  idNodeMap,
 } from './types';
 
 let _id = 1;
@@ -115,22 +117,38 @@ function serializeNode(n: Node, doc: Document): serializedNode | false {
   }
 }
 
-function _snapshot(n: Node, doc: Document): serializedNodeWithId | null {
+function serializeNodeWithId(
+  n: Node,
+  doc: Document,
+): serializedNodeWithId | null {
   const _serializedNode = serializeNode(n, doc);
   if (!_serializedNode) {
     // TODO: dev only
     console.warn(n, 'not serialized');
     return null;
   }
-  const serializedNode: serializedNodeWithId = Object.assign(_serializedNode, {
+  return Object.assign(_serializedNode, {
     id: genId(),
   });
+}
+
+function _snapshot(
+  n: Node,
+  doc: Document,
+  map: idNodeMap,
+): serializedNodeWithId | null {
+  const serializedNode = serializeNodeWithId(n, doc);
+  if (!serializedNode) {
+    return null;
+  }
+  (n as INode).__sn = serializedNode;
+  map[serializedNode.id] = n as INode;
   if (
     serializedNode.type === NodeType.Document ||
     serializedNode.type === NodeType.Element
   ) {
     for (const childN of Array.from(n.childNodes)) {
-      const serializedChildNode = _snapshot(childN, doc);
+      const serializedChildNode = _snapshot(childN, doc, map);
       if (serializedChildNode) {
         serializedNode.childNodes.push(serializedChildNode);
       }
@@ -139,9 +157,10 @@ function _snapshot(n: Node, doc: Document): serializedNodeWithId | null {
   return serializedNode;
 }
 
-function snapshot(n: Document): serializedNodeWithId | null {
+function snapshot(n: Document): [serializedNodeWithId | null, idNodeMap] {
   resetId();
-  return _snapshot(n, n);
+  const idNodeMap: idNodeMap = {};
+  return [_snapshot(n, n, idNodeMap), idNodeMap];
 }
 
 export default snapshot;
diff --git a/src/types.ts b/src/types.ts
index ea19a198..6fbe58e2 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -57,3 +57,11 @@ export type serializedNodeWithId = serializedNode & { id: number };
 export type tagMap = {
   [key: string]: string;
 };
+
+export interface INode extends Node {
+  __sn: serializedNodeWithId;
+}
+
+export type idNodeMap = {
+  [key: number]: INode;
+};
diff --git a/test/integration.ts b/test/integration.ts
index 40f9cc03..a6ed7da5 100644
--- a/test/integration.ts
+++ b/test/integration.ts
@@ -38,9 +38,9 @@ const server = () =>
       '.css': 'text/css',
     };
     const s = http.createServer((req, res) => {
-      const parsedUrl = url.parse(req.url);
+      const parsedUrl = url.parse(req.url!);
       const sanitizePath = path
-        .normalize(parsedUrl.pathname)
+        .normalize(parsedUrl.pathname!)
         .replace(/^(\.\.[\/\\])+/, '');
       let pathname = path.join(__dirname, sanitizePath);
       try {
@@ -98,7 +98,7 @@ describe('integration tests', () => {
       });
       const rebuildHtml = (await page.evaluate(`${this.code}
         const x = new XMLSerializer();
-        const snap = rrweb.snapshot(document);
+        const [snap] = rrweb.snapshot(document);
         x.serializeToString(rrweb.rebuild(snap));
       `)).replace(/\n\n/g, '');
       await page.goto(`http://localhost:3030/html`);