Infer type predicates from function bodies

Author: danvk

src/compiler/checker.ts

@@ -15458,9 +15458,17 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                         jsdocPredicate = getTypePredicateOfSignature(jsdocSignature);
                     }
                 }
-                signature.resolvedTypePredicate = type && isTypePredicateNode(type) ?
-                    createTypePredicateFromTypePredicateNode(type, signature) :
-                    jsdocPredicate || noTypePredicate;
+                if (type || jsdocPredicate) {
+                    signature.resolvedTypePredicate = type && isTypePredicateNode(type) ?
+                        createTypePredicateFromTypePredicateNode(type, signature) :
+                        jsdocPredicate || noTypePredicate;
+                } else if (signature.declaration && isFunctionLikeDeclaration(signature.declaration) && (!signature.resolvedReturnType || signature.resolvedReturnType === booleanType)) {
+                    const {declaration} = signature;
+                    signature.resolvedTypePredicate = noTypePredicate;  // avoid infinite loop
+                    signature.resolvedTypePredicate = getTypePredicateFromBody(declaration, signature) || noTypePredicate;
+                } else {
+                    signature.resolvedTypePredicate = noTypePredicate;
+                }
             }
             Debug.assert(!!signature.resolvedTypePredicate);
         }
@@ -37389,6 +37397,80 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
         }
     }
 
+    function getTypePredicateFromBody(func: FunctionLikeDeclaration, _sig: Signature): TypePredicate | undefined {
+        const functionFlags = getFunctionFlags(func);
+        if (functionFlags !== FunctionFlags.Normal) return undefined;
+
+        // Only attempt to infer a type predicate if there's exactly one return.
+        let singleReturn: Expression | undefined;
+        if (func.body && func.body.kind !== SyntaxKind.Block) {
+            singleReturn = func.body; // arrow function
+        } else {
+            if (functionHasImplicitReturn(func)) return undefined;
+
+            const bailedEarly = forEachReturnStatement(func.body as Block, returnStatement => {
+                if (singleReturn || !returnStatement.expression) return true;
+                singleReturn = returnStatement.expression;
+            });
+            if (bailedEarly || !singleReturn) return undefined;
+        }
+
+        const predicate = checkIfExpressionRefinesAnyParameter(singleReturn);
+        if (predicate) {
+            const [i, type] = predicate;
+            const param = func.parameters[i];
+            if (isIdentifier(param.name)) {
+                // TODO: is there an alternative to the "as string" here? (It's __String)
+                return createTypePredicate(TypePredicateKind.Identifier, param.name.escapedText as string, i, type);
+            }
+        }
+        return undefined;
+
+        function checkIfExpressionRefinesAnyParameter(expr: Expression): [number, Type] | undefined {
+            expr = skipParentheses(expr, /*excludeJSDocTypeAssertions*/ true);
+            const type = checkExpressionCached(expr, CheckMode.TypeOnly);
+            if (type !== booleanType || !func.body) return undefined;
+
+            return forEach(func.parameters, (param, i) => {
+                const initType = getSymbolLinks(param.symbol).type;
+                if (!initType || initType === booleanType || isSymbolAssigned(param.symbol)) {
+                    // Refining "x: boolean" to "x is true" or "x is false" isn't useful.
+                    return;
+                }
+                const trueType = checkIfExpressionRefinesParameter(expr, param, initType);
+                if (trueType) {
+                    return [i, trueType];
+                }
+            });
+        }
+
+        function checkIfExpressionRefinesParameter(expr: Expression, param: ParameterDeclaration, initType: Type): Type | undefined {
+            const antecedent = (expr as Expression & {flowNode?: FlowNode}).flowNode ?? { flags: FlowFlags.Start };
+            const trueCondition: FlowCondition = {
+                flags: FlowFlags.TrueCondition,
+                node: expr,
+                antecedent,
+            };
+
+            const trueType = getFlowTypeOfReference(param.name, initType, initType, func, trueCondition);
+            if (trueType === initType) return undefined;
+
+            // "x is T" means that x is T if and only if it returns true. If it returns false then x is not T.
+            // However, TS may not be able to represent "not T", in which case we can be more lax.
+            // It's safe to infer a type guard if falseType = Exclude<initType, trueType>
+            // This matches what you'd get if you called the type guard in an if/else statement.
+            const falseCondition: FlowCondition = {
+                ...trueCondition,
+                flags: FlowFlags.FalseCondition,
+            }
+            const falseType = getFlowTypeOfReference(param.name, initType, initType, func, falseCondition);
+            const candidateFalse = filterType(initType, t => !isTypeSubtypeOf(t, trueType));
+            if (isTypeIdenticalTo(candidateFalse, falseType)) {
+                return trueType;
+            }
+        }
+    }
+
     /**
      * TypeScript Specification 1.0 (6.3) - July 2014
      *   An explicitly typed function whose return type isn't the Void type,

tests/baselines/reference/findLast(target=esnext).types

@@ -3,133 +3,133 @@
 === findLast.ts ===
 const itemNumber: number | undefined = [0].findLast((item) => item === 0);
 >itemNumber : number
->[0].findLast((item) => item === 0) : number
+>[0].findLast((item) => item === 0) : 0
 >[0].findLast : { <S extends number>(predicate: (value: number, index: number, array: number[]) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: number[]) => unknown, thisArg?: any): number; }
 >[0] : number[]
 >0 : 0
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: number[]) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: number[]) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 const itemString: string | undefined = ["string"].findLast((item) => item === "string");
 >itemString : string
->["string"].findLast((item) => item === "string") : string
+>["string"].findLast((item) => item === "string") : "string"
 >["string"].findLast : { <S extends string>(predicate: (value: string, index: number, array: string[]) => value is S, thisArg?: any): S; (predicate: (value: string, index: number, array: string[]) => unknown, thisArg?: any): string; }
 >["string"] : string[]
 >"string" : "string"
 >findLast : { <S extends string>(predicate: (value: string, index: number, array: string[]) => value is S, thisArg?: any): S; (predicate: (value: string, index: number, array: string[]) => unknown, thisArg?: any): string; }
->(item) => item === "string" : (item: string) => boolean
+>(item) => item === "string" : (item: string) => item is "string"
 >item : string
 >item === "string" : boolean
 >item : string
 >"string" : "string"
 
 new Int8Array().findLast((item) => item === 0);
->new Int8Array().findLast((item) => item === 0) : number
+>new Int8Array().findLast((item) => item === 0) : 0
 >new Int8Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Int8Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Int8Array) => unknown, thisArg?: any): number; }
 >new Int8Array() : Int8Array
 >Int8Array : Int8ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Int8Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Int8Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Uint8Array().findLast((item) => item === 0);
->new Uint8Array().findLast((item) => item === 0) : number
+>new Uint8Array().findLast((item) => item === 0) : 0
 >new Uint8Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint8Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint8Array) => unknown, thisArg?: any): number; }
 >new Uint8Array() : Uint8Array
 >Uint8Array : Uint8ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint8Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint8Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Uint8ClampedArray().findLast((item) => item === 0);
->new Uint8ClampedArray().findLast((item) => item === 0) : number
+>new Uint8ClampedArray().findLast((item) => item === 0) : 0
 >new Uint8ClampedArray().findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint8ClampedArray) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint8ClampedArray) => unknown, thisArg?: any): number; }
 >new Uint8ClampedArray() : Uint8ClampedArray
 >Uint8ClampedArray : Uint8ClampedArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint8ClampedArray) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint8ClampedArray) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Int16Array().findLast((item) => item === 0);
->new Int16Array().findLast((item) => item === 0) : number
+>new Int16Array().findLast((item) => item === 0) : 0
 >new Int16Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Int16Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Int16Array) => unknown, thisArg?: any): number; }
 >new Int16Array() : Int16Array
 >Int16Array : Int16ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Int16Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Int16Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Uint16Array().findLast((item) => item === 0);
->new Uint16Array().findLast((item) => item === 0) : number
+>new Uint16Array().findLast((item) => item === 0) : 0
 >new Uint16Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint16Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint16Array) => unknown, thisArg?: any): number; }
 >new Uint16Array() : Uint16Array
 >Uint16Array : Uint16ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint16Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint16Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Int32Array().findLast((item) => item === 0);
->new Int32Array().findLast((item) => item === 0) : number
+>new Int32Array().findLast((item) => item === 0) : 0
 >new Int32Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Int32Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Int32Array) => unknown, thisArg?: any): number; }
 >new Int32Array() : Int32Array
 >Int32Array : Int32ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Int32Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Int32Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Uint32Array().findLast((item) => item === 0);
->new Uint32Array().findLast((item) => item === 0) : number
+>new Uint32Array().findLast((item) => item === 0) : 0
 >new Uint32Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint32Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint32Array) => unknown, thisArg?: any): number; }
 >new Uint32Array() : Uint32Array
 >Uint32Array : Uint32ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Uint32Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Uint32Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Float32Array().findLast((item) => item === 0);
->new Float32Array().findLast((item) => item === 0) : number
+>new Float32Array().findLast((item) => item === 0) : 0
 >new Float32Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Float32Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Float32Array) => unknown, thisArg?: any): number; }
 >new Float32Array() : Float32Array
 >Float32Array : Float32ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Float32Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Float32Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
 >0 : 0
 
 new Float64Array().findLast((item) => item === 0);
->new Float64Array().findLast((item) => item === 0) : number
+>new Float64Array().findLast((item) => item === 0) : 0
 >new Float64Array().findLast : { <S extends number>(predicate: (value: number, index: number, array: Float64Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Float64Array) => unknown, thisArg?: any): number; }
 >new Float64Array() : Float64Array
 >Float64Array : Float64ArrayConstructor
 >findLast : { <S extends number>(predicate: (value: number, index: number, array: Float64Array) => value is S, thisArg?: any): S; (predicate: (value: number, index: number, array: Float64Array) => unknown, thisArg?: any): number; }
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -170,7 +170,7 @@ const indexNumber: number = [0].findLastIndex((item) => item === 0);
 >[0] : number[]
 >0 : 0
 >findLastIndex : (predicate: (value: number, index: number, array: number[]) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -183,7 +183,7 @@ const indexString: number = ["string"].findLastIndex((item) => item === "string"
 >["string"] : string[]
 >"string" : "string"
 >findLastIndex : (predicate: (value: string, index: number, array: string[]) => unknown, thisArg?: any) => number
->(item) => item === "string" : (item: string) => boolean
+>(item) => item === "string" : (item: string) => item is "string"
 >item : string
 >item === "string" : boolean
 >item : string
@@ -195,7 +195,7 @@ new Int8Array().findLastIndex((item) => item === 0);
 >new Int8Array() : Int8Array
 >Int8Array : Int8ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Int8Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -207,7 +207,7 @@ new Uint8Array().findLastIndex((item) => item === 0);
 >new Uint8Array() : Uint8Array
 >Uint8Array : Uint8ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Uint8Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -219,7 +219,7 @@ new Uint8ClampedArray().findLastIndex((item) => item === 0);
 >new Uint8ClampedArray() : Uint8ClampedArray
 >Uint8ClampedArray : Uint8ClampedArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Uint8ClampedArray) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -231,7 +231,7 @@ new Int16Array().findLastIndex((item) => item === 0);
 >new Int16Array() : Int16Array
 >Int16Array : Int16ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Int16Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -243,7 +243,7 @@ new Uint16Array().findLastIndex((item) => item === 0);
 >new Uint16Array() : Uint16Array
 >Uint16Array : Uint16ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Uint16Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -255,7 +255,7 @@ new Int32Array().findLastIndex((item) => item === 0);
 >new Int32Array() : Int32Array
 >Int32Array : Int32ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Int32Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -267,7 +267,7 @@ new Uint32Array().findLastIndex((item) => item === 0);
 >new Uint32Array() : Uint32Array
 >Uint32Array : Uint32ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Uint32Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -279,7 +279,7 @@ new Float32Array().findLastIndex((item) => item === 0);
 >new Float32Array() : Float32Array
 >Float32Array : Float32ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Float32Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number
@@ -291,7 +291,7 @@ new Float64Array().findLastIndex((item) => item === 0);
 >new Float64Array() : Float64Array
 >Float64Array : Float64ArrayConstructor
 >findLastIndex : (predicate: (value: number, index: number, array: Float64Array) => unknown, thisArg?: any) => number
->(item) => item === 0 : (item: number) => boolean
+>(item) => item === 0 : (item: number) => item is 0
 >item : number
 >item === 0 : boolean
 >item : number