Improve optional chaining checker performance

Author: rbuckton

src/compiler/binder.ts

@@ -948,7 +948,7 @@ namespace ts {
             }
             if (expression.kind === SyntaxKind.TrueKeyword && flags & FlowFlags.FalseCondition ||
                 expression.kind === SyntaxKind.FalseKeyword && flags & FlowFlags.TrueCondition) {
-                if (!isOptionalChainRoot(expression.parent)) {
+                if (!isOptionalExpression(expression)) {
                     return unreachableFlow;
                 }
             }

src/compiler/checker.ts

@@ -333,10 +333,6 @@ namespace ts {
         /** This will be set during calls to `getResolvedSignature` where services determines an apparent number of arguments greater than what is actually provided. */
         let apparentArgumentCount: number | undefined;
 
-        // This object is reused for `checkOptionalExpression` return values to avoid frequent GC due to nursery object allocations.
-        // This object represents a pool-size of 1.
-        const pooledOptionalTypeResult: { isOptional: boolean, type: Type } = { isOptional: false, type: undefined! };
-
         // for public members that accept a Node or one of its subtypes, we must guard against
         // synthetic nodes created during transformations by calling `getParseTreeNode`.
         // for most of these, we perform the guard only on `checker` to avoid any possible
@@ -16677,51 +16673,50 @@ namespace ts {
         }
 
         function addOptionalTypeMarker(type: Type) {
-            return strictNullChecks ? getUnionType([type, optionalType]) : type;
+            if (strictNullChecks) {
+                if (isTypeAny(type)) return type;
+                if (type.optionalTypeOfType) {
+                    // The type aleady has an optional type cache.
+                    return type.optionalTypeOfType;
+                }
+                if (type.flags & TypeFlags.Union && type.nonOptionalTypeOfType) {
+                    // The type *is* the type produced by marking a type with the optional type.
+                    return type;
+                }
+                type.optionalTypeOfType = getUnionType([type, optionalType]);
+                type.optionalTypeOfType.nonOptionalTypeOfType = type;
+                return type.optionalTypeOfType;
+            }
+            return type;
         }
 
         function removeOptionalTypeMarker(type: Type): Type {
-            return strictNullChecks ? filterType(type, t => t !== optionalType) : type;
+            if (strictNullChecks) {
+                if (type.flags & TypeFlags.Union && type.nonOptionalTypeOfType) {
+                    return type.nonOptionalTypeOfType;
+                }
+                return filterType(type, t => t !== optionalType);
+            }
+            return type;
         }
 
         function propagateOptionalTypeMarker(type: Type, wasOptional: boolean) {
             return wasOptional ? addOptionalTypeMarker(type) : type;
         }
 
-        function createPooledOptionalTypeResult(isOptional: boolean, type: Type) {
-            pooledOptionalTypeResult.isOptional = isOptional;
-            pooledOptionalTypeResult.type = type;
-            return pooledOptionalTypeResult;
-        }
-
         function checkOptionalExpression(
-            parent: PropertyAccessExpression | QualifiedName | ElementAccessExpression | CallExpression,
             expression: Expression | QualifiedName,
             nullDiagnostic?: DiagnosticMessage,
             undefinedDiagnostic?: DiagnosticMessage,
             nullOrUndefinedDiagnostic?: DiagnosticMessage,
         ) {
-            let isOptional = false;
-            let type = checkExpression(expression);
-            if (isOptionalChain(parent)) {
-                if (parent.questionDotToken) {
-                    // If we have a questionDotToken then we are an OptionalExpression and should remove `null` and
-                    // `undefined` from the type and add the optionalType to the result, if needed.
-                    isOptional = isNullableType(type);
-                    return createPooledOptionalTypeResult(isOptional, isOptional ? getNonNullableType(type) : type);
-                }
-
-                // If we do not have a questionDotToken, then we are an OptionalChain and we remove the optionalType and
-                // indicate whether we need to add optionalType back into the result.
-                const nonOptionalType = removeOptionalTypeMarker(type);
-                if (nonOptionalType !== type) {
-                    isOptional = true;
-                    type = nonOptionalType;
-                }
-            }
-
-            type = checkNonNullType(type, expression, nullDiagnostic, undefinedDiagnostic, nullOrUndefinedDiagnostic);
-            return createPooledOptionalTypeResult(isOptional, type);
+            const exprType = checkExpression(expression);
+            const nonOptionalType =
+                isOptionalExpression(expression) ? getNonNullableType(exprType) : // 'a' in 'a?.b.c'
+                isOptionalChain(expression) ? removeOptionalTypeMarker(exprType) : // 'a?.b' in 'a?.b.c'
+                exprType;
+            const nonNullType = checkNonNullType(nonOptionalType, expression, nullDiagnostic, undefinedDiagnostic, nullOrUndefinedDiagnostic);
+            return exprType !== nonOptionalType ? addOptionalTypeMarker(nonNullType) : nonNullType;
         }
 
         /**
@@ -18756,7 +18751,7 @@ namespace ts {
                 // circularities in control flow analysis, we use getTypeOfDottedName when resolving the call
                 // target expression of an assertion.
                 const funcType = node.parent.kind === SyntaxKind.ExpressionStatement ? getTypeOfDottedName(node.expression, /*diagnostic*/ undefined) :
-                    node.expression.kind !== SyntaxKind.SuperKeyword ? checkOptionalExpression(node, node.expression).type :
+                    node.expression.kind !== SyntaxKind.SuperKeyword ? removeOptionalTypeMarker(checkOptionalExpression(node.expression)) :
                     undefined;
                 const signatures = getSignaturesOfType(funcType && getApparentType(funcType) || unknownType, SignatureKind.Call);
                 const candidate = signatures.length === 1 && !signatures[0].typeParameters ? signatures[0] :
@@ -19718,7 +19713,7 @@ namespace ts {
             // will be a subtype or the same type as the argument.
             function narrowType(type: Type, expr: Expression, assumeTrue: boolean): Type {
                 // for `a?.b`, we emulate a synthetic `a !== null && a !== undefined` condition for `a`
-                if (isOptionalChainRoot(expr.parent) ||
+                if (isOptionalExpression(expr) ||
                     isBinaryExpression(expr.parent) && expr.parent.operatorToken.kind === SyntaxKind.QuestionQuestionToken && expr.parent.left === expr) {
                     return narrowTypeByOptionality(type, expr, assumeTrue);
                 }
@@ -22687,7 +22682,8 @@ namespace ts {
         }
 
         function checkPropertyAccessExpressionOrQualifiedName(node: PropertyAccessExpression | QualifiedName, left: Expression | QualifiedName, right: Identifier) {
-            const { isOptional, type: leftType } = checkOptionalExpression(node, left);
+            const optionalLeftType = checkOptionalExpression(left);
+            const leftType = removeOptionalTypeMarker(optionalLeftType);
             const parentSymbol = getNodeLinks(left).resolvedSymbol;
             const assignmentKind = getAssignmentTargetKind(node);
             const apparentType = getApparentType(assignmentKind !== AssignmentKind.None || isMethodAccessForCall(node) ? getWidenedType(leftType) : leftType);
@@ -22741,7 +22737,7 @@ namespace ts {
                 }
                 propType = getConstraintForLocation(getTypeOfSymbol(prop), node);
             }
-            return propagateOptionalTypeMarker(getFlowTypeOfAccessExpression(node, prop, propType, right), isOptional);
+            return propagateOptionalTypeMarker(getFlowTypeOfAccessExpression(node, prop, propType, right), leftType !== optionalLeftType);
         }
 
         function getFlowTypeOfAccessExpression(node: ElementAccessExpression | PropertyAccessExpression | QualifiedName, prop: Symbol | undefined, propType: Type, errorNode: Node) {
@@ -23098,7 +23094,8 @@ namespace ts {
         }
 
         function checkIndexedAccess(node: ElementAccessExpression): Type {
-            const { isOptional, type: exprType } = checkOptionalExpression(node, node.expression);
+            const optionalExprType = checkOptionalExpression(node.expression);
+            const exprType = removeOptionalTypeMarker(optionalExprType);
             const objectType = getAssignmentTargetKind(node) !== AssignmentKind.None || isMethodAccessForCall(node) ? getWidenedType(exprType) : exprType;
             const indexExpression = node.argumentExpression;
             const indexType = checkExpression(indexExpression);
@@ -23117,7 +23114,7 @@ namespace ts {
                 AccessFlags.Writing | (isGenericObjectType(objectType) && !isThisTypeParameter(objectType) ? AccessFlags.NoIndexSignatures : 0) :
                 AccessFlags.None;
             const indexedAccessType = getIndexedAccessTypeOrUndefined(objectType, effectiveIndexType, node, accessFlags) || errorType;
-            return propagateOptionalTypeMarker(checkIndexedAccessIndexType(getFlowTypeOfAccessExpression(node, indexedAccessType.symbol, indexedAccessType, indexExpression), node), isOptional);
+            return propagateOptionalTypeMarker(checkIndexedAccessIndexType(getFlowTypeOfAccessExpression(node, indexedAccessType.symbol, indexedAccessType, indexExpression), node), exprType !== optionalExprType);
         }
 
         function checkThatExpressionIsProperSymbolReference(expression: Expression, expressionType: Type, reportError: boolean): boolean {
@@ -24304,13 +24301,13 @@ namespace ts {
                 return resolveUntypedCall(node);
             }
 
-            const { isOptional, type: funcType } = checkOptionalExpression(
-                node,
+            const optionalFuncType = checkOptionalExpression(
                 node.expression,
                 Diagnostics.Cannot_invoke_an_object_which_is_possibly_null,
                 Diagnostics.Cannot_invoke_an_object_which_is_possibly_undefined,
                 Diagnostics.Cannot_invoke_an_object_which_is_possibly_null_or_undefined
             );
+            const funcType = removeOptionalTypeMarker(optionalFuncType);
 
             if (funcType === silentNeverType) {
                 return silentNeverSignature;
@@ -24381,7 +24378,7 @@ namespace ts {
                 return resolveErrorCall(node);
             }
 
-            return resolveCall(node, callSignatures, candidatesOutArray, checkMode, isOptional);
+            return resolveCall(node, callSignatures, candidatesOutArray, checkMode, funcType !== optionalFuncType);
         }
 
         function isGenericFunctionReturningFunction(signature: Signature) {
@@ -27276,10 +27273,11 @@ namespace ts {
             // Optimize for the common case of a call to a function with a single non-generic call
             // signature where we can just fetch the return type without checking the arguments.
             if (isCallExpression(expr) && expr.expression.kind !== SyntaxKind.SuperKeyword && !isRequireCall(expr, /*checkArgumentIsStringLiteralLike*/ true) && !isSymbolOrSymbolForCall(expr)) {
-                const { isOptional, type: funcType } = checkOptionalExpression(expr, expr.expression);
+                const optionalFuncType = checkOptionalExpression(expr.expression);
+                const funcType = removeOptionalTypeMarker(optionalFuncType);
                 const signature = getSingleCallSignature(funcType);
                 if (signature && !signature.typeParameters) {
-                    return propagateOptionalTypeMarker(getReturnTypeOfSignature(signature), isOptional);
+                    return propagateOptionalTypeMarker(getReturnTypeOfSignature(signature), funcType !== optionalFuncType);
                 }
             }
             else if (isAssertionExpression(expr) && !isConstTypeReference(expr.type)) {

src/compiler/types.ts

@@ -4265,6 +4265,10 @@ namespace ts {
         immediateBaseConstraint?: Type;  // Immediate base constraint cache
         /* @internal */
         widened?: Type; // Cached widened form of the type
+        /* @internal */
+        optionalTypeOfType?: Type;       // A cached copy of the type unioned with the optionalType marker.
+        /* @internal */
+        nonOptionalTypeOfType?: Type;    // A cached copy of the type with the optionalType marker removed.
     }
 
     /* @internal */

src/compiler/utilities.ts

@@ -5913,6 +5913,14 @@ namespace ts {
                 || kind === SyntaxKind.CallExpression);
     }
 
+    /**
+     * Determines whether a node is the expression preceding an optional chain (i.e. `a` in `a?.b`).
+     */
+    /* @internal */
+    export function isOptionalExpression(node: Node): node is Expression & { parent: OptionalChainRoot } {
+        return isOptionalChainRoot(node.parent) && node === node.parent.expression;
+    }
+
     export function isNewExpression(node: Node): node is NewExpression {
         return node.kind === SyntaxKind.NewExpression;
     }