Fix emit for optional chain with non-null assertion

src/compiler/binder.ts

@@ -833,6 +833,9 @@ namespace ts {
                 case SyntaxKind.CallExpression:
                     bindCallExpressionFlow(<CallExpression>node);
                     break;
+                case SyntaxKind.NonNullExpression:
+                    bindNonNullExpressionFlow(<NonNullExpression>node);
+                    break;
                 case SyntaxKind.JSDocTypedefTag:
                 case SyntaxKind.JSDocCallbackTag:
                 case SyntaxKind.JSDocEnumTag:
@@ -1668,18 +1671,22 @@ namespace ts {
         }
 
         function bindOptionalChainRest(node: OptionalChain) {
-            bind(node.questionDotToken);
             switch (node.kind) {
                 case SyntaxKind.PropertyAccessExpression:
+                    bind(node.questionDotToken);
                     bind(node.name);
                     break;
                 case SyntaxKind.ElementAccessExpression:
+                    bind(node.questionDotToken);
                     bind(node.argumentExpression);
                     break;
                 case SyntaxKind.CallExpression:
+                    bind(node.questionDotToken);
                     bindEach(node.typeArguments);
                     bindEach(node.arguments);
                     break;
+                case SyntaxKind.NonNullExpression:
+                    break;
             }
         }
 
@@ -1695,7 +1702,7 @@ namespace ts {
             // and build it's CFA graph as if it were the first condition (`a && ...`). Then we bind the rest
             // of the node as part of the "true" branch, and continue to do so as we ascend back up to the outermost
             // chain node. We then treat the entire node as the right side of the expression.
-            const preChainLabel = node.questionDotToken ? createBranchLabel() : undefined;
+            const preChainLabel = isOptionalChainRoot(node) ? createBranchLabel() : undefined;
             bindOptionalExpression(node.expression, preChainLabel || trueTarget, falseTarget);
             if (preChainLabel) {
                 currentFlow = finishFlowLabel(preChainLabel);
@@ -1718,7 +1725,16 @@ namespace ts {
             }
         }
 
-        function bindAccessExpressionFlow(node: AccessExpression) {
+        function bindNonNullExpressionFlow(node: NonNullExpression | NonNullChain) {
+            if (isOptionalChain(node)) {
+                bindOptionalChainFlow(node);
+            }
+            else {
+                bindEachChild(node);
+            }
+        }
+
+        function bindAccessExpressionFlow(node: AccessExpression | PropertyAccessChain | ElementAccessChain) {
             if (isOptionalChain(node)) {
                 bindOptionalChainFlow(node);
             }
@@ -1727,7 +1743,7 @@ namespace ts {
             }
         }
 
-        function bindCallExpressionFlow(node: CallExpression) {
+        function bindCallExpressionFlow(node: CallExpression | CallChain) {
             if (isOptionalChain(node)) {
                 bindOptionalChainFlow(node);
             }

src/compiler/checker.ts

@@ -26158,8 +26158,15 @@ namespace ts {
             return targetType;
         }
 
+        function checkNonNullChain(node: NonNullChain) {
+            const leftType = checkExpression(node.expression);
+            const nonOptionalType = getOptionalExpressionType(leftType, node.expression);
+            return propagateOptionalTypeMarker(getNonNullableType(nonOptionalType), node, nonOptionalType !== leftType);
+        }
+
         function checkNonNullAssertion(node: NonNullExpression) {
-            return getNonNullableType(checkExpression(node.expression));
+            return node.flags & NodeFlags.OptionalChain ? checkNonNullChain(node as NonNullChain) :
+                getNonNullableType(checkExpression(node.expression));
         }
 
         function checkMetaProperty(node: MetaProperty): Type {

src/compiler/debug.ts

@@ -166,7 +166,7 @@ namespace ts {
             }
         }
 
-        export function assertNotNode<T extends Node, U extends T>(node: T | undefined, test: (node: T) => node is U, message?: string, stackCrawlMark?: AnyFunction): asserts node is Exclude<T, U>;
+        export function assertNotNode<T extends Node, U extends T>(node: T | undefined, test: (node: Node) => node is U, message?: string, stackCrawlMark?: AnyFunction): asserts node is Exclude<T, U>;
         export function assertNotNode(node: Node | undefined, test: ((node: Node) => boolean) | undefined, message?: string, stackCrawlMark?: AnyFunction): void;
         export function assertNotNode(node: Node | undefined, test: ((node: Node) => boolean) | undefined, message?: string, stackCrawlMark?: AnyFunction) {
             if (shouldAssertFunction(AssertionLevel.Normal, "assertNotNode")) {

src/compiler/factory.ts

@@ -1335,9 +1335,11 @@ namespace ts {
 
     export const enum OuterExpressionKinds {
         Parentheses = 1 << 0,
-        Assertions = 1 << 1,
-        PartiallyEmittedExpressions = 1 << 2,
+        TypeAssertions = 1 << 1,
+        NonNullAssertions = 1 << 2,
+        PartiallyEmittedExpressions = 1 << 3,
 
+        Assertions = TypeAssertions | NonNullAssertions,
         All = Parentheses | Assertions | PartiallyEmittedExpressions
     }
 
@@ -1349,8 +1351,9 @@ namespace ts {
                 return (kinds & OuterExpressionKinds.Parentheses) !== 0;
             case SyntaxKind.TypeAssertionExpression:
             case SyntaxKind.AsExpression:
+                return (kinds & OuterExpressionKinds.TypeAssertions) !== 0;
             case SyntaxKind.NonNullExpression:
-                return (kinds & OuterExpressionKinds.Assertions) !== 0;
+                return (kinds & OuterExpressionKinds.NonNullAssertions) !== 0;
             case SyntaxKind.PartiallyEmittedExpression:
                 return (kinds & OuterExpressionKinds.PartiallyEmittedExpressions) !== 0;
         }
@@ -1360,34 +1363,16 @@ namespace ts {
     export function skipOuterExpressions(node: Expression, kinds?: OuterExpressionKinds): Expression;
     export function skipOuterExpressions(node: Node, kinds?: OuterExpressionKinds): Node;
     export function skipOuterExpressions(node: Node, kinds = OuterExpressionKinds.All) {
-        let previousNode: Node;
-        do {
-            previousNode = node;
-            if (kinds & OuterExpressionKinds.Parentheses) {
-                node = skipParentheses(node);
-            }
-
-            if (kinds & OuterExpressionKinds.Assertions) {
-                node = skipAssertions(node);
-            }
-
-            if (kinds & OuterExpressionKinds.PartiallyEmittedExpressions) {
-                node = skipPartiallyEmittedExpressions(node);
-            }
+        while (isOuterExpression(node, kinds)) {
+            node = node.expression;
         }
-        while (previousNode !== node);
-
         return node;
     }
 
     export function skipAssertions(node: Expression): Expression;
     export function skipAssertions(node: Node): Node;
     export function skipAssertions(node: Node): Node {
-        while (isAssertionExpression(node) || node.kind === SyntaxKind.NonNullExpression) {
-            node = (<AssertionExpression | NonNullExpression>node).expression;
-        }
-
-        return node;
+        return skipOuterExpressions(node, OuterExpressionKinds.Assertions);
     }
 
     function updateOuterExpression(outerExpression: OuterExpression, expression: Expression) {

src/compiler/factoryPublic.ts

@@ -1076,10 +1076,8 @@ namespace ts {
     }
 
     export function updatePropertyAccess(node: PropertyAccessExpression, expression: Expression, name: Identifier | PrivateIdentifier) {
-        if (isOptionalChain(node) && isIdentifier(node.name) && isIdentifier(name)) {
-            // Not sure why this cast was necessary: the previous line should already establish that node.name is an identifier
-            const theNode = node as (typeof node & { name: Identifier });
-            return updatePropertyAccessChain(theNode, expression, node.questionDotToken, name);
+        if (isPropertyAccessChain(node)) {
+            return updatePropertyAccessChain(node, expression, node.questionDotToken, cast(name, isIdentifier));
         }
         // Because we are updating existed propertyAccess we want to inherit its emitFlags
         // instead of using the default from createPropertyAccess
@@ -1653,11 +1651,28 @@ namespace ts {
     }
 
     export function updateNonNullExpression(node: NonNullExpression, expression: Expression) {
+        if (isNonNullChain(node)) {
+            return updateNonNullChain(node, expression);
+        }
         return node.expression !== expression
             ? updateNode(createNonNullExpression(expression), node)
             : node;
     }
 
+    export function createNonNullChain(expression: Expression) {
+        const node = <NonNullChain>createSynthesizedNode(SyntaxKind.NonNullExpression);
+        node.flags |= NodeFlags.OptionalChain;
+        node.expression = parenthesizeForAccess(expression);
+        return node;
+    }
+
+    export function updateNonNullChain(node: NonNullChain, expression: Expression) {
+        Debug.assert(!!(node.flags & NodeFlags.OptionalChain), "Cannot update a NonNullExpression using updateNonNullChain. Use updateNonNullExpression instead.");
+        return node.expression !== expression
+            ? updateNode(createNonNullChain(expression), node)
+            : node;
+    }
+
     export function createMetaProperty(keywordToken: MetaProperty["keywordToken"], name: Identifier) {
         const node = <MetaProperty>createSynthesizedNode(SyntaxKind.MetaProperty);
         node.keywordToken = keywordToken;

src/compiler/parser.ts

@@ -4752,12 +4752,34 @@ namespace ts {
                 && lookAhead(nextTokenIsIdentifierOrKeywordOrOpenBracketOrTemplate);
         }
 
+        function tryReparseOptionalChain(node: Expression) {
+            if (node.flags & NodeFlags.OptionalChain) {
+                return true;
+            }
+            // check for an optional chain in a non-null expression
+            if (isNonNullExpression(node)) {
+                let expr = node.expression;
+                while (isNonNullExpression(expr) && !(expr.flags & NodeFlags.OptionalChain)) {
+                    expr = expr.expression;
+                }
+                if (expr.flags & NodeFlags.OptionalChain) {
+                    // this is part of an optional chain. Walk down from `node` to `expression` and set the flag.
+                    while (isNonNullExpression(node)) {
+                        node.flags |= NodeFlags.OptionalChain;
+                        node = node.expression;
+                    }
+                    return true;
+                }
+            }
+            return false;
+        }
+
         function parsePropertyAccessExpressionRest(expression: LeftHandSideExpression, questionDotToken: QuestionDotToken | undefined) {
             const propertyAccess = <PropertyAccessExpression>createNode(SyntaxKind.PropertyAccessExpression, expression.pos);
             propertyAccess.expression = expression;
             propertyAccess.questionDotToken = questionDotToken;
             propertyAccess.name = parseRightSideOfDot(/*allowIdentifierNames*/ true, /*allowPrivateIdentifiers*/ true);
-            if (questionDotToken || expression.flags & NodeFlags.OptionalChain) {
+            if (questionDotToken || tryReparseOptionalChain(expression)) {
                 propertyAccess.flags |= NodeFlags.OptionalChain;
                 if (isPrivateIdentifier(propertyAccess.name)) {
                     parseErrorAtRange(propertyAccess.name, Diagnostics.An_optional_chain_cannot_contain_private_identifiers);
@@ -4783,7 +4805,7 @@ namespace ts {
             }
 
             parseExpected(SyntaxKind.CloseBracketToken);
-            if (questionDotToken || expression.flags & NodeFlags.OptionalChain) {
+            if (questionDotToken || tryReparseOptionalChain(expression)) {
                 indexedAccess.flags |= NodeFlags.OptionalChain;
             }
             return finishNode(indexedAccess);
@@ -4870,7 +4892,7 @@ namespace ts {
                         callExpr.questionDotToken = questionDotToken;
                         callExpr.typeArguments = typeArguments;
                         callExpr.arguments = parseArgumentList();
-                        if (questionDotToken || expression.flags & NodeFlags.OptionalChain) {
+                        if (questionDotToken || tryReparseOptionalChain(expression)) {
                             callExpr.flags |= NodeFlags.OptionalChain;
                         }
                         expression = finishNode(callExpr);
@@ -4882,7 +4904,7 @@ namespace ts {
                     callExpr.expression = expression;
                     callExpr.questionDotToken = questionDotToken;
                     callExpr.arguments = parseArgumentList();
-                    if (questionDotToken || expression.flags & NodeFlags.OptionalChain) {
+                    if (questionDotToken || tryReparseOptionalChain(expression)) {
                         callExpr.flags |= NodeFlags.OptionalChain;
                     }
                     expression = finishNode(callExpr);

src/compiler/transformers/es2020.ts

@@ -42,9 +42,11 @@ namespace ts {
       }
 
       function flattenChain(chain: OptionalChain) {
+          Debug.assertNotNode(chain, isNonNullChain);
           const links: OptionalChain[] = [chain];
           while (!chain.questionDotToken && !isTaggedTemplateExpression(chain)) {
-              chain = cast(chain.expression, isOptionalChain);
+              chain = cast(skipPartiallyEmittedExpressions(chain.expression), isOptionalChain);
+              Debug.assertNotNode(chain, isNonNullChain);
               links.unshift(chain);
           }
           return { expression: chain.expression, chain: links };

src/compiler/types.ts

@@ -1933,6 +1933,7 @@ namespace ts {
         | PropertyAccessChain
         | ElementAccessChain
         | CallChain
+        | NonNullChain
         ;
 
     /* @internal */
@@ -2027,6 +2028,10 @@ namespace ts {
         expression: Expression;
     }
 
+    export interface NonNullChain extends NonNullExpression {
+        _optionalChainBrand: any;
+    }
+
     // NOTE: MetaProperty is really a MemberExpression, but we consider it a PrimaryExpression
     //       for the same reasons we treat NewExpression as a PrimaryExpression.
     export interface MetaProperty extends PrimaryExpression {

src/compiler/utilities.ts

@@ -2617,11 +2617,7 @@ namespace ts {
     export function skipParentheses(node: Expression): Expression;
     export function skipParentheses(node: Node): Node;
     export function skipParentheses(node: Node): Node {
-        while (node.kind === SyntaxKind.ParenthesizedExpression) {
-            node = (node as ParenthesizedExpression).expression;
-        }
-
-        return node;
+        return skipOuterExpressions(node, OuterExpressionKinds.Parentheses);
     }
 
     function skipParenthesesUp(node: Node): Node {

src/compiler/utilitiesPublic.ts

@@ -1086,17 +1086,18 @@ namespace ts {
         return isCallExpression(node) && !!(node.flags & NodeFlags.OptionalChain);
     }
 
-    export function isOptionalChain(node: Node): node is PropertyAccessChain | ElementAccessChain | CallChain {
+    export function isOptionalChain(node: Node): node is PropertyAccessChain | ElementAccessChain | CallChain | NonNullChain {
         const kind = node.kind;
         return !!(node.flags & NodeFlags.OptionalChain) &&
             (kind === SyntaxKind.PropertyAccessExpression
                 || kind === SyntaxKind.ElementAccessExpression
-                || kind === SyntaxKind.CallExpression);
+                || kind === SyntaxKind.CallExpression
+                || kind === SyntaxKind.NonNullExpression);
     }
 
     /* @internal */
     export function isOptionalChainRoot(node: Node): node is OptionalChainRoot {
-        return isOptionalChain(node) && !!node.questionDotToken;
+        return isOptionalChain(node) && !isNonNullExpression(node) && !!node.questionDotToken;
     }
 
     /**
@@ -1111,17 +1112,18 @@ namespace ts {
      * Determines whether a node is the outermost `OptionalChain` in an ECMAScript `OptionalExpression`:
      *
      * 1. For `a?.b.c`, the outermost chain is `a?.b.c` (`c` is the end of the chain starting at `a?.`)
-     * 2. For `(a?.b.c).d`, the outermost chain is `a?.b.c` (`c` is the end of the chain starting at `a?.` since parens end the chain)
-     * 3. For `a?.b.c?.d`, both `a?.b.c` and `a?.b.c?.d` are outermost (`c` is the end of the chain starting at `a?.`, and `d` is
+     * 2. For `a?.b!`, the outermost chain is `a?.b!` (`c` is the end of the chain starting at `a?.`)
+     * 3. For `(a?.b.c).d`, the outermost chain is `a?.b.c` (`c` is the end of the chain starting at `a?.` since parens end the chain)
+     * 4. For `a?.b.c?.d`, both `a?.b.c` and `a?.b.c?.d` are outermost (`c` is the end of the chain starting at `a?.`, and `d` is
      *   the end of the chain starting at `c?.`)
-     * 4. For `a?.(b?.c).d`, both `b?.c` and `a?.(b?.c)d` are outermost (`c` is the end of the chain starting at `b`, and `d` is
+     * 5. For `a?.(b?.c).d`, both `b?.c` and `a?.(b?.c)d` are outermost (`c` is the end of the chain starting at `b`, and `d` is
      *   the end of the chain starting at `a?.`)
      */
     /* @internal */
     export function isOutermostOptionalChain(node: OptionalChain) {
-        return !isOptionalChain(node.parent) // cases 1 and 2
-            || isOptionalChainRoot(node.parent) // case 3
-            || node !== node.parent.expression; // case 4
+        return !isOptionalChain(node.parent) // cases 1, 2, and 3
+            || isOptionalChainRoot(node.parent) // case 4
+            || node !== node.parent.expression; // case 5
     }
 
     export function isNullishCoalesce(node: Node) {
@@ -1152,11 +1154,7 @@ namespace ts {
     export function skipPartiallyEmittedExpressions(node: Expression): Expression;
     export function skipPartiallyEmittedExpressions(node: Node): Node;
     export function skipPartiallyEmittedExpressions(node: Node) {
-        while (node.kind === SyntaxKind.PartiallyEmittedExpression) {
-            node = (<PartiallyEmittedExpression>node).expression;
-        }
-
-        return node;
+        return skipOuterExpressions(node, OuterExpressionKinds.PartiallyEmittedExpressions);
     }
 
     export function isFunctionExpression(node: Node): node is FunctionExpression {
@@ -1231,6 +1229,10 @@ namespace ts {
         return node.kind === SyntaxKind.NonNullExpression;
     }
 
+    export function isNonNullChain(node: Node): node is NonNullChain {
+        return isNonNullExpression(node) && !!(node.flags & NodeFlags.OptionalChain);
+    }
+
     export function isMetaProperty(node: Node): node is MetaProperty {
         return node.kind === SyntaxKind.MetaProperty;
     }

src/services/utilities.ts

@@ -1355,7 +1355,7 @@ namespace ts {
     export function getPossibleGenericSignatures(called: Expression, typeArgumentCount: number, checker: TypeChecker): readonly Signature[] {
         let type = checker.getTypeAtLocation(called);
         if (isOptionalChain(called.parent)) {
-            type = removeOptionality(type, !!called.parent.questionDotToken, /*isOptionalChain*/ true);
+            type = removeOptionality(type, isOptionalChainRoot(called.parent), /*isOptionalChain*/ true);
         }
 
         const signatures = isNewExpression(called.parent) ? type.getConstructSignatures() : type.getCallSignatures();