Correctly print parens around in in for heads (babel#16630)

packages/babel-generator/src/generators/base.ts

@@ -46,7 +46,9 @@ export function BlockStatement(this: Printer, node: t.BlockStatement) {
     }
   }
 
+  const exit = this.enterForStatementInit(false);
   this.printSequence(node.body, node, { indent: true });
+  exit();
 
   this.rightBrace(node);
 }

packages/babel-generator/src/generators/classes.ts

@@ -75,7 +75,9 @@ export function ClassBody(this: Printer, node: t.ClassBody) {
   } else {
     this.newline();
 
+    const exit = this.enterForStatementInit(false);
     this.printSequence(node.body, node, { indent: true });
+    exit();
 
     if (!this.endsWith(charCodes.lineFeed)) this.newline();
 

packages/babel-generator/src/generators/expressions.ts

@@ -6,7 +6,6 @@ import {
   isNewExpression,
 } from "@babel/types";
 import type * as t from "@babel/types";
-import * as n from "../node/index.ts";
 
 export function UnaryExpression(this: Printer, node: t.UnaryExpression) {
   const { operator } = node;
@@ -97,7 +96,9 @@ export function NewExpression(
     this.token("?.");
   }
   this.token("(");
+  const exit = this.enterForStatementInit(false);
   this.printList(node.arguments, node);
+  exit();
   this.rightParens(node);
 }
 
@@ -179,7 +180,9 @@ export function OptionalCallExpression(
   this.print(node.typeArguments, node); // Flow
 
   this.token("(");
+  const exit = this.enterForStatementInit(false);
   this.printList(node.arguments, node);
+  exit();
   this.rightParens(node);
 }
 
@@ -189,7 +192,9 @@ export function CallExpression(this: Printer, node: t.CallExpression) {
   this.print(node.typeArguments, node); // Flow
   this.print(node.typeParameters, node); // TS
   this.token("(");
+  const exit = this.enterForStatementInit(false);
   this.printList(node.arguments, node);
+  exit();
   this.rightParens(node);
 }
 
@@ -251,19 +256,7 @@ export function AssignmentPattern(this: Printer, node: t.AssignmentPattern) {
 export function AssignmentExpression(
   this: Printer,
   node: t.AssignmentExpression,
-  parent: t.Node,
 ) {
-  // Somewhere inside a for statement `init` node but doesn't usually
-  // needs a paren except for `in` expressions: `for (a in b ? a : b;;)`
-  const parens =
-    this.inForStatementInitCounter &&
-    node.operator === "in" &&
-    !n.needsParens(node, parent);
-
-  if (parens) {
-    this.token("(");
-  }
-
   this.print(node.left, node);
 
   this.space();
@@ -275,10 +268,6 @@ export function AssignmentExpression(
   this.space();
 
   this.print(node.right, node);
-
-  if (parens) {
-    this.token(")");
-  }
 }
 
 export function BindExpression(this: Printer, node: t.BindExpression) {
@@ -306,9 +295,11 @@ export function MemberExpression(this: Printer, node: t.MemberExpression) {
   }
 
   if (computed) {
+    const exit = this.enterForStatementInit(false);
     this.token("[");
     this.print(node.property, node);
     this.token("]");
+    exit();
   } else {
     this.token(".");
     this.print(node.property, node);

packages/babel-generator/src/generators/methods.ts

@@ -38,6 +38,8 @@ export function _parameters(
     | t.TSFunctionType
     | t.TSConstructorType,
 ) {
+  const exit = this.enterForStatementInit(false);
+
   const paramLength = parameters.length;
   for (let i = 0; i < paramLength; i++) {
     this._param(parameters[i], parent);
@@ -47,6 +49,8 @@ export function _parameters(
       this.space();
     }
   }
+
+  exit();
 }
 
 export function _param(

packages/babel-generator/src/generators/statements.ts

@@ -68,9 +68,12 @@ export function ForStatement(this: Printer, node: t.ForStatement) {
   this.space();
   this.token("(");
 
-  this.inForStatementInitCounter++;
-  this.print(node.init, node);
-  this.inForStatementInitCounter--;
+  {
+    const exit = this.enterForStatementInit(true);
+    this.print(node.init, node);
+    exit();
+  }
+
   this.token(";");
 
   if (node.test) {
@@ -107,7 +110,11 @@ function ForXStatement(this: Printer, node: t.ForXStatement) {
   }
   this.noIndentInnerCommentsHere();
   this.token("(");
-  this.print(node.left, node);
+  {
+    const exit = isForOf ? null : this.enterForStatementInit(true);
+    this.print(node.left, node);
+    exit?.();
+  }
   this.space();
   this.word(isForOf ? "of" : "in");
   this.space();

packages/babel-generator/src/generators/types.ts

@@ -25,9 +25,11 @@ export function ObjectExpression(this: Printer, node: t.ObjectExpression) {
   this.token("{");
 
   if (props.length) {
+    const exit = this.enterForStatementInit(false);
     this.space();
     this.printList(props, node, { indent: true, statement: true });
     this.space();
+    exit();
   }
 
   this.sourceWithOffset("end", node.loc, -1);
@@ -87,6 +89,8 @@ export function ArrayExpression(this: Printer, node: t.ArrayExpression) {
 
   this.token("[");
 
+  const exit = this.enterForStatementInit(false);
+
   for (let i = 0; i < elems.length; i++) {
     const elem = elems[i];
     if (elem) {
@@ -103,6 +107,8 @@ export function ArrayExpression(this: Printer, node: t.ArrayExpression) {
     }
   }
 
+  exit();
+
   this.token("]");
 }
 

packages/babel-generator/src/node/index.ts

@@ -21,6 +21,7 @@ type NodeHandler<R> = (
   //   : t.Node,
   parent: t.Node,
   stack?: t.Node[],
+  inForStatementInit?: boolean,
 ) => R;
 
 export type NodeHandlers<R> = {
@@ -35,8 +36,11 @@ function expandAliases<R>(obj: NodeHandlers<R>) {
     map.set(
       type,
       fn
-        ? function (node, parent, stack) {
-            return fn(node, parent, stack) ?? func(node, parent, stack);
+        ? function (node, parent, stack, inForInit) {
+            return (
+              fn(node, parent, stack, inForInit) ??
+              func(node, parent, stack, inForInit)
+            );
           }
         : func,
     );
@@ -101,6 +105,7 @@ export function needsParens(
   node: t.Node,
   parent: t.Node,
   printStack?: t.Node[],
+  inForInit?: boolean,
 ) {
   if (!parent) return false;
 
@@ -116,7 +121,7 @@ export function needsParens(
     );
   }
 
-  return expandedParens.get(node.type)?.(node, parent, printStack);
+  return expandedParens.get(node.type)?.(node, parent, printStack, inForInit);
 }
 
 function isDecoratorMemberExpression(node: t.Node): boolean {

packages/babel-generator/src/node/parentheses.ts

@@ -252,19 +252,12 @@ export function TSInstantiationExpression(
 export function BinaryExpression(
   node: t.BinaryExpression,
   parent: t.Node,
+  stack: unknown,
+  inForStatementInit: boolean,
 ): boolean {
-  // let i = (1 in []);
   // for ((1 in []);;);
-  if (node.operator === "in") {
-    const parentType = parent.type;
-    return (
-      parentType === "VariableDeclarator" ||
-      parentType === "ForStatement" ||
-      parentType === "ForInStatement" ||
-      parentType === "ForOfStatement"
-    );
-  }
-  return false;
+  // for (var x = (1 in []) in 2);
+  return node.operator === "in" && inForStatementInit;
 }
 
 export function SequenceExpression(

packages/babel-generator/src/printer.ts

@@ -119,7 +119,16 @@ class Printer {
   declare _inputMap: TraceMap;
 
   declare format: Format;
-  inForStatementInitCounter: number = 0;
+
+  inForStatementInit: boolean = false;
+  enterForStatementInit(val: boolean) {
+    const old = this.inForStatementInit;
+    if (old === val) return () => {};
+    this.inForStatementInit = val;
+    return () => {
+      this.inForStatementInit = old;
+    };
+  }
 
   declare _buf: Buffer;
   _printStack: Array<t.Node> = [];
@@ -658,7 +667,7 @@ class Printer {
       (parenthesized &&
         format.retainFunctionParens &&
         nodeType === "FunctionExpression") ||
-      needsParens(node, parent, this._printStack);
+      needsParens(node, parent, this._printStack, this.inForStatementInit);
 
     if (
       !shouldPrintParens &&
@@ -683,9 +692,11 @@ class Printer {
       }
     }
 
+    let exitInForStatementInit;
     if (shouldPrintParens) {
       this.token("(");
       this._endsWithInnerRaw = false;
+      exitInForStatementInit = this.enterForStatementInit(false);
     }
 
     this._lastCommentLine = 0;
@@ -707,6 +718,7 @@ class Printer {
       this._printTrailingComments(node, parent);
       this.token(")");
       this._noLineTerminator = noLineTerminatorAfter;
+      exitInForStatementInit();
     } else if (noLineTerminatorAfter && !this._noLineTerminator) {
       this._noLineTerminator = true;
       this._printTrailingComments(node, parent);

packages/babel-generator/test/fixtures/parentheses/in-inside-for/input.js

@@ -0,0 +1,47 @@
+for (var a = (b in c) in {});
+for (var a = 1 || (b in c) in {});
+for (var a = 1 + (2 || (b in c)) in {});
+for (var a = (() => b in c) in {});
+for (var a = 1 || (() => b in c) in {});
+for (var a = (() => { b in c; }) in {});
+for (var a = [b in c] in {});
+for (var a = {b: b in c} in {});
+for (var a = (x = b in c) => {} in {});
+for (var a = class extends (b in c) {} in {});
+for (var a = function (x = b in c) {} in {});
+
+for (var a = (b in c);;);
+for (var a = 1 || (b in c);;);
+for (var a = 1 + (2 || (b in c));;);
+for (var a = (() => b in c);;);
+for (var a = 1 || (() => b in c);;);
+for (var a = (() => { b in c; });;);
+for (var a = [b in c];;);
+for (var a = {b: b in c};;);
+for (var a = (x = b in c) => {};;);
+for (var a = class extends (b in c) {};;);
+for (var a = function (x = b in c) {};;);
+
+for (var a in (b in c));
+for (var a in 1 || (b in c));
+for (var a in 1 + (2 || (b in c)));
+for (var a in (() => b in c));
+for (var a in 1 || (() => b in c));
+for (var a in (() => { b in c; }));
+for (var a in [b in c]);
+for (var a in {b: b in c});
+for (var a in (x = b in c) => {});
+for (var a in class extends (b in c) {});
+for (var a in function (x = b in c) {});
+
+for (;a = (b in c););
+for (;a = 1 || (b in c););
+for (;a = 1 + (2 || (b in c)););
+for (;a = (() => b in c););
+for (;a = 1 || (() => b in c););
+for (;a = (() => { b in c; }););
+for (;a = [b in c];);
+for (;a = {b: b in c};);
+for (;a = (x = b in c) => {};);
+for (;a = class extends (b in c) {};);
+for (;a = function (x = b in c) {};);

packages/babel-generator/test/fixtures/parentheses/in-inside-for/options.json

@@ -0,0 +1,4 @@
+{
+  "sourceType": "script",
+  "strictMode": false
+}