JavaScript - fixing parsing of whitespace before function expressions

rewrite-javascript/rewrite/src/javascript/format.ts

@@ -15,7 +15,7 @@
  */
 import {JS} from "./tree";
 import {JavaScriptVisitor} from "./visitor";
-import {Comment, emptySpace, J, Statement} from "../java";
+import {Comment, J, Statement} from "../java";
 import {Draft, produce} from "immer";
 import {Cursor, isScope, Tree} from "../tree";
 import {
@@ -445,6 +445,10 @@ export class SpacesVisitor<P> extends JavaScriptVisitor<P> {
     }
     protected async visitVariable(variable: J.VariableDeclarations.NamedVariable, p: P): Promise<J | undefined> {
         const ret = await super.visitVariable(variable, p) as J.VariableDeclarations.NamedVariable;
+        if (variable.initializer?.element?.kind == JS.Kind.StatementExpression
+            && (variable.initializer.element as JS.StatementExpression).statement.kind == J.Kind.MethodDeclaration) {
+            return ret;
+        }
         return produceAsync(ret, async draft => {
             if (draft.initializer) {
                 draft.initializer.before.whitespace = this.style.aroundOperators.assignment ? " " : "";
@@ -968,7 +972,11 @@ export class BlankLinesVisitor<P> extends JavaScriptVisitor<P> {
             });
             return super.visit(cu, p, cursor);
         }
-        if (tree.kind === J.Kind.MethodDeclaration) {
+        if (tree.kind === JS.Kind.StatementExpression && (tree as JS.StatementExpression).statement.kind == J.Kind.MethodDeclaration) {
+            tree = produce(tree as JS.StatementExpression, draft => {
+                this.ensurePrefixHasNewLine(draft);
+            });
+        } else if (tree.kind === J.Kind.MethodDeclaration && this.cursor.value.kind != JS.Kind.StatementExpression) {
             tree = produce(tree as J.MethodDeclaration, draft => {
                 this.ensurePrefixHasNewLine(draft);
             });
@@ -1149,7 +1157,7 @@ export class TabsAndIndentsVisitor<P> extends JavaScriptVisitor<P> {
         let indentShouldIncrease =
             tree.kind === J.Kind.Block
             || this.cursor.parent?.parent?.parent?.value.kind == J.Kind.Case
-            || (tree.kind === J.Kind.MethodDeclaration && this.cursor.parent?.value.kind === JS.Kind.StatementExpression);
+            || (tree.kind === JS.Kind.StatementExpression && (tree as JS.StatementExpression).statement.kind == J.Kind.MethodDeclaration);
 
         const previousIndent = this.currentIndent;
 

rewrite-javascript/rewrite/src/javascript/parser.ts

@@ -3131,12 +3131,15 @@ export class JavaScriptParserVisitor {
     }
 
     visitFunctionExpression(node: ts.FunctionExpression): JS.StatementExpression {
+        const delegate = this.mapFunctionDeclaration(node);
         return {
             kind: JS.Kind.StatementExpression,
             id: randomId(),
-            prefix: emptySpace,
+            prefix: delegate.prefix,
             markers: emptyMarkers,
-            statement: this.mapFunctionDeclaration(node)
+            statement: produce(delegate, draft => {
+              draft.prefix = emptySpace;
+            })
         };
     }
 

rewrite-javascript/rewrite/test/javascript/format/format.test.ts

@@ -192,9 +192,8 @@ describe('AutoformatVisitor', () => {
             //language=typescript
             typescript(
                 `const fn = function () {return 99;};`,
-                 // TODO the space after `const fn=` is excessive
                  `
-                const fn = 
+                const fn =
                     function () {
                         return 99;
                     };

rewrite-javascript/rewrite/test/javascript/parser/variable-declarations.test.ts

@@ -252,4 +252,23 @@ describe('variable declaration mapping', () => {
                 }
             `)
         ));
+
+    test.each([
+        "const c =  function(): number { return 116; };",
+        "const c =  136;",
+        "const c =  (1 > 0) ? 116 : 119;"
+    ])('double space: %s', (code) =>
+        spec.rewriteRun({
+            //language=javascript
+            ...typescript(code),
+            afterRecipe: (cu: JS.CompilationUnit) => {
+                expect(cu.statements).toHaveLength(1);
+                cu.statements.forEach(statement => {
+                    const varDecl = statement.element as J.VariableDeclarations;
+                    const initializer = varDecl.variables[0].element.initializer!;
+                    expect(initializer.before.whitespace).toBe(" ");
+                    expect(initializer.element.prefix.whitespace).toBe("  ");
+                });
+            }
+        }));
 });