Distinguish between ternary's : and arrow fn's return type (#573)

* Distinguish between ternary's : and arrow fn's return type

* Correctly parse nested arrow functions inside conditional expressions

Defer the conversion of arrow function parameters to assignable nodes so that
it is possible to use the (invalid) ast to get the exact position of the (wrong)
arrow functions.

* Check params of arrow fns w/ type params or w/o return type

* Fix also async functions

* Add test from prettier

prettier/prettier#2194

* Don't check arrow params if they are valid at the first attemp

* Use state instead of relying on the "noArrowParamsConversion" parameter

* Remove noArrowParamsConversion

src/parser/expression.js

@@ -1081,11 +1081,15 @@ export default class ExpressionParser extends LValParser {
 
   parseArrowExpression(node: N.ArrowFunctionExpression, params: N.Expression[], isAsync?: boolean): N.ArrowFunctionExpression {
     this.initFunction(node, isAsync);
-    node.params = this.toAssignableList(params, true, "arrow function parameters");
+    this.setArrowFunctionParameters(node, params);
     this.parseFunctionBody(node, true);
     return this.finishNode(node, "ArrowFunctionExpression");
   }
 
+  setArrowFunctionParameters(node: N.ArrowFunctionExpression, params: N.Expression[]): N.ArrowFunctionExpression {
+    node.params = this.toAssignableList(params, true, "arrow function parameters");
+  }
+
   isStrictBody(node: { body: N.BlockStatement }, isExpression?: boolean): boolean {
     if (!isExpression && node.body.directives.length) {
       for (const directive of node.body.directives) {
@@ -1120,12 +1124,19 @@ export default class ExpressionParser extends LValParser {
     }
     this.state.inAsync = oldInAsync;
 
+    this.checkFunctionNameAndParams(node, allowExpression);
+  }
+
+  checkFunctionNameAndParams(
+    node: N.Function,
+    isArrowFunction?: boolean
+  ): void {
     // If this is a strict mode function, verify that argument names
     // are not repeated, and it does not try to bind the words `eval`
     // or `arguments`.
-    const isStrict = this.isStrictBody(node, isExpression);
-    // Also check when allowExpression === true for arrow functions
-    const checkLVal = this.state.strict || allowExpression || isStrict;
+    const isStrict = this.isStrictBody(node, node.expression);
+    // Also check for arrow functions
+    const checkLVal = this.state.strict || isStrict || isArrowFunction;
 
     if (isStrict && node.id && node.id.type === "Identifier" && node.id.name === "yield") {
       this.raise(node.id.start, "Binding yield in strict mode");

src/plugins/flow.js

@@ -37,6 +37,16 @@ const exportSuggestions = {
   interface: "export interface",
 };
 
+// Like Array#filter, but returns a tuple [ acceptedElements, discardedElements ]
+function partition<T>(list: T[], test: (T, number, T[]) => ?boolean): [ T[], T[] ] {
+  const list1 = [];
+  const list2 = [];
+  for (let i = 0; i < list.length; i++) {
+    (test(list[i], i, list) ? list1 : list2).push(list[i]);
+  }
+  return [ list1, list2 ];
+}
+
 export default (superClass: Class<Parser>): Class<Parser> => class extends superClass {
   flowParseTypeInitialiser(tok?: TokenType): N.FlowType {
     const oldInType = this.state.inType;
@@ -1007,9 +1017,15 @@ export default (superClass: Class<Parser>): Class<Parser> => class extends super
   // Overrides
   // ==================================
 
-  // plain function return types: function name(): string {}
   parseFunctionBody(node: N.Function, allowExpression?: boolean): void {
-    if (this.match(tt.colon) && !allowExpression) {
+    if (allowExpression && this.state.noArrowParamsConversionAt.indexOf(node.start) !== -1) {
+      this.state.noArrowParamsConversionAt.push(this.state.start);
+      super.parseFunctionBody(node, allowExpression);
+      this.state.noArrowParamsConversionAt.pop();
+
+      return;
+    } else if (this.match(tt.colon)) {
+      // plain function return types: function name(): string {}
       // if allowExpression is true then we're parsing an arrow function and if
       // there's a return type then it's been handled elsewhere
       const typeNode = this.startNode();
@@ -1074,9 +1090,11 @@ export default (superClass: Class<Parser>): Class<Parser> => class extends super
   }
 
   parseConditional(expr: N.Expression, noIn: ?boolean, startPos: number, startLoc: Position, refNeedsArrowPos?: ?Pos): N.Expression {
+    if (!this.match(tt.question)) return expr;
+
     // only do the expensive clone if there is a question mark
     // and if we come from inside parens
-    if (refNeedsArrowPos && this.match(tt.question)) {
+    if (refNeedsArrowPos) {
       const state = this.state.clone();
       try {
         return super.parseConditional(expr, noIn, startPos, startLoc);
@@ -1092,7 +1110,129 @@ export default (superClass: Class<Parser>): Class<Parser> => class extends super
       }
     }
 
-    return super.parseConditional(expr, noIn, startPos, startLoc);
+    this.expect(tt.question);
+    const state = this.state.clone();
+    const originalNoArrowAt = this.state.noArrowAt;
+    const node = this.startNodeAt(startPos, startLoc);
+    let { consequent, failed } = this.tryParseConditionalConsequent();
+    let [ valid, invalid ] = this.getArrowLikeExpressions(consequent);
+
+    if (failed || invalid.length > 0) {
+      const noArrowAt = [ ...originalNoArrowAt ];
+
+      if (invalid.length > 0) {
+        this.state = state;
+        this.state.noArrowAt = noArrowAt;
+
+        for (let i = 0; i < invalid.length; i++) {
+          noArrowAt.push(invalid[i].start);
+        }
+
+        ({ consequent, failed } = this.tryParseConditionalConsequent());
+        [ valid, invalid ] = this.getArrowLikeExpressions(consequent);
+      }
+
+      if (failed && valid.length > 1) {
+        // if there are two or more possible correct ways of parsing, throw an
+        // error.
+        // e.g.   Source: a ? (b): c => (d): e => f
+        //      Result 1: a ? b : (c => ((d): e => f))
+        //      Result 2: a ? ((b): c => d) : (e => f)
+        this.raise(
+          state.start,
+          "Ambiguous expression: wrap the arrow functions in parentheses to disambiguate."
+        );
+      }
+
+      if (failed && valid.length === 1) {
+        this.state = state;
+        this.state.noArrowAt = noArrowAt.concat(valid[0].start);
+        ({ consequent, failed } = this.tryParseConditionalConsequent());
+      }
+
+      this.getArrowLikeExpressions(consequent, true);
+    }
+
+    this.state.noArrowAt = originalNoArrowAt;
+    this.expect(tt.colon);
+
+    node.test = expr;
+    node.consequent = consequent;
+    node.alternate = this.forwardNoArrowParamsConversionAt(node, () => (
+      this.parseMaybeAssign(noIn, undefined, undefined, undefined)
+    ));
+
+    return this.finishNode(node, "ConditionalExpression");
+  }
+
+  tryParseConditionalConsequent(): { consequent: N.Expression, failed: boolean } {
+    this.state.noArrowParamsConversionAt.push(this.state.start);
+
+    const consequent = this.parseMaybeAssign();
+    const failed = !this.match(tt.colon);
+
+    this.state.noArrowParamsConversionAt.pop();
+
+    return { consequent, failed };
+  }
+
+  // Given an expression, walks throught its arrow functions whose body is
+  // an expression and throught conditional expressions. It returns every
+  // function which has been parsed with a return type but could have been
+  // parenthesized expressions.
+  // These functions are separated into two arrays: one containing the ones
+  // whose parameters can be converted to assignable lists, one containing the
+  // others.
+  getArrowLikeExpressions(node: N.Expression, disallowInvalid?: boolean): [ N.ArrowFunctionExpression[], N.ArrowFunctionExpression[] ] {
+    const stack = [ node ];
+    const arrows: N.ArrowFunctionExpression[] = [];
+
+    while (stack.length !== 0) {
+      const node = stack.pop();
+      if (node.type === "ArrowFunctionExpression") {
+        if (node.typeParameters || !node.returnType) {
+          // This is an arrow expression without ambiguity, so check its parameters
+          this.toAssignableList(node.params, true, "arrow function parameters");
+          // Use super's method to force the parameters to be checked
+          super.checkFunctionNameAndParams(node, true);
+        } else {
+          arrows.push(node);
+        }
+        stack.push(node.body);
+      } else if (node.type === "ConditionalExpression") {
+        stack.push(node.consequent);
+        stack.push(node.alternate);
+      }
+    }
+
+    if (disallowInvalid) {
+      for (let i = 0; i < arrows.length; i++) {
+        this.toAssignableList(arrows[i].params, true, "arrow function parameters");
+      }
+      return [ arrows, [] ];
+    }
+
+    return partition(arrows, (node) => {
+      try {
+        this.toAssignableList(node.params, true, "arrow function parameters");
+        return true;
+      } catch (err) {
+        return false;
+      }
+    });
+  }
+
+  forwardNoArrowParamsConversionAt<T>(node: N.Node, parse: () => T): T {
+    let result: T;
+    if (this.state.noArrowParamsConversionAt.indexOf(node.start) !== -1) {
+      this.state.noArrowParamsConversionAt.push(this.state.start);
+      result = parse();
+      this.state.noArrowParamsConversionAt.pop(this.state.start);
+    } else {
+      result = parse();
+    }
+
+    return result;
   }
 
   parseParenItem(node: N.Expression, startPos: number, startLoc: Position): N.Expression {
@@ -1510,8 +1650,9 @@ export default (superClass: Class<Parser>): Class<Parser> => class extends super
       let typeParameters;
       try {
         typeParameters = this.flowParseTypeParameterDeclaration();
-
-        arrowExpression = super.parseMaybeAssign(noIn, refShorthandDefaultPos, afterLeftParse, refNeedsArrowPos);
+        arrowExpression = this.forwardNoArrowParamsConversionAt(typeParameters, () => (
+          super.parseMaybeAssign(noIn, refShorthandDefaultPos, afterLeftParse, refNeedsArrowPos)
+        ));
         arrowExpression.typeParameters = typeParameters;
         this.resetStartLocationFromNode(arrowExpression, typeParameters);
       } catch (err) {
@@ -1570,4 +1711,43 @@ export default (superClass: Class<Parser>): Class<Parser> => class extends super
   shouldParseArrow(): boolean {
     return this.match(tt.colon) || super.shouldParseArrow();
   }
+
+  setArrowFunctionParameters(node: N.ArrowFunctionExpression, params: N.Expression[]): N.ArrowFunctionExpression {
+    if (this.state.noArrowParamsConversionAt.indexOf(node.start) !== -1) {
+      node.params = params;
+      return node;
+    }
+
+    return super.setArrowFunctionParameters(node, params);
+  }
+
+  checkFunctionNameAndParams(node: N.Function, isArrowFunction?: boolean): void {
+    if (isArrowFunction && this.state.noArrowParamsConversionAt.indexOf(node.start) !== -1) {
+      return;
+    }
+
+    return super.checkFunctionNameAndParams(node, isArrowFunction);
+  }
+
+  parseParenAndDistinguishExpression(canBeArrow: boolean): N.Expression {
+    return super.parseParenAndDistinguishExpression(
+      canBeArrow && this.state.noArrowAt.indexOf(this.state.start) === -1
+    );
+  }
+
+  parseSubscripts(base: N.Expression, startPos: number, startLoc: Position, noCalls?: boolean): N.Expression {
+    if (
+      base.type === "Identifier" && base.name === "async" &&
+      this.state.noArrowAt.indexOf(startPos) !== -1
+    ) {
+      this.next();
+
+      const node = this.startNodeAt(startPos, startLoc);
+      node.callee = base;
+      node.arguments = this.parseCallExpressionArguments(tt.parenR, false);
+      base = this.finishNode(node, "CallExpression");
+    }
+
+    return super.parseSubscripts(base, startPos, startLoc, noCalls);
+  }
 };

src/tokenizer/state.js

@@ -18,6 +18,9 @@ export default class State {
 
     this.potentialArrowAt = -1;
 
+    this.noArrowAt = [];
+    this.noArrowParamsConversionAt = [];
+
     this.inMethod =
       this.inFunction =
       this.inGenerator =
@@ -76,6 +79,20 @@ export default class State {
   // Used to signify the start of a potential arrow function
   potentialArrowAt: number;
 
+  // Used to signify the start of an expression which looks like a
+  // typed arrow function, but it isn't
+  // e.g. a ? (b) : c => d
+  //          ^
+  noArrowAt: number[];
+
+  // Used to signify the start of an expression whose params, if it looks like
+  // an arrow function, shouldn't be converted to assignable nodes.
+  // This is used to defer the validation of typed arrow functions inside
+  // conditional expressions.
+  // e.g. a ? (b) : c => d
+  //          ^
+  noArrowParamsConversionAt: number[];
+
   // Flags to track whether we are in a function, a generator.
   inFunction: boolean;
   inGenerator: boolean;

test/fixtures/flow/regression/issue-58-ambiguous/actual.js

@@ -0,0 +1,4 @@
+// This can be parsed in two ways:
+// a ? b : (c => ((d): e => f))
+// a ? ((b): c => d) : (e => f)
+a ? (b) : c => (d) : e => f;

test/fixtures/flow/regression/issue-58-ambiguous/options.json

@@ -0,0 +1,3 @@
+{
+    "throws": "Ambiguous expression: wrap the arrow functions in parentheses to disambiguate. (4:4)"
+}

test/fixtures/flow/regression/issue-58-failing-1/actual.js

@@ -0,0 +1,2 @@
+// Function which looks like a return type
+a ? (b) : (c => d) => e;

test/fixtures/flow/regression/issue-58-failing-1/options.json

@@ -0,0 +1,3 @@
+{
+    "throws": "Invalid left-hand side in arrow function parameters (2:11)"
+}

test/fixtures/flow/regression/issue-58-failing-2/actual.js

@@ -0,0 +1,2 @@
+// Invalid LHS parameter after type parameters
+a ? <T>(b => c) : d => (e) : f => g;

test/fixtures/flow/regression/issue-58-failing-2/options.json

@@ -0,0 +1,3 @@
+{
+    "throws": "Invalid left-hand side in arrow function parameters (2:8)"
+}

test/fixtures/flow/regression/issue-58-failing-3/actual.js

@@ -0,0 +1,2 @@
+// Invalid LHS parameter after type parameters
+a ? (b => c) => (e) : f => g;

test/fixtures/flow/regression/issue-58-failing-3/options.json

@@ -0,0 +1,3 @@
+{
+    "throws": "Invalid left-hand side in arrow function parameters (2:5)"
+}

test/fixtures/flow/regression/issue-58-failing-4/actual.js

@@ -0,0 +1,2 @@
+// Invalid LHS parameter
+a ? async (b => c) => (d) : f => g;

test/fixtures/flow/regression/issue-58-failing-4/options.json

@@ -0,0 +1,3 @@
+{
+    "throws": "Invalid left-hand side in arrow function parameters (2:11)"
+}

test/fixtures/flow/regression/issue-58/actual.js

@@ -0,0 +1,44 @@
+// Valid lhs value inside parentheses
+a ? (b) : c => d; // a ? b : (c => d)
+a ? (b) : c => d : e; // a ? ((b): c => d) : e
+a ? (b) : (c) : d => e; // a ? b : ((c): d => e)
+
+// Nested arrow function inside parentheses
+a ? (b = (c) => d) : e => f; // a ? (b = (c) => d) : (e => f)
+a ? (b = (c) => d) : e => f : g; // a ? ((b = (c) => d): e => f) : g
+
+// Nested conditional expressions
+    b ? c ? (d) : e => (f) : g : h; // b ? (c ? ((d): e => f) : g) : h
+a ? b ? c ? (d) : e => (f) : g : h; // a ? (b ? (c ? d : (e => f)) : g) : h
+
+a ? b ? (c) : (d) : (e) => f : g; // a ? (b ? c : ((d): e => f)) : g
+
+// Multiple arrow functions
+a ? (b) : c => d : (e) : f => g; // a ? ((b): c => d) : ((e): f => g)
+
+// Multiple nested arrow functions (<T> is needed to avoid ambiguities)
+a ? (b) : c => (d) : e => f : g; // a ? ((b): c => ((d): e => f)) : g
+a ? (b) : c => <T>(d) : e => f; // a ? b : (c => (<T>(d): e => f))
+a ? <T>(b) : c => (d) : e => f; // a ? (<T>(b): c => d) : (e => f)
+
+// Invalid lhs value inside parentheses
+a ? (b => c) : d => e; // a ? (b => c) : (d => e)
+a ? b ? (c => d) : e => f : g; // a ? (b ? (c => d) : (e => f)) : g
+
+// Invalid lhs value inside parentheses inside arrow function
+a ? (b) : c => (d => e) : f => g; // a ? ((b): c => (d => e)) : (f => g)
+a ? b ? (c => d) : e => (f => g) : h => i; // a ? (b ? (c => d) : (e => (f => g))) : (h => i)
+
+// Function as type annotation
+a ? (b) : (c => d) => e : f; // a ? ((b): (c => d) => e) : f
+
+// Async functions or calls
+a ? async (b) : c => d; // a ? (async(b)) : (c => d)
+a ? async (b) : c => d : e; // a ? (async (b): c => d) : e
+a ? async (b => c) : d => e; // a ? (async(b => c)) : (d => e)
+a ? async (b) => (c => d) : e => f; // a ? (async (b) => c => d) : (e => f)
+
+// https://github.com/prettier/prettier/issues/2194
+let icecream = what == "cone"
+  ? p => (!!p ? `here's your ${p} cone` : `just the empty cone for you`)
+  : p => `here's your ${p} ${what}`;