Do not create multiple copies of comments when cloning nodes

packages/babel-types/src/clone/cloneNode.ts

@@ -5,19 +5,19 @@ import { isFile, isIdentifier } from "../validators/generated";
 const has = Function.call.bind(Object.prototype.hasOwnProperty);
 
 // This function will never be called for comments, only for real nodes.
-function cloneIfNode(obj, deep, withoutLoc) {
+function cloneIfNode(obj, deep, withoutLoc, commentsCache) {
   if (obj && typeof obj.type === "string") {
-    return cloneNode(obj, deep, withoutLoc);
+    return cloneNodeInternal(obj, deep, withoutLoc, commentsCache);
   }
 
   return obj;
 }
 
-function cloneIfNodeOrArray(obj, deep, withoutLoc) {
+function cloneIfNodeOrArray(obj, deep, withoutLoc, commentsCache) {
   if (Array.isArray(obj)) {
-    return obj.map(node => cloneIfNode(node, deep, withoutLoc));
+    return obj.map(node => cloneIfNode(node, deep, withoutLoc, commentsCache));
   }
-  return cloneIfNode(obj, deep, withoutLoc);
+  return cloneIfNode(obj, deep, withoutLoc, commentsCache);
 }
 
 /**
@@ -29,6 +29,15 @@ export default function cloneNode<T extends t.Node>(
   node: T,
   deep: boolean = true,
   withoutLoc: boolean = false,
+): T {
+  return cloneNodeInternal(node, deep, withoutLoc, new Map());
+}
+
+function cloneNodeInternal<T extends t.Node>(
+  node: T,
+  deep: boolean = true,
+  withoutLoc: boolean = false,
+  commentsCache: Map<t.Comment, t.Comment>,
 ): T {
   if (!node) return node;
 
@@ -45,7 +54,12 @@ export default function cloneNode<T extends t.Node>(
 
     if (has(node, "typeAnnotation")) {
       newNode.typeAnnotation = deep
-        ? cloneIfNodeOrArray(node.typeAnnotation, true, withoutLoc)
+        ? cloneIfNodeOrArray(
+            node.typeAnnotation,
+            true,
+            withoutLoc,
+            commentsCache,
+          )
         : node.typeAnnotation;
     }
   } else if (!has(NODE_FIELDS, type)) {
@@ -56,8 +70,18 @@ export default function cloneNode<T extends t.Node>(
         if (deep) {
           newNode[field] =
             isFile(node) && field === "comments"
-              ? maybeCloneComments(node.comments, deep, withoutLoc)
-              : cloneIfNodeOrArray(node[field], true, withoutLoc);
+              ? maybeCloneComments(
+                  node.comments,
+                  deep,
+                  withoutLoc,
+                  commentsCache,
+                )
+              : cloneIfNodeOrArray(
+                  node[field],
+                  true,
+                  withoutLoc,
+                  commentsCache,
+                );
         } else {
           newNode[field] = node[field];
         }
@@ -77,20 +101,23 @@ export default function cloneNode<T extends t.Node>(
       node.leadingComments,
       deep,
       withoutLoc,
+      commentsCache,
     );
   }
   if (has(node, "innerComments")) {
     newNode.innerComments = maybeCloneComments(
       node.innerComments,
       deep,
       withoutLoc,
+      commentsCache,
     );
   }
   if (has(node, "trailingComments")) {
     newNode.trailingComments = maybeCloneComments(
       node.trailingComments,
       deep,
       withoutLoc,
+      commentsCache,
     );
   }
   if (has(node, "extra")) {
@@ -106,14 +133,24 @@ function maybeCloneComments<T extends t.Comment>(
   comments: ReadonlyArray<T> | null,
   deep: boolean,
   withoutLoc: boolean,
+  commentsCache: Map<T, T>,
 ): ReadonlyArray<T> | null {
   if (!comments || !deep) {
     return comments;
   }
-  return comments.map(({ type, value, loc }) => {
+  return comments.map(comment => {
+    const cache = commentsCache.get(comment);
+    if (cache) return cache;
+
+    const { type, value, loc } = comment;
+
+    const ret = { type, value, loc } as T;
     if (withoutLoc) {
-      return { type, value, loc: null } as T;
+      ret.loc = null;
     }
-    return { type, value, loc } as T;
+
+    commentsCache.set(comment, ret);
+
+    return ret;
   });
 }

packages/babel-types/test/cloning.js

@@ -1,5 +1,6 @@
 import * as t from "../lib/index.js";
 import { parse } from "@babel/parser";
+import { CodeGenerator } from "@babel/generator";
 
 describe("cloneNode", function () {
   it("should handle undefined", function () {
@@ -151,4 +152,24 @@ describe("cloneNode", function () {
     expect(cloned.declarations[0].id.innerComments[0].loc).toBe(null);
     expect(cloned.declarations[0].id.trailingComments[0].loc).toBe(null);
   });
+
+  it("should generate same code after deep cloning", function () {
+    let code = `//test1
+    /*test2*/var/*test3*/ a = 1/*test4*/;//test5
+    //test6
+    var b;
+    `;
+    code = new CodeGenerator(parse(code), { retainLines: true }).generate()
+      .code;
+
+    const ast = t.cloneNode(
+      parse(code),
+      /* deep */ true,
+      /* withoutLoc */ false,
+    );
+    const newCode = new CodeGenerator(ast, { retainLines: true }).generate()
+      .code;
+
+    expect(newCode).toBe(code);
+  });
 });