From 5a70d1425187e8fc69220c876fd12904cbfc36f8 Mon Sep 17 00:00:00 2001
From: Alan Pierce <alangpierce@gmail.com>
Date: Mon, 30 Dec 2019 22:08:55 -0800
Subject: [PATCH] Support TS type argument syntax in optional chaining

Progress toward #461

The previous implemented worked for normal JS syntax, but `f?.<T>()` needs a
special case when working with tokens. Also, regular type argument syntax wasn't
marking the open-paren correctly, so this fixes that as well.
---
 src/parser/plugins/typescript.ts                    | 12 ++++++++++++
 .../OptionalChainingNullishTransformer.ts           |  5 ++++-
 test/typescript-test.ts                             | 13 ++++++++++++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/parser/plugins/typescript.ts b/src/parser/plugins/typescript.ts
index 0e7e3bb1..0a9d14cf 100644
--- a/src/parser/plugins/typescript.ts
+++ b/src/parser/plugins/typescript.ts
@@ -1079,6 +1079,8 @@ export function tsParseSubscript(
     }
     tsParseTypeArguments();
     if (!noCalls && eat(tt.parenL)) {
+      // With f<T>(), the subscriptStartIndex marker is on the ( token.
+      state.tokens[state.tokens.length - 1].subscriptStartIndex = startTokenIndex;
       parseCallExpressionArguments();
     } else if (match(tt.backQuote)) {
       // Tagged template with a type argument.
@@ -1092,6 +1094,16 @@ export function tsParseSubscript(
     } else {
       return;
     }
+  } else if (!noCalls && match(tt.questionDot) && lookaheadType() === tt.lessThan) {
+    // If we see f?.<, then this must be an optional call with a type argument.
+    next();
+    state.tokens[startTokenIndex].isOptionalChainStart = true;
+    // With f?.<T>(), the subscriptStartIndex marker is on the ?. token.
+    state.tokens[state.tokens.length - 1].subscriptStartIndex = startTokenIndex;
+
+    tsParseTypeArguments();
+    expect(tt.parenL);
+    parseCallExpressionArguments();
   }
   baseParseSubscript(startTokenIndex, noCalls, stopState);
 }
diff --git a/src/transformers/OptionalChainingNullishTransformer.ts b/src/transformers/OptionalChainingNullishTransformer.ts
index c4bd173c..37f0b7c6 100644
--- a/src/transformers/OptionalChainingNullishTransformer.ts
+++ b/src/transformers/OptionalChainingNullishTransformer.ts
@@ -46,7 +46,10 @@ export default class OptionalChainingNullishTransformer extends Transformer {
       } else {
         arrowStartSnippet = `${param} => ${param}`;
       }
-      if (this.tokens.matches2(tt.questionDot, tt.parenL)) {
+      if (
+        this.tokens.matches2(tt.questionDot, tt.parenL) ||
+        this.tokens.matches2(tt.questionDot, tt.lessThan)
+      ) {
         this.tokens.replaceTokenTrimmingLeftWhitespace(`, 'optionalCall', ${arrowStartSnippet}`);
       } else if (this.tokens.matches2(tt.questionDot, tt.bracketL)) {
         this.tokens.replaceTokenTrimmingLeftWhitespace(`, 'optionalAccess', ${arrowStartSnippet}`);
diff --git a/test/typescript-test.ts b/test/typescript-test.ts
index 1b7b0f73..1395ef60 100644
--- a/test/typescript-test.ts
+++ b/test/typescript-test.ts
@@ -1901,7 +1901,7 @@ describe("typescript transform", () => {
       example.inner?.greet<string>()
     `,
       `"use strict";${OPTIONAL_CHAIN_PREFIX}
-      _optionalChain([example, 'access', _ => _.inner, 'optionalAccess', _2 => _2.greet()])
+      _optionalChain([example, 'access', _ => _.inner, 'optionalAccess', _2 => _2.greet, 'call', _3 => _3()])
     `,
     );
   });
@@ -1931,4 +1931,15 @@ describe("typescript transform", () => {
     `,
     );
   });
+
+  it("supports type arguments with optional chaining", () => {
+    assertTypeScriptResult(
+      `
+      const x = a.b?.<number>();
+    `,
+      `"use strict";${OPTIONAL_CHAIN_PREFIX}
+      const x = _optionalChain([a, 'access', _ => _.b, 'optionalCall', _2 => _2()]);
+    `,
+    );
+  });
 });