Narrow generic conditional and indexed access return types when checking return statements

DRANK

Fixes #33912. Fixes #33014. Motivation Sometimes we want to write functions whose return type is picked between different options, depending on the type of a parameter. For instance: declare const record: Record<string, string>; declare const array: string[]; function getObject(group) { if (group === undefined) { return record; } return array; } const arrayResult = getObject("group"); const recordResult = getObject(undefined); If we want to precisely express this dependency between the return type and the type of nameOrId, we have a few options. The first one is to use overloads: declare const record: Record<string, string[]>; declare const array: string[]; function getObject(group: undefined): Record<string, string[]>; function getObject(group: string): string[]; function getObject(group: string | undefined): string[] | Record<string, string[]>; function getObject(group: string | undefined): string[] | Record<string, string[]> { if (group === undefined) { return record; } return array; } const arrayResult = getObject("group"); const recordResult = getObject(undefined); However, if you make a mistake in the implementation of the function and return the wrong type, TypeScript will not warn you. For instance, if instead you implement the function like this: declare const record: Record<string, string[]>; declare const array: string[]; function getObject(group: undefined): Record<string, string[]>; function getObject(group: string): string[]; function getObject(group: string | undefined): string[] | Record<string, string[]>; function getObject(group: string | undefined): string[] | Record<string, string[]> { if (!group) { // An empty string is falsy return record; } return array; } const badResult = getObject(""); // Type says this returns `string[]`, but actually it returns a record. then your function implementation doesn't respect the overload signatures, but TypeScript will not error. The alternative to overloads is to use conditional types, like so: declare const record: Record<string, string[]>; declare const array: string[]; function getObject<T extends string | undefined>(group: T): T extends string ? string[] : T extends undefined ? Record<string, string[]> : never { if (group === undefined) { return record; // Error! Type 'Record<string, string[]>' is not assignable to type 'T extends string ? string[] : T extends undefined ? Record<string, string[]> : never'. } return array; // Error! Type 'string[]' is not assignable to type 'T extends string ? string[] : T extends undefined ? Record<string, string[]> : never' } const arrayResult = getObject("group"); const recordResult = getObject(undefined); However, while everything works out for the callers of getObject, in the implementation TypeScript errors on the return statements, because it compares the type of the return expression to the annotated conditional return type, and Record<string, string[]> is not assignable to T extends undefined ? Record<string, string[]> : never. Solution: conditional return type narrowing For this PR, I propose a way of checking return statements that understands cases like above. The idea is that, for the function above: function getObject<T extends string | undefined>(group: T): T extends string ? string[] : T extends undefined ? Record<string, string[]> : never { if (group === undefined) { return record; } return array; } when checking the return statement return record, TS will know that group has type undefined. TS will also know that type parameter T corresponds exactly to the type of group. Combining those two pieces of information, TS will know that, inside that branch, the expected return type has to be Record<string, string[]> (or a supertype). Then, instead of checking return record's type against T extends string ? string[] : T extends undefined ? Record<string, string[]> : never, it will check return record's type against the Record<string, string[]> branch of the conditional, i.e. a narrowed version of the conditional type. Then there will be no error on the return statement. In the same manner, when checking the return statement return array, TS will know that group has type string, and that therefore it can check the type of the return statement against the string[] branch of the conditional return type. For now, we can think of it like this: when we check return statement return record, we see that group has narrowed type undefined. Then, we plug that information back into the return type by instantiating T extends string ? string[] : T extends undefined ? Record<string, string[]> : never with T replaced by undefined. Restrictions Conditional types Reasoning about conditional types is pretty tricky. In general, given a function f(...args): SomeConditionalType whose return type is some (generic) conditional type, we are not able to do the special check proposed above, because it wouldn't be safe. We need to place some restrictions on what the conditional return type looks like in order for TS to safely analyze it. We can safely analyze a conditional return type that has the shape T extends A ? AType : T extends B ? BType : never. This means the conditional type needs to be distributive (i.e. its check type is a naked type parameter T) and have never as its false-most type, and it cannot have infer type parameters (e.g. it cannot be T extends [infer A] ? AType : T extends B ? BType : never). Intuitively, we can think of this conditional type shape as reflecting the kind of code one would write in the implementation of such a function. In addition to the previous restrictions on the conditional type, the extends types of the conditional (A and B above), have to be constituents of the type parameter's constraint (T above), like so: function fun<T extends A | B>(param: T): T extends A ? AType : T extends B ? BType : never { if (isA(param)) { ... } else { ... } } This is because, to narrow the return type, we first need to narrow the type of param (more on that below). When we narrow the type of param, in a lot of scenarios, we will start from its type, T, or in this case its type constraint, A | B. Then, we will further narrow that type based on information from control flow analysis, e.g. to pick either A or B (see getNarrowableTypeForReference in checker.ts). Therefore, in this typical case, narrowing param means we will end up with a type that is either A, B, or a subtype of those. In turn, when we plug this narrowed type back into the conditional return type, this means we will be able to pick a branch of the conditional type and resolve it. e.g. if the narrowed type of param is A, the conditional type will resolve to AType. This restriction also reinforces that you want your return type to closely match the structure of your implementation. Aside: why never A common way of trying to write a conditional return type is like the following: function stringOrNumber<T extends string | number>(param: T): T extends string ? string : number { if (typeof param === "string") { return "some string"; } return 123; } const num = stringOrNumber(123); const str = stringOrNumber("string"); declare let strOrNum: string | number; const both = stringOrNumber(strOrNum); This example works fine and it would be safe for TS to allow that function implementation. However, in general, it is not safe to allow this pattern of conditional return type. Consider this case: function aStringOrANumber<T extends { a: string } | { a: undefined }>(param: T): T extends { a: string } ? string : number { if (typeof param.a === "string") { return "some string"; } return 123; } const aNum = aStringOrANumber({ a: undefined }); const aStr = aStringOrANumber({ a: "" }); // Bad! Type says `number`, but actually should say `string | number` const aNotBoth = aStringOrANumber({ a: strOrUndef }); The problem boils down to the fact that, when a conditional return type resolves to its false branch, we can't know if the check type is related or not to the extends type. For the example above, when we plug in { a: undefined } for T in T extends { a: string } ? string : number, then we fall into the false branch of the conditional, which is desired because { a: undefined } does not overlap { a: string }. However, when we plug in { a: string | undefined } for T in T extends { a: string } ? string : number, we fall into the false branch of the conditional, but this is not desired because { a: string | undefined } overlaps { a: string }, and therefore the return type could actually be string. Resolving a conditional type to its false-most branch of a conditional type doesn't provide TS with enough information to safely determine what the return type should be, and because of that, narrowing a conditional return type requires the false-most branch of the conditional to be never. Type parameter references As hinted at above, to narrow a conditional return type, we first need to narrow a parameter, and we need to know that the type of that parameter uniquely corresponds to a type parameter. Revisiting our example: function getObject<T extends string | undefined>(group: T): T extends string ? string[] : T extends undefined ? Record<string, string[]> : never { if (group === undefined) { return record; } return array; } To narrow the return type T extends string ? string[] : T extends undefined ? Record<string, string[]> : never, we first narrow the type of group inside the if branch, and then we can use that information to reason about what type T could be replaced with. This only works because the declared type of group is exactly T, and also because there are no other parameters that use type T. So in the following cases, TS would not be able to narrow the return type, because there is no unique parameter to which T is linked: function badGetObject1<T extends string | undefined>(group: T, someOtherParam: T): T extends string ? string[] : T extends undefined ? Record<string, string[]> : never { ... } function badGetObject2<T extends string | undefined>(group: T, options: { a: number, b: T }): T extends string ? string[] : T extends undefined ? Record<string, string[]> : never { ... } Indexed access types The reasoning explained above for conditional types applies in a similar manner to indexed access types that look like this: interface F { "t": number, "f": boolean, } function depLikeFun<T extends "t" | "f">(str: T): F[T] { if (str === "t") { return 1; } else { return true; } } depLikeFun("t"); // has type number depLikeFun("f"); // has type boolean So cases like this, where the return type is an indexed access type with a type parameter index, are also supported by this PR. The thinking is similar: in the code above, when we are in the if (str === "t") { ... } branch, we know str has type "t", and we can plug that information back into the return type F[T] which resolves to type number, and similarly for the else branch. Implementation The implementation works roughly like this: When checking a return statement expression: We check if the type of the expression is assignable to the return type. If it is, good, nothing else needs to be done. This makes sure we don't introduce new errors to existing code. If the type of the expression is not assignable to the return type, we might need to narrow the return type and check again. We proceed to attempt narrowing the return type. We check if the return type has the right shape, i.e. it has a shape of either SomeType[T] or T extends A ? AType : T extends B ? BType : never, and what parameters the type parameters are uniquely linked to, among other requirements. If any of those requirements is not met, we don't continue with narrowing. For every type parameter that is uniquely linked to a parameter, we obtain its narrowed type. Say we have type parameter T that is uniquely linked to parameter param. To narrow the return type, we first need to obtain the narrowed type for param at the return statement position. Because, in the source code, there might not be an occurrence of param at the return statement position, we create a synthetic reference to param at that position and obtain its narrowed type via regular control-flow analysis. We then obtain a narrowed type N for param. (If we don't, we'll just ignore that type parameter). Once we have the narrowed type for each type parameter, we need to plug that information back into the return type. We do this by instantiating the return type with a substitution type. For instance, if the return type is T extends A ? AType : T extends B ? BType : never, T is linked to param, and param has narrowed type N, we will instantiate the return type with T replaced by T & N (as a substitution type). Once we obtain this narrowed return type, we get the type of the return expression, this time contextually checked by the narrowed return type, and then check if this type is assignable to the narrowed return type. Instantiation As mentioned above, the process of narrowing a return type is implemented as instantiating that return type with the narrowed type parameters replaced by substitution types. Substitution types can be thought of as T & A, where T is the base type and A the constraint type. There are a few changes made to instantiation to make this work: There is now a new kind of substitution type, indicated by the ObjectFlags.IsNarrowingType flag, which corresponds to substitution types created by return type narrowing. These narrowing substitution types are handled in a special way by conditional type instantiation functions, i.e. getConditionalTypeInstantiation, and getConditionalType. getConditionalTypeInstantiation is responsible for distributing a conditional type over its check type. When instantiating a distributive conditional type in getConditionalTypeInstantiation, if the conditional's check type is a substitution type like T & (A | B), the usual logic would not distribute over this type, because it's a substitution type and not a union type. So, for distribution to happen, we have to take apart the T & (A | B) into (T & A) | (T & B), and distribute over that. The other special thing that we have to do in getConditionalTypeInstantiation is to take the intersection of the distribution result, as opposed to the union. This is because, if we narrow a type parameter T to A | B, and we have a conditional return type T extends A ? R1 : T extends B ? R2 : T extends C ? R3 : never, then we don't know which branch of the conditional return to pick, if branch T extends A ? R1, or branch T extends B ? R2, so we have to check whether the return expression's type is assignable to both, i.e. assignable to R1 & R2. Validating whether the conditional type has the right shape to be narrowed happens on-demand in getConditionalTypeInstantiation (and also at first as an optimization in checkReturnExpression when we're deciding whether to narrow the return type). This is because, as we instantiate a type, we may produce new conditional types that we need to then decide are safe to narrow or not (see nested types case in dependentReturnType6.ts test). Conditional expression checking To support conditional expression checking in return statements, this PR changes how we check a conditional expression in a return statement. Before this PR, when checking a return statement return cond ? exp1 : exp2, we obtain the type of the whole expression cond ? exp1 : exp2, and then compare that type to the return type. With this PR, we now separately check each branch: we first obtain the type of exp1, and compare that type to the return type, then obtain the type of exp2 and compare that type to the return type. This allows us to properly check a conditional expression when return type narrowing is needed. This a breaking change, and the only change that affects existing code, but this change finds bugs. Analysis of extended tests changes: https://github.com/microsoft/TypeScript/pull/56941#issuecomment-2406411716. This change also slightly affects performance because we do more checks. Latest perf results here: https://github.com/microsoft/TypeScript/pull/56941#issuecomment-2403351320. Performance results This feature is opt-in. Currently, virtually no code has functions whose return types are conditional or indexed access types that satisfy the restrictions above, so no code goes down the new code path for narrowing return types. This means for existing code, there's no performance impact from the return type narrowing. The only current performance impact is from the conditional expression checking change (see a version of this PR without the change for conditional expressions: https://github.com/microsoft/TypeScript/pull/60268#issuecomment-2421064260). Assessing the performance of the new code path is tricky, as there are no baselines. The existing alternative to conditional return types is to use overloads. In one scenario I tested, checking a function written with conditional return types with this PR took ~+16% check time compared to the same function using overloads and the main branch. However, that's for checking the function declaration. In a different scenario where I included a lot of function calls though, the version with conditional return types + this PR took ~-15% compared to overloads + main branch. So I'd say the performance is acceptable, especially considering you get stronger checks when using conditional return types, and also only a small number of functions in a codebase should be written using this feature. Unsupported things Inference: TS will not infer a conditional return type or an indexed access type for any function or expression. Inferring such a type is more complicated than checking, and inferring such a type could also be surprising for users. This is out of scope. Contextually-typed anonymous functions: type GetObjectCallback = <T extends string | undefined>(group: T) => T extends string ? string[] : T extends undefined ? Record<string, string[]> : never; const getObjectBad1: GetObjectCallback = (group) => { return group === undefined ? record : array }; // Error declare function outerFun(callback: GetObjectCallback); outerFun((group) => { return group === undefined ? record : array }); // Error This is because, if your function does not have an explicitly annotated return type, we will infer one from the returns. Detection of link between parameter and type parameter in more complicated scenarios: // All cases below are not recognized // Type alias type Id<X> = X; function f2<T extends boolean>(arg: Id<T>): T extends true ? string : T extends false ? number : never { if (arg) { return "someString"; } return 123; } // Property type function f3<T extends boolean>(arg: { prop: T }): T extends true ? string : T extends false ? number : never { if (arg.prop) { return "someString"; } return 123; } // Destructuring function f4<T extends boolean>({ arg }: { arg: T }): T extends true ? string : T extends false ? number : never { if (arg.prop) { return "someString"; } return 123; } // Combinations of the above, e.g.: type Opts<X> = { prop: X }; function f5<T extends boolean>(arg: Opts<T>): T extends true ? string : T extends false ? number : never { if (arg.prop) { return "someString"; } return 123; } This could be supported in the future.