Infer from context sensitive return expressions

src/compiler/checker.ts

@@ -38938,7 +38938,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
         return getTypeOfSymbol(getSymbolOfDeclaration(node));
     }
 
-    function contextuallyCheckFunctionExpressionOrObjectLiteralMethod(node: FunctionExpression | ArrowFunction | MethodDeclaration, checkMode?: CheckMode) {
+    function contextuallyCheckFunctionExpressionOrObjectLiteralMethod(node: FunctionExpression | ArrowFunction | MethodDeclaration, checkMode = CheckMode.Normal) {
         const links = getNodeLinks(node);
         // Check if function expression is contextually typed and assign parameter types if so.
         if (!(links.flags & NodeCheckFlags.ContextChecked)) {
@@ -38956,7 +38956,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                     if (contextualSignature) {
                         const inferenceContext = getInferenceContext(node);
                         let instantiatedContextualSignature: Signature | undefined;
-                        if (checkMode && checkMode & CheckMode.Inferential) {
+                        if (checkMode & CheckMode.Inferential) {
                             inferFromAnnotatedParameters(signature, contextualSignature, inferenceContext!);
                             const restType = getEffectiveRestType(contextualSignature);
                             if (restType && restType.flags & TypeFlags.TypeParameter) {
@@ -38974,12 +38974,26 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                 }
                 else if (contextualSignature && !node.typeParameters && contextualSignature.parameters.length > node.parameters.length) {
                     const inferenceContext = getInferenceContext(node);
-                    if (checkMode && checkMode & CheckMode.Inferential) {
+                    if (checkMode & CheckMode.Inferential) {
                         inferFromAnnotatedParameters(signature, contextualSignature, inferenceContext!);
                     }
                 }
                 if (contextualSignature && !getReturnTypeFromAnnotation(node) && !signature.resolvedReturnType) {
-                    const returnType = getReturnTypeFromBody(node, checkMode);
+                    let contextualReturnType: Type;
+                    let returnType: Type;
+
+                    if (checkMode & CheckMode.Inferential && couldContainTypeVariables(contextualReturnType = getReturnTypeOfSignature(contextualSignature))) {
+                        const inferenceContext = getInferenceContext(node);
+                        const isReturnContextSensitive = !!node.body && (node.body.kind === SyntaxKind.Block ? forEachReturnStatement(node.body as Block, statement => !!statement.expression && isContextSensitive(statement.expression)) : isContextSensitive(node.body));
+                        returnType = getReturnTypeFromBody(node, checkMode | (isReturnContextSensitive ? CheckMode.SkipContextSensitive : 0));
+                        inferTypes(inferenceContext!.inferences, returnType, contextualReturnType);
+                        if (isReturnContextSensitive) {
+                            returnType = getReturnTypeFromBody(node, checkMode);
+                        }
+                    } else {
+                        returnType = getReturnTypeFromBody(node, checkMode);
+                    }
+
                     if (!signature.resolvedReturnType) {
                         signature.resolvedReturnType = returnType;
                     }

tests/baselines/reference/intraExpressionInferencesInContextSensitive1.js

@@ -0,0 +1,145 @@
+//// [tests/cases/conformance/types/typeRelationships/typeInference/intraExpressionInferencesInContextSensitive1.ts] ////
+
+//// [intraExpressionInferencesInContextSensitive1.ts]
+// https://github.com/microsoft/TypeScript/issues/60720
+
+type Options<TContext> = {
+  onStart?: () => TContext;
+  onEnd?: (context: TContext) => void;
+};
+
+function create<TContext>(builder: (arg: boolean) => Options<TContext>) {
+  return builder(true);
+}
+
+create((arg: boolean) => ({
+  onStart: () => ({ time: new Date() }),
+  onEnd: (context) => {},
+}));
+
+create(() => ({
+  onStart: () => ({ time: new Date() }),
+  onEnd: (context) => {},
+}));
+
+create((arg) => ({
+  onStart: () => ({ time: new Date() }),
+  onEnd: (context) => {},
+}));
+
+// https://github.com/microsoft/TypeScript/issues/57021
+
+type Schema = Record<string, unknown>;
+
+type StepFunction<TSchema extends Schema = Schema> = (anything: unknown) => {
+  readonly schema: TSchema;
+  readonly toAnswers?: (keys: keyof TSchema) => unknown;
+};
+
+function step1<TSchema extends Schema = Schema>(
+  stepVal: StepFunction<TSchema>,
+): StepFunction<TSchema> {
+  return stepVal;
+}
+
+const stepResult1 = step1((_something) => ({
+  schema: {
+    attribute: "anything",
+  },
+  toAnswers: (keys) => {
+    type Test = string extends typeof keys ? never : "true";
+    const test: Test = "true"; // ok
+    return { test };
+  },
+}));
+
+type StepFunction2<TSchema extends Schema = Schema> = (anything: unknown) => {
+  readonly schema: (thing: number) => TSchema;
+  readonly toAnswers?: (keys: keyof TSchema) => unknown;
+};
+
+function step2<TSchema extends Schema = Schema>(
+  stepVal: StepFunction2<TSchema>,
+): StepFunction2<TSchema> {
+  return stepVal;
+}
+
+const stepResult2 = step2((_something) => ({
+  schema: (thing) => ({
+    attribute: "anything",
+  }),
+  toAnswers: (keys) => {
+    type Test = string extends typeof keys ? never : "true";
+    const test: Test = "true"; // ok
+    return { test };
+  },
+}));
+
+
+//// [intraExpressionInferencesInContextSensitive1.js]
+"use strict";
+// https://github.com/microsoft/TypeScript/issues/60720
+function create(builder) {
+    return builder(true);
+}
+create(function (arg) { return ({
+    onStart: function () { return ({ time: new Date() }); },
+    onEnd: function (context) { },
+}); });
+create(function () { return ({
+    onStart: function () { return ({ time: new Date() }); },
+    onEnd: function (context) { },
+}); });
+create(function (arg) { return ({
+    onStart: function () { return ({ time: new Date() }); },
+    onEnd: function (context) { },
+}); });
+function step1(stepVal) {
+    return stepVal;
+}
+var stepResult1 = step1(function (_something) { return ({
+    schema: {
+        attribute: "anything",
+    },
+    toAnswers: function (keys) {
+        var test = "true"; // ok
+        return { test: test };
+    },
+}); });
+function step2(stepVal) {
+    return stepVal;
+}
+var stepResult2 = step2(function (_something) { return ({
+    schema: function (thing) { return ({
+        attribute: "anything",
+    }); },
+    toAnswers: function (keys) {
+        var test = "true"; // ok
+        return { test: test };
+    },
+}); });
+
+
+//// [intraExpressionInferencesInContextSensitive1.d.ts]
+type Options<TContext> = {
+    onStart?: () => TContext;
+    onEnd?: (context: TContext) => void;
+};
+declare function create<TContext>(builder: (arg: boolean) => Options<TContext>): Options<TContext>;
+type Schema = Record<string, unknown>;
+type StepFunction<TSchema extends Schema = Schema> = (anything: unknown) => {
+    readonly schema: TSchema;
+    readonly toAnswers?: (keys: keyof TSchema) => unknown;
+};
+declare function step1<TSchema extends Schema = Schema>(stepVal: StepFunction<TSchema>): StepFunction<TSchema>;
+declare const stepResult1: StepFunction<{
+    attribute: string;
+}>;
+type StepFunction2<TSchema extends Schema = Schema> = (anything: unknown) => {
+    readonly schema: (thing: number) => TSchema;
+    readonly toAnswers?: (keys: keyof TSchema) => unknown;
+};
+declare function step2<TSchema extends Schema = Schema>(stepVal: StepFunction2<TSchema>): StepFunction2<TSchema>;
+declare const stepResult2: StepFunction2<{
+    attribute: string;
+}>;