Update pathCache in NodePath#_replaceWith (#12391)

* add tests on regression 12386

* fix: update cache on _replaceWith

* fix: pathCache in replaceWithMultiple could be nullish

* Update packages/babel-traverse/src/path/replacement.js

Co-authored-by: Nicolò Ribaudo <nicolo.ribaudo@gmail.com>

* test: add replaceWith test to traverse

Co-authored-by: Brian Ng <bng412@gmail.com>
Co-authored-by: Nicolò Ribaudo <nicolo.ribaudo@gmail.com>

packages/babel-plugin-transform-typescript/test/fixtures/variable-declaration/non-null-in-optional-chain/input.ts

@@ -0,0 +1,2 @@
+const a = 1;
+const b = () => a?.b?.c!.d;

packages/babel-plugin-transform-typescript/test/fixtures/variable-declaration/non-null-in-optional-chain/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": ["transform-typescript", "proposal-optional-chaining"]
+}

packages/babel-plugin-transform-typescript/test/fixtures/variable-declaration/non-null-in-optional-chain/output.js

@@ -0,0 +1,7 @@
+const a = 1;
+
+const b = () => {
+  var _a$b;
+
+  return a === null || a === void 0 ? void 0 : (_a$b = a.b) === null || _a$b === void 0 ? void 0 : _a$b.c.d;
+};

packages/babel-traverse/src/path/replacement.js

@@ -50,7 +50,7 @@ export function replaceWithMultiple(nodes: Array<Object>) {
   nodes = this._verifyNodeList(nodes);
   t.inheritLeadingComments(nodes[0], this.node);
   t.inheritTrailingComments(nodes[nodes.length - 1], this.node);
-  pathCache.get(this.parent).delete(this.node);
+  pathCache.get(this.parent)?.delete(this.node);
   this.node = this.container[this.key] = null;
   const paths = this.insertAfter(nodes);
 
@@ -199,6 +199,7 @@ export function _replaceWith(node) {
   }
 
   this.debug(`Replace with ${node?.type}`);
+  pathCache.get(this.parent)?.set(node, this).delete(this.node);
 
   this.node = this.container[this.key] = node;
 }

packages/babel-traverse/test/conversion.js

@@ -51,6 +51,14 @@ describe("conversion", function () {
       expect(generateCode(rootPath)).toBe("() => {\n  return false;\n};");
     });
 
+    it("preserves arrow function body's context", function () {
+      const rootPath = getPath("() => true").get("expression");
+      const body = rootPath.get("body");
+      rootPath.ensureBlock();
+      body.replaceWithMultiple([t.booleanLiteral(false), t.emptyStatement()]);
+      expect(generateCode(rootPath)).toBe("() => {\n  return false;\n};");
+    });
+
     it("converts for loop with statement body to block", function () {
       const rootPath = getPath("for (;;) true;");
       rootPath.ensureBlock();

packages/babel-traverse/test/replacement.js

@@ -112,6 +112,33 @@ describe("path/replacement", function () {
       });
       expect(visitCounter).toBe(1);
     });
+
+    // https://github.com/babel/babel/issues/12386
+    it("updates pathCache with the replaced node", () => {
+      const ast = parse(`() => (a?.b)?.c`, {
+        createParenthesizedExpressions: true,
+      });
+      traverse(ast, {
+        OptionalMemberExpression(path) {
+          path.node.type = "MemberExpression";
+          // force `replaceWith` to replace `path.node`
+          path.replaceWith(path.node.__clone());
+          path.parentPath.ensureBlock();
+
+          const aQuestionDotBNode = path.node.object.expression;
+          // avoid traversing to a?.b
+          aQuestionDotBNode.type = "MemberExpression";
+        },
+        ParenthesizedExpression(path) {
+          path.replaceWith(path.node.expression);
+        },
+      });
+      expect(generate(ast).code).toMatchInlineSnapshot(`
+        "() => {
+          return a.b.c;
+        };"
+      `);
+    });
   });
   describe("replaceWithMultiple", () => {
     it("does not add extra parentheses for a JSXElement with a JSXElement parent", () => {