From bb901135958ac3683f72b65ce0e3157d6775f9aa Mon Sep 17 00:00:00 2001
From: Alan Pierce <alan@benchling.com>
Date: Sun, 6 Jan 2019 09:18:15 -0800
Subject: [PATCH] Move close-paren to before arrow when removing arrow return
 types (#393)

Fixes #391

It's a bit ugly, but we now do a lookahead whenever we see a `):` with the `:`
as a type token. If the next non-type is an arrow, we move the `)` to just
before the arrow. This seems to work and seems to agree with
`processBalancedCode`.

I also added a new type of test that actually executes the resulting code and
asserts the output value from the code, which will probably be useful in the
future as well.
---
 src/transformers/FlowTransformer.ts       |  5 ++++-
 src/transformers/RootTransformer.ts       | 26 +++++++++++++++++++++++
 src/transformers/TypeScriptTransformer.ts |  6 ++++--
 test/types-test.ts                        | 21 +++++++++++++++++-
 test/util.ts                              | 25 +++++++++++++++++++++-
 5 files changed, 78 insertions(+), 5 deletions(-)

diff --git a/src/transformers/FlowTransformer.ts b/src/transformers/FlowTransformer.ts
index aa449c77..36eaa32f 100644
--- a/src/transformers/FlowTransformer.ts
+++ b/src/transformers/FlowTransformer.ts
@@ -8,6 +8,9 @@ export default class FlowTransformer extends Transformer {
   }
 
   process(): boolean {
-    return this.rootTransformer.processPossibleTypeRange();
+    return (
+      this.rootTransformer.processPossibleArrowParamEnd() ||
+      this.rootTransformer.processPossibleTypeRange()
+    );
   }
 }
diff --git a/src/transformers/RootTransformer.ts b/src/transformers/RootTransformer.ts
index 6a5e64cc..78e7511e 100644
--- a/src/transformers/RootTransformer.ts
+++ b/src/transformers/RootTransformer.ts
@@ -312,6 +312,32 @@ export default class RootTransformer {
     ].join(";");
   }
 
+  /**
+   * Normally it's ok to simply remove type tokens, but we need to be more careful when dealing with
+   * arrow function return types since they can confuse the parser. In that case, we want to move
+   * the close-paren to the same line as the arrow.
+   *
+   * See https://github.com/alangpierce/sucrase/issues/391 for more details.
+   */
+  processPossibleArrowParamEnd(): boolean {
+    if (this.tokens.matches2(tt.parenR, tt.colon) && this.tokens.tokenAtRelativeIndex(1).isType) {
+      let nextNonTypeIndex = this.tokens.currentIndex() + 1;
+      // Look ahead to see if this is an arrow function or something else.
+      while (this.tokens.tokens[nextNonTypeIndex].isType) {
+        nextNonTypeIndex++;
+      }
+      if (this.tokens.matches1AtIndex(nextNonTypeIndex, tt.arrow)) {
+        this.tokens.removeInitialToken();
+        while (this.tokens.currentIndex() < nextNonTypeIndex) {
+          this.tokens.removeToken();
+        }
+        this.tokens.replaceTokenTrimmingLeftWhitespace(") =>");
+        return true;
+      }
+    }
+    return false;
+  }
+
   processPossibleTypeRange(): boolean {
     if (this.tokens.currentToken().isType) {
       this.tokens.removeInitialToken();
diff --git a/src/transformers/TypeScriptTransformer.ts b/src/transformers/TypeScriptTransformer.ts
index 12ece2c1..31f9f130 100644
--- a/src/transformers/TypeScriptTransformer.ts
+++ b/src/transformers/TypeScriptTransformer.ts
@@ -14,8 +14,10 @@ export default class TypeScriptTransformer extends Transformer {
   }
 
   process(): boolean {
-    const processedType = this.rootTransformer.processPossibleTypeRange();
-    if (processedType) {
+    if (
+      this.rootTransformer.processPossibleArrowParamEnd() ||
+      this.rootTransformer.processPossibleTypeRange()
+    ) {
       return true;
     }
     if (
diff --git a/test/types-test.ts b/test/types-test.ts
index 7d2e1e2f..8ccb18a9 100644
--- a/test/types-test.ts
+++ b/test/types-test.ts
@@ -1,5 +1,10 @@
 import {ESMODULE_PREFIX} from "./prefixes";
-import {assertResult} from "./util";
+import {assertExpectations, assertResult, Expectations} from "./util";
+
+function assertTypeScriptAndFlowExpectations(code: string, expectations: Expectations): void {
+  assertExpectations(code, expectations, {transforms: ["jsx", "imports", "typescript"]});
+  assertExpectations(code, expectations, {transforms: ["jsx", "imports", "flow"]});
+}
 
 function assertTypeScriptAndFlowResult(code: string, expectedResult: string): void {
   assertResult(code, expectedResult, {transforms: ["jsx", "imports", "typescript"]});
@@ -487,4 +492,18 @@ describe("type transforms", () => {
     `,
     );
   });
+
+  it("does not produce code with a syntax error on multiline return types", () => {
+    assertTypeScriptAndFlowExpectations(
+      `
+      const multiLineReturn = (
+        x: number
+      ): {
+        value: number;
+      } => ({value: x}); 
+      setOutput(multiLineReturn(5).value)
+    `,
+      {expectedOutput: 5},
+    );
+  });
 });
diff --git a/test/util.ts b/test/util.ts
index 608570b1..6ff628ab 100644
--- a/test/util.ts
+++ b/test/util.ts
@@ -1,13 +1,36 @@
 import * as assert from "assert";
+import vm from "vm";
 
 import {Options, transform} from "../src";
 
+export interface Expectations {
+  expectedResult?: string;
+  expectedOutput?: unknown;
+}
+
+export function assertExpectations(
+  code: string,
+  expectations: Expectations,
+  options: Options,
+): void {
+  const resultCode = transform(code, options).code;
+  if ("expectedResult" in expectations) {
+    assert.strictEqual(resultCode, expectations.expectedResult);
+  }
+  if ("expectedOutput" in expectations) {
+    const outputs: Array<unknown> = [];
+    vm.runInNewContext(resultCode, {setOutput: (value: unknown) => outputs.push(value)});
+    assert.strictEqual(outputs.length, 1, "setOutput should be called exactly once");
+    assert.deepStrictEqual(outputs[0], expectations.expectedOutput);
+  }
+}
+
 export function assertResult(
   code: string,
   expectedResult: string,
   options: Options = {transforms: ["jsx", "imports"]},
 ): void {
-  assert.strictEqual(transform(code, options).code, expectedResult);
+  assertExpectations(code, {expectedResult}, options);
 }
 
 export function devProps(lineNumber: number): string {