Only extract IDs for TDZ checks in assign when necessary

packages/babel-plugin-transform-block-scoping/src/index.ts

@@ -1,6 +1,6 @@
 import { declare } from "@babel/helper-plugin-utils";
 import type { NodePath, Visitor, Scope, Binding } from "@babel/traverse";
-import { visitor as tdzVisitor } from "./tdz";
+import { skipTDZChecks, visitor as tdzVisitor } from "./tdz";
 import type { TDZVisitorState } from "./tdz";
 import { traverse, template, types as t } from "@babel/core";
 import type { PluginPass } from "@babel/core";
@@ -43,8 +43,7 @@ export default declare((api, opts: Options) => {
               t.cloneNode(decl.id),
               decl.init || scope.buildUndefinedNode(),
             );
-            // @ts-expect-error todo(flow->ts): avoid mutations
-            assign._ignoreBlockScopingTDZ = true;
+            skipTDZChecks.add(assign);
             nodes.push(t.expressionStatement(assign));
             decl.init = this.addHelper("temporalUndefined");
           }

packages/babel-plugin-transform-block-scoping/src/tdz.ts

@@ -1,10 +1,7 @@
-import { types as t, template, type PluginPass } from "@babel/core";
+import { types as t, type PluginPass } from "@babel/core";
 import type { NodePath, Scope, Visitor } from "@babel/traverse";
 
-function getTDZStatus(
-  refPath: NodePath<t.Identifier | t.JSXIdentifier>,
-  bindingPath: NodePath,
-) {
+function getTDZStatus(refPath: NodePath, bindingPath: NodePath) {
   const executionStatus = bindingPath._guessExecutionStatusRelativeTo(refPath);
 
   if (executionStatus === "before") {
@@ -16,15 +13,26 @@ function getTDZStatus(
   }
 }
 
+export const skipTDZChecks = new WeakSet();
+
 function buildTDZAssert(
+  status: "maybe" | "inside",
   node: t.Identifier | t.JSXIdentifier,
   state: TDZVisitorState,
 ) {
-  return t.callExpression(state.addHelper("temporalRef"), [
-    // @ts-expect-error Fixme: we may need to handle JSXIdentifier
-    node,
-    t.stringLiteral(node.name),
-  ]);
+  if (status === "maybe") {
+    const clone = t.cloneNode(node);
+    skipTDZChecks.add(clone);
+    return t.callExpression(state.addHelper("temporalRef"), [
+      // @ts-expect-error Fixme: we may need to handle JSXIdentifier
+      clone,
+      t.stringLiteral(node.name),
+    ]);
+  } else {
+    return t.callExpression(state.addHelper("tdz"), [
+      t.stringLiteral(node.name),
+    ]);
+  }
 }
 
 function isReference(
@@ -39,7 +47,41 @@ function isReference(
   return scope.getBindingIdentifier(node.name) === declared;
 }
 
-const visitedMaybeTDZNodes = new WeakSet();
+type TDZReplacement = { status: "maybe" | "inside"; node: t.Expression };
+function getTDZReplacement(
+  path: NodePath<t.Identifier | t.JSXIdentifier>,
+  state: TDZVisitorState,
+): TDZReplacement | undefined;
+function getTDZReplacement(
+  path: NodePath,
+  state: TDZVisitorState,
+  id: t.Identifier | t.JSXIdentifier,
+): TDZReplacement | undefined;
+function getTDZReplacement(
+  path: NodePath,
+  state: TDZVisitorState,
+  id: t.Identifier | t.JSXIdentifier = path.node as any,
+): TDZReplacement | undefined {
+  if (!isReference(id, path.scope, state)) return;
+
+  if (skipTDZChecks.has(id)) return;
+  skipTDZChecks.add(id);
+
+  const bindingPath = path.scope.getBinding(id.name).path;
+
+  if (bindingPath.isFunctionDeclaration()) return;
+
+  const status = getTDZStatus(path, bindingPath);
+  if (status === "outside") return;
+
+  if (status === "maybe") {
+    // add tdzThis to parent variable declarator so it's exploded
+    // @ts-expect-error todo(flow->ts): avoid mutations
+    bindingPath.parent._tdzThis = true;
+  }
+
+  return { status, node: buildTDZAssert(status, id, state) };
+}
 
 export interface TDZVisitorState {
   tdzEnabled: boolean;
@@ -50,72 +92,54 @@ export interface TDZVisitorState {
 export const visitor: Visitor<TDZVisitorState> = {
   ReferencedIdentifier(path, state) {
     if (!state.tdzEnabled) return;
+    if (path.parentPath.isUpdateExpression()) return;
+    // It will be handled after transforming the loop
+    if (path.parentPath.isFor({ left: path.node })) return;
 
-    const { node, parent, scope } = path;
+    const replacement = getTDZReplacement(path, state);
+    if (!replacement) return;
 
-    if (path.parentPath.isFor({ left: node })) return;
-    if (!isReference(node, scope, state)) return;
+    path.replaceWith(replacement.node);
+  },
 
-    const bindingPath = scope.getBinding(node.name).path;
+  UpdateExpression(path, state) {
+    if (!state.tdzEnabled) return;
 
-    if (bindingPath.isFunctionDeclaration()) return;
+    const { node } = path;
+    if (skipTDZChecks.has(node)) return;
+    skipTDZChecks.add(node);
 
-    const status = getTDZStatus(path, bindingPath);
-    if (status === "outside") return;
+    const arg = path.get("argument");
+    if (!arg.isIdentifier()) return;
 
-    if (status === "maybe") {
-      if (visitedMaybeTDZNodes.has(node)) {
-        return;
-      }
-      visitedMaybeTDZNodes.add(node);
-      const assert = buildTDZAssert(node, state);
-
-      // add tdzThis to parent variable declarator so it's exploded
-      // @ts-expect-error todo(flow->ts): avoid mutations
-      bindingPath.parent._tdzThis = true;
-
-      if (path.parentPath.isUpdateExpression()) {
-        // @ts-expect-error todo(flow->ts): avoid node mutations
-        if (parent._ignoreBlockScopingTDZ) return;
-        path.parentPath.replaceWith(
-          t.sequenceExpression([assert, parent as t.UpdateExpression]),
-        );
-      } else {
-        path.replaceWith(assert);
-      }
-    } else if (status === "inside") {
-      path.replaceWith(
-        template.ast`${state.addHelper("tdz")}("${node.name}")` as t.Statement,
-      );
+    const replacement = getTDZReplacement(path, state, arg.node);
+    if (!replacement) return;
+
+    if (replacement.status === "maybe") {
+      path.insertBefore(replacement.node);
+    } else {
+      path.replaceWith(replacement.node);
     }
   },
 
-  AssignmentExpression: {
-    exit(path, state) {
-      if (!state.tdzEnabled) return;
-
-      const { node } = path;
-
-      // @ts-expect-error todo(flow->ts): avoid node mutations
-      if (node._ignoreBlockScopingTDZ) return;
+  AssignmentExpression(path, state) {
+    if (!state.tdzEnabled) return;
 
-      const nodes = [];
-      const ids = path.getBindingIdentifiers();
+    const { node } = path;
+    if (skipTDZChecks.has(node)) return;
+    skipTDZChecks.add(node);
 
-      for (const name of Object.keys(ids)) {
-        const id = ids[name];
+    const nodes = [];
+    const ids = path.getBindingIdentifiers();
 
-        if (isReference(id, path.scope, state)) {
-          nodes.push(id);
-        }
+    for (const name of Object.keys(ids)) {
+      const replacement = getTDZReplacement(path, state, ids[name]);
+      if (replacement) {
+        nodes.push(t.expressionStatement(replacement.node));
+        if (replacement.status === "inside") break;
       }
+    }
 
-      if (nodes.length) {
-        // @ts-expect-error todo(flow->ts): avoid mutations
-        node._ignoreBlockScopingTDZ = true;
-        nodes.push(node);
-        path.replaceWithMultiple(nodes.map(n => t.expressionStatement(n)));
-      }
-    },
+    if (nodes.length > 0) path.insertBefore(nodes);
   },
 };

packages/babel-plugin-transform-block-scoping/test/fixtures/tdz/simple-assign-no-tdz/exec.js

@@ -0,0 +1,2 @@
+let i
+i = 2;

packages/babel-plugin-transform-block-scoping/test/fixtures/tdz/simple-assign-no-tdz/input.js

@@ -0,0 +1,2 @@
+let i
+i = 2;

packages/babel-plugin-transform-block-scoping/test/fixtures/tdz/simple-assign-no-tdz/output.js

@@ -0,0 +1,2 @@
+var i;
+i = 2;

packages/babel-plugin-transform-block-scoping/test/fixtures/tdz/update-expression/input.js

@@ -0,0 +1,13 @@
+maybeCallLater(function f() {
+  x++;
+  ++x;
+  x.p++;
+  ++x.p;
+});
+
+x++;
+++x;
+x.p++;
+++x.p;
+
+let x;

packages/babel-plugin-transform-block-scoping/test/fixtures/tdz/update-expression/output.js

@@ -0,0 +1,15 @@
+var x = babelHelpers.temporalUndefined;
+maybeCallLater(function f() {
+  babelHelpers.temporalRef(x, "x")
+  x++;
+  babelHelpers.temporalRef(x, "x")
+  ++x;
+  babelHelpers.temporalRef(x, "x").p++;
+  ++babelHelpers.temporalRef(x, "x").p;
+});
+babelHelpers.tdz("x");
+babelHelpers.tdz("x");
+babelHelpers.tdz("x").p++;
+++babelHelpers.tdz("x").p;
+x = void 0;
+void 0;