Don't lose inference session in all the possible locations

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -16969,6 +16969,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
             }
 
+            @Test
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @Test
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
@@ -16987,18 +16993,6 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @Test
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @Test
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
@@ -17058,6 +17052,24 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @Test
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @Test
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @Test
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/extensions/PartialAnalysisHandlerExtension.kt

@@ -73,12 +73,12 @@ open class PartialAnalysisHandlerExtension : AnalysisHandlerExtension {
                     ForceResolveUtil.forceResolveAllContents(descriptor.typeConstructor.supertypes)
 
                     if (declaration is KtClassOrObject && descriptor is ClassDescriptorWithResolutionScopes) {
-                        bodyResolver.resolveSuperTypeEntryList(DataFlowInfo.EMPTY,
-                                                               declaration,
-                                                               descriptor,
-                                                               descriptor.unsubstitutedPrimaryConstructor,
-                                                               descriptor.scopeForConstructorHeaderResolution,
-                                                               descriptor.scopeForMemberDeclarationResolution)
+                        bodyResolver.resolveSuperTypeEntryList(
+                            DataFlowInfo.EMPTY, declaration, descriptor, descriptor.unsubstitutedPrimaryConstructor,
+                            descriptor.scopeForConstructorHeaderResolution,
+                            descriptor.scopeForMemberDeclarationResolution,
+                            resolveSession.inferenceSession
+                        )
                     }
                 }
                 is PropertyDescriptor -> {
@@ -99,7 +99,8 @@ open class PartialAnalysisHandlerExtension : AnalysisHandlerExtension {
                         if (containingScope != null) {
                             bodyResolver.resolveConstructorParameterDefaultValues(
                                 topDownAnalysisContext.outerDataFlowInfo, bindingTrace,
-                                declaration, descriptor as ConstructorDescriptor, containingScope
+                                declaration, descriptor as ConstructorDescriptor, containingScope,
+                                resolveSession.inferenceSession
                             )
                         }
                     } else if (declaration is KtFunction && !declaration.hasDeclaredReturnType() && !declaration.hasBlockBody()) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/AnnotationResolverImpl.java

@@ -163,7 +163,8 @@ public class AnnotationResolverImpl extends AnnotationResolver {
                 CallMaker.makeCall(null, null, annotationEntry),
                 NO_EXPECTED_TYPE,
                 DataFlowInfoFactory.EMPTY,
-                true
+                true,
+                null // specific calls in terms of inference, can't be inside annotation calls
         );
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/BodyResolver.java

@@ -157,7 +157,8 @@ public class BodyResolver {
         resolveFunctionBody(
                 outerDataFlowInfo, trace, constructor, descriptor, declaringScope,
                 headerInnerScope -> resolveSecondaryConstructorDelegationCall(
-                        outerDataFlowInfo, trace, headerInnerScope, constructor, descriptor
+                        outerDataFlowInfo, trace, headerInnerScope, constructor,
+                        descriptor, localContext != null ? localContext.inferenceSession : null
                 ),
                 scope -> new LexicalScopeImpl(
                         scope, descriptor, scope.isOwnerDescriptorAccessibleByLabel(), scope.getImplicitReceiver(),
@@ -173,7 +174,8 @@ public class BodyResolver {
             @NotNull BindingTrace trace,
             @NotNull LexicalScope scope,
             @NotNull KtSecondaryConstructor constructor,
-            @NotNull ClassConstructorDescriptor descriptor
+            @NotNull ClassConstructorDescriptor descriptor,
+            @Nullable InferenceSession inferenceSession
     ) {
         if (descriptor.isExpect() || isEffectivelyExternal(descriptor)) {
             // For expected and external classes, we do not resolve constructor delegation calls because they are prohibited
@@ -182,7 +184,7 @@ public class BodyResolver {
 
         OverloadResolutionResults<?> results = callResolver.resolveConstructorDelegationCall(
                 trace, scope, outerDataFlowInfo,
-                descriptor, constructor.getDelegationCall());
+                descriptor, constructor.getDelegationCall(), inferenceSession);
 
         if (results != null && results.isSingleResult()) {
             ResolvedCall<? extends CallableDescriptor> resolvedCall = results.getResultingCall();
@@ -260,11 +262,13 @@ public class BodyResolver {
         for (Map.Entry<KtClassOrObject, ClassDescriptorWithResolutionScopes> entry : c.getDeclaredClasses().entrySet()) {
             KtClassOrObject classOrObject = entry.getKey();
             ClassDescriptorWithResolutionScopes descriptor = entry.getValue();
+            ExpressionTypingContext localContext = c.getLocalContext();
 
             resolveSuperTypeEntryList(c.getOuterDataFlowInfo(), classOrObject, descriptor,
                                       descriptor.getUnsubstitutedPrimaryConstructor(),
                                       descriptor.getScopeForConstructorHeaderResolution(),
-                                      descriptor.getScopeForMemberDeclarationResolution());
+                                      descriptor.getScopeForMemberDeclarationResolution(),
+                                      localContext != null ? localContext.inferenceSession : null);
         }
     }
 
@@ -274,7 +278,8 @@ public class BodyResolver {
             @NotNull ClassDescriptor descriptor,
             @Nullable ConstructorDescriptor primaryConstructor,
             @NotNull LexicalScope scopeForConstructorResolution,
-            @NotNull LexicalScope scopeForMemberResolution
+            @NotNull LexicalScope scopeForMemberResolution,
+            @Nullable InferenceSession inferenceSession
     ) {
         ProgressManager.checkCanceled();
 
@@ -316,7 +321,8 @@ public class BodyResolver {
                     LexicalScope scope = scopeForConstructor == null ? scopeForMemberResolution : scopeForConstructor;
                     KotlinType expectedType = supertype != null ? supertype : NO_EXPECTED_TYPE;
                     typeInferrer.getType(
-                            scope, delegateExpression, expectedType, outerDataFlowInfo, InferenceSession.Companion.getDefault(), trace
+                            scope, delegateExpression, expectedType, outerDataFlowInfo,
+                            inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault(), trace
                     );
                 }
 
@@ -348,8 +354,9 @@ public class BodyResolver {
                     return;
                 }
                 OverloadResolutionResults<FunctionDescriptor> results = callResolver.resolveFunctionCall(
-                        trace, scopeForConstructor,
-                        CallMaker.makeConstructorCallWithoutTypeArguments(call), NO_EXPECTED_TYPE, outerDataFlowInfo, false);
+                        trace, scopeForConstructor, CallMaker.makeConstructorCallWithoutTypeArguments(call),
+                        NO_EXPECTED_TYPE, outerDataFlowInfo, false, inferenceSession
+                );
                 if (results.isSingleResult()) {
                     KotlinType supertype = results.getResultingDescriptor().getReturnType();
                     recordSupertype(typeReference, supertype);
@@ -404,8 +411,8 @@ public class BodyResolver {
         if (ktClass instanceof KtEnumEntry && DescriptorUtils.isEnumEntry(descriptor) && ktClass.getSuperTypeListEntries().isEmpty()) {
             assert scopeForConstructor != null : "Scope for enum class constructor should be non-null: " + descriptor;
             resolveConstructorCallForEnumEntryWithoutInitializer(
-                    (KtEnumEntry) ktClass, descriptor,
-                    scopeForConstructor, outerDataFlowInfo, primaryConstructorDelegationCall
+                    (KtEnumEntry) ktClass, descriptor, scopeForConstructor,
+                    outerDataFlowInfo, primaryConstructorDelegationCall, inferenceSession
             );
         }
 
@@ -450,7 +457,8 @@ public class BodyResolver {
             @NotNull ClassDescriptor enumEntryDescriptor,
             @NotNull LexicalScope scopeForConstructor,
             @NotNull DataFlowInfo outerDataFlowInfo,
-            @NotNull ResolvedCall<?>[] primaryConstructorDelegationCall
+            @NotNull ResolvedCall<?>[] primaryConstructorDelegationCall,
+            @Nullable InferenceSession inferenceSession
     ) {
         assert enumEntryDescriptor.getKind() == ClassKind.ENUM_ENTRY : "Enum entry expected: " + enumEntryDescriptor;
         ClassDescriptor enumClassDescriptor = (ClassDescriptor) enumEntryDescriptor.getContainingDeclaration();
@@ -468,8 +476,15 @@ public class BodyResolver {
         Call call = CallMaker.makeConstructorCallWithoutTypeArguments(ktCallEntry);
         trace.record(BindingContext.TYPE, ktCallEntry.getTypeReference(), enumClassDescriptor.getDefaultType());
         trace.record(BindingContext.CALL, ktEnumEntry, call);
-        OverloadResolutionResults<FunctionDescriptor> results =
-                callResolver.resolveFunctionCall(trace, scopeForConstructor, call, NO_EXPECTED_TYPE, outerDataFlowInfo, false);
+        OverloadResolutionResults<FunctionDescriptor> results = callResolver.resolveFunctionCall(
+                trace,
+                scopeForConstructor,
+                call,
+                NO_EXPECTED_TYPE,
+                outerDataFlowInfo,
+                false,
+                inferenceSession
+        );
         if (primaryConstructorDelegationCall[0] == null) {
             primaryConstructorDelegationCall[0] = results.getResultingCall();
         }
@@ -666,14 +681,16 @@ public class BodyResolver {
         for (Map.Entry<KtAnonymousInitializer, ClassDescriptorWithResolutionScopes> entry : c.getAnonymousInitializers().entrySet()) {
             KtAnonymousInitializer initializer = entry.getKey();
             ClassDescriptorWithResolutionScopes descriptor = entry.getValue();
-            resolveAnonymousInitializer(c.getOuterDataFlowInfo(), initializer, descriptor);
+            ExpressionTypingContext context = c.getLocalContext();
+            resolveAnonymousInitializer(c.getOuterDataFlowInfo(), initializer, descriptor, context != null ? context.inferenceSession : null);
         }
     }
 
     public void resolveAnonymousInitializer(
             @NotNull DataFlowInfo outerDataFlowInfo,
             @NotNull KtAnonymousInitializer anonymousInitializer,
-            @NotNull ClassDescriptorWithResolutionScopes classDescriptor
+            @NotNull ClassDescriptorWithResolutionScopes classDescriptor,
+            @Nullable InferenceSession inferenceSession
     ) {
         ProgressManager.checkCanceled();
 
@@ -684,7 +701,7 @@ public class BodyResolver {
                     (KtDeclaration) anonymousInitializer.getParent().getParent(), trace, languageVersionSettings);
             expressionTypingServices.getTypeInfo(
                     scopeForInitializers, body, NO_EXPECTED_TYPE, outerDataFlowInfo,
-                    InferenceSession.Companion.getDefault(), trace, /*isStatement = */true
+                    inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault(), trace, /*isStatement = */true
             );
         }
         processModifiersOnInitializer(anonymousInitializer, scopeForInitializers);
@@ -713,11 +730,13 @@ public class BodyResolver {
             if (unsubstitutedPrimaryConstructor != null) {
                 ForceResolveUtil.forceResolveAllContents(unsubstitutedPrimaryConstructor.getAnnotations());
 
+                ExpressionTypingContext localContext = c.getLocalContext();
                 LexicalScope parameterScope = getPrimaryConstructorParametersScope(classDescriptor.getScopeForConstructorHeaderResolution(),
                                                                                    unsubstitutedPrimaryConstructor);
-                valueParameterResolver.resolveValueParameters(klass.getPrimaryConstructorParameters(),
-                                                              unsubstitutedPrimaryConstructor.getValueParameters(),
-                                                              parameterScope, c.getOuterDataFlowInfo(), trace);
+                valueParameterResolver.resolveValueParameters(
+                        klass.getPrimaryConstructorParameters(), unsubstitutedPrimaryConstructor.getValueParameters(),
+                        parameterScope, c.getOuterDataFlowInfo(), trace, localContext != null ? localContext.inferenceSession : null
+                );
                 // Annotations on value parameter and constructor parameter could be splitted
                 resolveConstructorPropertyDescriptors(klass);
             }
@@ -769,15 +788,22 @@ public class BodyResolver {
         PreliminaryDeclarationVisitor.Companion.createForDeclaration(property, trace, languageVersionSettings);
         KtExpression initializer = property.getInitializer();
         LexicalScope propertyHeaderScope = ScopeUtils.makeScopeForPropertyHeader(getScopeForProperty(c, property), propertyDescriptor);
+        ExpressionTypingContext context = c.getLocalContext();
 
         if (initializer != null) {
-            resolvePropertyInitializer(c.getOuterDataFlowInfo(), property, propertyDescriptor, initializer, propertyHeaderScope);
+            resolvePropertyInitializer(
+                    c.getOuterDataFlowInfo(), property, propertyDescriptor,
+                    initializer, propertyHeaderScope, context != null ? context.inferenceSession : null
+            );
         }
 
         KtExpression delegateExpression = property.getDelegateExpression();
         if (delegateExpression != null) {
             assert initializer == null : "Initializer should be null for delegated property : " + property.getText();
-            resolvePropertyDelegate(c.getOuterDataFlowInfo(), property, propertyDescriptor, delegateExpression, propertyHeaderScope);
+            resolvePropertyDelegate(
+                    c.getOuterDataFlowInfo(), property, propertyDescriptor,
+                    delegateExpression, propertyHeaderScope, context != null ? context.inferenceSession : null
+            );
         }
 
         resolvePropertyAccessors(c, property, propertyDescriptor);
@@ -881,14 +907,15 @@ public class BodyResolver {
             @NotNull KtProperty property,
             @NotNull PropertyDescriptor propertyDescriptor,
             @NotNull KtExpression delegateExpression,
-            @NotNull LexicalScope propertyHeaderScope
+            @NotNull LexicalScope propertyHeaderScope,
+            @Nullable InferenceSession inferenceSession
     ) {
         delegatedPropertyResolver.resolvePropertyDelegate(outerDataFlowInfo,
                                                           property,
                                                           propertyDescriptor,
                                                           delegateExpression,
                                                           propertyHeaderScope,
-                                                          InferenceSession.Companion.getDefault(),
+                                                          inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault(),
                                                           trace);
     }
 
@@ -897,13 +924,16 @@ public class BodyResolver {
             @NotNull KtProperty property,
             @NotNull PropertyDescriptor propertyDescriptor,
             @NotNull KtExpression initializer,
-            @NotNull LexicalScope propertyHeader
+            @NotNull LexicalScope propertyHeader,
+            @Nullable InferenceSession inferenceSession
     ) {
         LexicalScope propertyDeclarationInnerScope = ScopeUtils.makeScopeForPropertyInitializer(propertyHeader, propertyDescriptor);
         KotlinType expectedTypeForInitializer = property.getTypeReference() != null ? propertyDescriptor.getType() : NO_EXPECTED_TYPE;
         if (propertyDescriptor.getCompileTimeInitializer() == null) {
-            expressionTypingServices.getType(propertyDeclarationInnerScope, initializer, expectedTypeForInitializer,
-                                             outerDataFlowInfo, InferenceSession.Companion.getDefault(), trace);
+            expressionTypingServices.getType(
+                    propertyDeclarationInnerScope, initializer, expectedTypeForInitializer,
+                    outerDataFlowInfo, inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault(), trace
+            );
         }
     }
 
@@ -966,7 +996,8 @@ public class BodyResolver {
 
         LexicalScope headerScope = headerScopeFactory != null ? headerScopeFactory.invoke(innerScope) : innerScope;
         valueParameterResolver.resolveValueParameters(
-                valueParameters, valueParameterDescriptors, headerScope, outerDataFlowInfo, trace
+                valueParameters, valueParameterDescriptors, headerScope, outerDataFlowInfo, trace,
+                localContext != null ? localContext.inferenceSession : null
         );
 
         // Synthetic "field" creation
@@ -1008,14 +1039,15 @@ public class BodyResolver {
             @NotNull BindingTrace trace,
             @NotNull KtPrimaryConstructor constructor,
             @NotNull ConstructorDescriptor constructorDescriptor,
-            @NotNull LexicalScope declaringScope
+            @NotNull LexicalScope declaringScope,
+            @Nullable InferenceSession inferenceSession
     ) {
         List<KtParameter> valueParameters = constructor.getValueParameters();
         List<ValueParameterDescriptor> valueParameterDescriptors = constructorDescriptor.getValueParameters();
 
         LexicalScope scope = getPrimaryConstructorParametersScope(declaringScope, constructorDescriptor);
 
-        valueParameterResolver.resolveValueParameters(valueParameters, valueParameterDescriptors, scope, outerDataFlowInfo, trace);
+        valueParameterResolver.resolveValueParameters(valueParameters, valueParameterDescriptors, scope, outerDataFlowInfo, trace, inferenceSession);
     }
 
     public static void computeDeferredType(KotlinType type) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/DelegatedPropertyResolver.kt

@@ -94,10 +94,14 @@ class DelegatedPropertyResolver(
             delegateExpression, property, variableDescriptor, initializerScope, trace, outerDataFlowInfo, inferenceSession
         )
 
-        resolveProvideDelegateMethod(variableDescriptor, delegateExpression, byExpressionType, trace, initializerScope, outerDataFlowInfo)
+        resolveProvideDelegateMethod(
+            variableDescriptor, delegateExpression, byExpressionType, trace, initializerScope, outerDataFlowInfo, inferenceSession
+        )
         val delegateType = getResolvedDelegateType(variableDescriptor, delegateExpression, byExpressionType, trace)
 
-        resolveGetValueMethod(variableDescriptor, delegateExpression, delegateType, trace, initializerScope, outerDataFlowInfo)
+        resolveGetValueMethod(
+            variableDescriptor, delegateExpression, delegateType, trace, initializerScope, outerDataFlowInfo, inferenceSession
+        )
         if (property.isVar) {
             resolveSetValueMethod(variableDescriptor, delegateExpression, delegateType, trace, initializerScope, outerDataFlowInfo)
         }
@@ -123,9 +127,12 @@ class DelegatedPropertyResolver(
         byExpressionType: KotlinType,
         trace: BindingTrace,
         initializerScope: LexicalScope,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession
     ): KotlinType? {
-        resolveProvideDelegateMethod(variableDescriptor, delegateExpression, byExpressionType, trace, initializerScope, dataFlowInfo)
+        resolveProvideDelegateMethod(
+            variableDescriptor, delegateExpression, byExpressionType, trace, initializerScope, dataFlowInfo, inferenceSession
+        )
         val delegateType = getResolvedDelegateType(variableDescriptor, delegateExpression, byExpressionType, trace)
         resolveGetSetValueMethod(variableDescriptor, delegateExpression, delegateType, trace, initializerScope, dataFlowInfo, true)
 
@@ -142,10 +149,12 @@ class DelegatedPropertyResolver(
         delegateType: KotlinType,
         trace: BindingTrace,
         initializerScope: LexicalScope,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession
     ) {
-        val returnType =
-            getGetValueMethodReturnType(variableDescriptor, delegateExpression, delegateType, trace, initializerScope, dataFlowInfo)
+        val returnType = getGetValueMethodReturnType(
+            variableDescriptor, delegateExpression, delegateType, trace, initializerScope, dataFlowInfo, inferenceSession
+        )
         val propertyType = variableDescriptor.type
 
         /* Do not check return type of get() method of delegate for properties with DeferredType because property type is taken from it */
@@ -288,7 +297,8 @@ class DelegatedPropertyResolver(
         byExpressionType: KotlinType,
         trace: BindingTrace,
         initializerScope: LexicalScope,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession
     ) {
         if (!isOperatorProvideDelegateSupported) return
         if (trace.bindingContext.get(PROVIDE_DELEGATE_CALL, propertyDescriptor) != null) return
@@ -297,7 +307,7 @@ class DelegatedPropertyResolver(
 
         val provideDelegateResults = getProvideDelegateMethod(
             propertyDescriptor, byExpression, byExpressionType,
-            traceForProvideDelegate, initializerScope, dataFlowInfo
+            traceForProvideDelegate, initializerScope, dataFlowInfo, inferenceSession
         )
         if (!provideDelegateResults.isSuccess) {
             val call = traceForProvideDelegate.bindingContext.get(PROVIDE_DELEGATE_CALL, propertyDescriptor)
@@ -421,7 +431,8 @@ class DelegatedPropertyResolver(
         delegateExpressionType: KotlinType,
         trace: BindingTrace,
         initializerScope: LexicalScope,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession?
     ): OverloadResolutionResults<FunctionDescriptor> {
         val context = ExpressionTypingContext.newContext(
             trace,
@@ -429,7 +440,8 @@ class DelegatedPropertyResolver(
             dataFlowInfo,
             NO_EXPECTED_TYPE,
             languageVersionSettings,
-            dataFlowValueFactory
+            dataFlowValueFactory,
+            inferenceSession
         )
         return getProvideDelegateMethod(propertyDescriptor, delegateExpression, delegateExpressionType, context)
     }
@@ -567,7 +579,8 @@ class DelegatedPropertyResolver(
             val delegateTypeWithoutNonFixedVariables = nonFixedVariablesToStubTypesSubstitutor.safeSubstitute(delegateType.unwrap())
 
             val contextForProvideDelegate = createContextForProvideDelegateMethod(
-                scopeForDelegate, delegateDataFlow, traceForProvideDelegate, InferenceSessionForExistingCandidates(substitutionMap != null)
+                scopeForDelegate, delegateDataFlow, traceForProvideDelegate,
+                InferenceSessionForExistingCandidates(substitutionMap != null)
             )
 
             val provideDelegateResults = getProvideDelegateMethod(
@@ -589,7 +602,7 @@ class DelegatedPropertyResolver(
         return inferDelegateTypeFromGetSetValueMethods(
             delegateExpression, variableDescriptor, scopeForDelegate,
             traceToResolveDelegatedProperty, delegateType, delegateTypeForProperType,
-            delegateDataFlow
+            delegateDataFlow, inferenceSession
         )
     }
 
@@ -632,10 +645,11 @@ class DelegatedPropertyResolver(
         trace: TemporaryBindingTrace,
         delegateType: KotlinType,
         delegateTypeForProperType: KotlinType?,
-        delegateDataFlow: DataFlowInfo
+        delegateDataFlow: DataFlowInfo,
+        inferenceSession: InferenceSession
     ): UnwrappedType {
         val expectedType = if (variableDescriptor.type !is DeferredType) variableDescriptor.type.unwrap() else null
-        val inferenceSession = DelegatedPropertyInferenceSession(
+        val newInferenceSession = DelegatedPropertyInferenceSession(
             variableDescriptor, expectedType, psiCallResolver,
             postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter,
             callComponents, builtIns
@@ -643,7 +657,7 @@ class DelegatedPropertyResolver(
 
         val receiver = createReceiverForGetSetValueMethods(delegateExpression, delegateType, trace)
         val context = createContextForGetSetValueMethods(
-            variableDescriptor, scopeForDelegate, delegateDataFlow, trace, inferenceSession
+            variableDescriptor, scopeForDelegate, delegateDataFlow, trace, newInferenceSession
         )
 
         fun recordResolvedDelegateOrReportError(
@@ -678,8 +692,8 @@ class DelegatedPropertyResolver(
             )
         }
 
-        val resolutionCallbacks = psiCallResolver.createResolutionCallbacks(trace, inferenceSession, context = null)
-        val resolutionResults = inferenceSession.resolveCandidates(resolutionCallbacks)
+        val resolutionCallbacks = psiCallResolver.createResolutionCallbacks(trace, newInferenceSession, context = null)
+        val resolutionResults = newInferenceSession.resolveCandidates(resolutionCallbacks)
 
         for ((name, isGet) in listOf(OperatorNameConventions.GET_VALUE to true, OperatorNameConventions.SET_VALUE to false)) {
             val result = resolutionResults.firstOrNull {
@@ -779,7 +793,8 @@ class DelegatedPropertyResolver(
             if (isOperatorProvideDelegateSupported) {
                 val provideDelegateResults = getProvideDelegateMethod(
                     variableDescriptor, delegateExpression, byExpressionType,
-                    traceToResolveConventionMethods, scopeForDelegate, dataFlowInfo
+                    traceToResolveConventionMethods, scopeForDelegate,
+                    dataFlowInfo, null // it's used only from the old type inference
                 )
                 if (conventionMethodFound(provideDelegateResults)) {
                     val provideDelegateDescriptor = provideDelegateResults.resultingDescriptor

compiler/frontend/src/org/jetbrains/kotlin/resolve/DescriptorResolver.java

@@ -307,7 +307,8 @@ public class DescriptorResolver {
             int index,
             @NotNull KotlinType type,
             @NotNull BindingTrace trace,
-            @NotNull Annotations additionalAnnotations
+            @NotNull Annotations additionalAnnotations,
+            @Nullable InferenceSession inferenceSession
     ) {
         KotlinType varargElementType = null;
         KotlinType variableType = type;
@@ -340,7 +341,7 @@ public class DescriptorResolver {
                                 destructuringDeclaration, new TransientReceiver(type), /* initializer = */ null,
                                 ExpressionTypingContext.newContext(
                                         trace, scopeForDestructuring, DataFlowInfoFactory.EMPTY, TypeUtils.NO_EXPECTED_TYPE,
-                                        languageVersionSettings, dataFlowValueFactory
+                                        languageVersionSettings, dataFlowValueFactory, inferenceSession
                                 )
                         );
 
@@ -813,7 +814,8 @@ public class DescriptorResolver {
         KtExpression initializer = destructuringDeclaration.getInitializer();
 
         ExpressionTypingContext context = ExpressionTypingContext.newContext(
-                trace, scopeForDeclarationResolution, dataFlowInfo, TypeUtils.NO_EXPECTED_TYPE, languageVersionSettings, dataFlowValueFactory
+                trace, scopeForDeclarationResolution, dataFlowInfo, TypeUtils.NO_EXPECTED_TYPE,
+                languageVersionSettings, dataFlowValueFactory, inferenceSession
         );
 
         ExpressionReceiver receiver = createReceiverForDestructuringDeclaration(destructuringDeclaration, context);
@@ -874,7 +876,7 @@ public class DescriptorResolver {
             @NotNull KtVariableDeclaration variableDeclaration,
             @NotNull BindingTrace trace,
             @NotNull DataFlowInfo dataFlowInfo,
-            @NotNull InferenceSession inferenceSession,
+            @Nullable InferenceSession inferenceSession,
             @NotNull VariableAsPropertyInfo propertyInfo
     ) {
         KtModifierList modifierList = variableDeclaration.getModifierList();
@@ -987,7 +989,9 @@ public class DescriptorResolver {
                 trace,
                 typeIfKnown,
                 propertyInfo.getPropertyGetter(),
-                propertyInfo.getHasDelegate());
+                propertyInfo.getHasDelegate(),
+                inferenceSession
+        );
 
         KotlinType type = typeIfKnown != null ? typeIfKnown : getter.getReturnType();
 
@@ -1006,7 +1010,9 @@ public class DescriptorResolver {
                 annotationSplitter,
                 trace,
                 propertyInfo.getPropertySetter(),
-                propertyInfo.getHasDelegate());
+                propertyInfo.getHasDelegate(),
+                inferenceSession
+        );
 
         propertyDescriptor.initialize(
                 getter, setter,
@@ -1064,7 +1070,8 @@ public class DescriptorResolver {
             @NotNull AnnotationSplitter annotationSplitter,
             @NotNull BindingTrace trace,
             @Nullable KtPropertyAccessor setter,
-            boolean hasDelegate
+            boolean hasDelegate,
+            @Nullable InferenceSession inferenceSession
     ) {
         PropertySetterDescriptorImpl setterDescriptor = null;
         Annotations setterTargetedAnnotations = annotationSplitter.getAnnotationsForTarget(PROPERTY_SETTER);
@@ -1114,7 +1121,7 @@ public class DescriptorResolver {
                 }
 
                 ValueParameterDescriptorImpl valueParameterDescriptor = resolveValueParameterDescriptor(
-                        scopeWithTypeParameters, setterDescriptor, parameter, 0, type, trace, parameterTargetedAnnotations
+                        scopeWithTypeParameters, setterDescriptor, parameter, 0, type, trace, parameterTargetedAnnotations, inferenceSession
                 );
                 setterDescriptor.initialize(valueParameterDescriptor);
             }
@@ -1150,7 +1157,8 @@ public class DescriptorResolver {
             @NotNull BindingTrace trace,
             @Nullable KotlinType propertyTypeIfKnown,
             @Nullable KtPropertyAccessor getter,
-            boolean hasDelegate
+            boolean hasDelegate,
+            @Nullable InferenceSession inferenceSession
     ) {
         PropertyGetterDescriptorImpl getterDescriptor;
         KotlinType getterType;
@@ -1170,7 +1178,9 @@ public class DescriptorResolver {
                     property.hasModifier(KtTokens.INLINE_KEYWORD) || getter.hasModifier(KtTokens.INLINE_KEYWORD),
                     CallableMemberDescriptor.Kind.DECLARATION, null, KotlinSourceElementKt.toSourceElement(getter)
             );
-            getterType = determineGetterReturnType(scopeForDeclarationResolution, trace, getterDescriptor, getter, propertyTypeIfKnown);
+            getterType = determineGetterReturnType(
+                    scopeForDeclarationResolution, trace, getterDescriptor, getter, propertyTypeIfKnown, inferenceSession
+            );
             trace.record(BindingContext.PROPERTY_ACCESSOR, getter, getterDescriptor);
         }
         else {
@@ -1193,7 +1203,8 @@ public class DescriptorResolver {
             @NotNull BindingTrace trace,
             @NotNull PropertyGetterDescriptor getterDescriptor,
             @NotNull KtPropertyAccessor getter,
-            @Nullable KotlinType propertyTypeIfKnown
+            @Nullable KotlinType propertyTypeIfKnown,
+            @Nullable InferenceSession inferenceSession
     ) {
         KtTypeReference returnTypeReference = getter.getReturnTypeReference();
         if (returnTypeReference != null) {
@@ -1210,7 +1221,7 @@ public class DescriptorResolver {
         KtProperty property = getter.getProperty();
         if (!property.hasDelegateExpressionOrInitializer() && property.getTypeReference() == null &&
             getter.hasBody() && !getter.hasBlockBody()) {
-            return inferReturnTypeFromExpressionBody(trace, scope, DataFlowInfoFactory.EMPTY, getter, getterDescriptor);
+            return inferReturnTypeFromExpressionBody(trace, scope, DataFlowInfoFactory.EMPTY, getter, getterDescriptor, inferenceSession);
         }
 
         return propertyTypeIfKnown;
@@ -1222,11 +1233,14 @@ public class DescriptorResolver {
             @NotNull LexicalScope scope,
             @NotNull DataFlowInfo dataFlowInfo,
             @NotNull KtDeclarationWithBody function,
-            @NotNull FunctionDescriptor functionDescriptor
+            @NotNull FunctionDescriptor functionDescriptor,
+            @Nullable InferenceSession inferenceSession
     ) {
         return wrappedTypeFactory.createRecursionIntolerantDeferredType(trace, () -> {
             PreliminaryDeclarationVisitor.Companion.createForDeclaration(function, trace, languageVersionSettings);
-            KotlinType type = expressionTypingServices.getBodyExpressionType(trace, scope, dataFlowInfo, function, functionDescriptor);
+            KotlinType type = expressionTypingServices.getBodyExpressionType(
+                    trace, scope, dataFlowInfo, function, functionDescriptor, inferenceSession
+            );
             KotlinType publicType = transformAnonymousTypeIfNeeded(
                     functionDescriptor, function, type, trace, anonymousTypeTransformers, languageVersionSettings
             );

compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.resolve.ModifiersChecker.resolveMemberModalityFromMo
 import org.jetbrains.kotlin.resolve.ModifiersChecker.resolveVisibilityFromModifiers
 import org.jetbrains.kotlin.resolve.bindingContextUtil.recordScope
 import org.jetbrains.kotlin.resolve.calls.DslMarkerUtils
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.util.createValueParametersForInvokeInFunctionType
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
@@ -84,12 +85,14 @@ class FunctionDescriptorResolver(
         scope: LexicalScope,
         function: KtNamedFunction,
         trace: BindingTrace,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession?
     ): SimpleFunctionDescriptor {
         if (function.name == null) trace.report(FUNCTION_DECLARATION_WITH_NO_NAME.on(function))
 
         return resolveFunctionDescriptor(
-            SimpleFunctionDescriptorImpl::create, containingDescriptor, scope, function, trace, dataFlowInfo, TypeUtils.NO_EXPECTED_TYPE
+            SimpleFunctionDescriptorImpl::create, containingDescriptor, scope,
+            function, trace, dataFlowInfo, TypeUtils.NO_EXPECTED_TYPE, inferenceSession
         )
     }
 
@@ -100,9 +103,10 @@ class FunctionDescriptorResolver(
         function: KtNamedFunction,
         trace: BindingTrace,
         dataFlowInfo: DataFlowInfo,
-        expectedFunctionType: KotlinType
+        expectedFunctionType: KotlinType,
+        inferenceSession: InferenceSession?
     ): SimpleFunctionDescriptor = resolveFunctionDescriptor(
-        ::FunctionExpressionDescriptor, containingDescriptor, scope, function, trace, dataFlowInfo, expectedFunctionType
+        ::FunctionExpressionDescriptor, containingDescriptor, scope, function, trace, dataFlowInfo, expectedFunctionType, inferenceSession
     )
 
     private fun resolveFunctionDescriptor(
@@ -112,7 +116,8 @@ class FunctionDescriptorResolver(
         function: KtNamedFunction,
         trace: BindingTrace,
         dataFlowInfo: DataFlowInfo,
-        expectedFunctionType: KotlinType
+        expectedFunctionType: KotlinType,
+        inferenceSession: InferenceSession?
     ): SimpleFunctionDescriptor {
         val functionDescriptor = functionConstructor(
             containingDescriptor,
@@ -128,9 +133,10 @@ class FunctionDescriptorResolver(
             functionDescriptor,
             trace,
             expectedFunctionType,
-            dataFlowInfo
+            dataFlowInfo,
+            inferenceSession
         )
-        initializeFunctionReturnTypeBasedOnFunctionBody(scope, function, functionDescriptor, trace, dataFlowInfo)
+        initializeFunctionReturnTypeBasedOnFunctionBody(scope, function, functionDescriptor, trace, dataFlowInfo, inferenceSession)
         BindingContextUtils.recordFunctionDeclarationToDescriptor(trace, function, functionDescriptor)
         return functionDescriptor
     }
@@ -140,7 +146,8 @@ class FunctionDescriptorResolver(
         function: KtNamedFunction,
         functionDescriptor: SimpleFunctionDescriptorImpl,
         trace: BindingTrace,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession?
     ) {
         if (functionDescriptor.returnType != null) return
         assert(function.typeReference == null) {
@@ -153,7 +160,9 @@ class FunctionDescriptorResolver(
             function.hasBlockBody() ->
                 builtIns.unitType
             function.hasBody() ->
-                descriptorResolver.inferReturnTypeFromExpressionBody(trace, scope, dataFlowInfo, function, functionDescriptor)
+                descriptorResolver.inferReturnTypeFromExpressionBody(
+                    trace, scope, dataFlowInfo, function, functionDescriptor, inferenceSession
+                )
             else ->
                 ErrorUtils.createErrorType("No type, no body")
         }
@@ -167,7 +176,8 @@ class FunctionDescriptorResolver(
         functionDescriptor: SimpleFunctionDescriptorImpl,
         trace: BindingTrace,
         expectedFunctionType: KotlinType,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        inferenceSession: InferenceSession?
     ) {
         val headerScope = LexicalWritableScope(
             scope, functionDescriptor, true,
@@ -188,7 +198,7 @@ class FunctionDescriptorResolver(
 
 
         val valueParameterDescriptors =
-            createValueParameterDescriptors(function, functionDescriptor, headerScope, trace, expectedFunctionType)
+            createValueParameterDescriptors(function, functionDescriptor, headerScope, trace, expectedFunctionType, inferenceSession)
 
         headerScope.freeze()
 
@@ -200,7 +210,7 @@ class FunctionDescriptorResolver(
             trace.bindingContext, container
         )
 
-        val contractProvider = getContractProvider(functionDescriptor, trace, scope, dataFlowInfo, function)
+        val contractProvider = getContractProvider(functionDescriptor, trace, scope, dataFlowInfo, function, inferenceSession)
         val userData = mutableMapOf<CallableDescriptor.UserDataKey<*>, Any>().apply {
             if (contractProvider != null) {
                 put(ContractProviderKey, contractProvider)
@@ -250,7 +260,8 @@ class FunctionDescriptorResolver(
         trace: BindingTrace,
         scope: LexicalScope,
         dataFlowInfo: DataFlowInfo,
-        function: KtFunction
+        function: KtFunction,
+        inferenceSession: InferenceSession?
     ): LazyContractProvider? {
         if (function !is KtNamedFunction) return null
 
@@ -261,7 +272,7 @@ class FunctionDescriptorResolver(
 
         return LazyContractProvider(storageManager) {
             AstLoadingFilter.forceAllowTreeLoading(function.containingFile, ThrowableComputable {
-                expressionTypingServices.getBodyExpressionType(trace, scope, dataFlowInfo, function, functionDescriptor)
+                expressionTypingServices.getBodyExpressionType(trace, scope, dataFlowInfo, function, functionDescriptor, inferenceSession)
             })
         }
     }
@@ -271,7 +282,8 @@ class FunctionDescriptorResolver(
         functionDescriptor: SimpleFunctionDescriptorImpl,
         innerScope: LexicalWritableScope,
         trace: BindingTrace,
-        expectedFunctionType: KotlinType
+        expectedFunctionType: KotlinType,
+        inferenceSession: InferenceSession?
     ): List<ValueParameterDescriptor> {
         val expectedValueParameters = expectedFunctionType.getValueParameters(functionDescriptor)
         val expectedParameterTypes = expectedValueParameters?.map { it.type.removeParameterNameAnnotation() }
@@ -300,7 +312,8 @@ class FunctionDescriptorResolver(
             innerScope,
             function.valueParameters,
             trace,
-            expectedParameterTypes
+            expectedParameterTypes,
+            inferenceSession
         )
     }
 
@@ -324,7 +337,8 @@ class FunctionDescriptorResolver(
         classDescriptor: ClassDescriptor,
         classElement: KtPureClassOrObject,
         trace: BindingTrace,
-        languageVersionSettings: LanguageVersionSettings
+        languageVersionSettings: LanguageVersionSettings,
+        inferenceSession: InferenceSession?
     ): ClassConstructorDescriptorImpl? {
         if (classDescriptor.kind == ClassKind.ENUM_ENTRY || !classElement.hasPrimaryConstructor()) return null
         return createConstructorDescriptor(
@@ -335,7 +349,8 @@ class FunctionDescriptorResolver(
             classElement.primaryConstructor ?: classElement,
             classElement.primaryConstructorParameters,
             trace,
-            languageVersionSettings
+            languageVersionSettings,
+            inferenceSession
         )
     }
 
@@ -344,7 +359,8 @@ class FunctionDescriptorResolver(
         classDescriptor: ClassDescriptor,
         constructor: KtSecondaryConstructor,
         trace: BindingTrace,
-        languageVersionSettings: LanguageVersionSettings
+        languageVersionSettings: LanguageVersionSettings,
+        inferenceSession: InferenceSession?
     ): ClassConstructorDescriptorImpl {
         return createConstructorDescriptor(
             scope,
@@ -354,7 +370,8 @@ class FunctionDescriptorResolver(
             constructor,
             constructor.valueParameters,
             trace,
-            languageVersionSettings
+            languageVersionSettings,
+            inferenceSession
         )
     }
 
@@ -366,7 +383,8 @@ class FunctionDescriptorResolver(
         declarationToTrace: KtPureElement,
         valueParameters: List<KtParameter>,
         trace: BindingTrace,
-        languageVersionSettings: LanguageVersionSettings
+        languageVersionSettings: LanguageVersionSettings,
+        inferenceSession: InferenceSession?
     ): ClassConstructorDescriptorImpl {
         val constructorDescriptor = ClassConstructorDescriptorImpl.create(
             classDescriptor,
@@ -375,8 +393,7 @@ class FunctionDescriptorResolver(
             declarationToTrace.toSourceElement()
         )
         constructorDescriptor.isExpect = classDescriptor.isExpect
-        constructorDescriptor.isActual =
-                modifierList?.hasActualModifier() == true ||
+        constructorDescriptor.isActual = modifierList?.hasActualModifier() == true ||
                 // We don't require 'actual' for constructors of actual annotations
                 classDescriptor.kind == ClassKind.ANNOTATION_CLASS && classDescriptor.isActual
         if (declarationToTrace is PsiElement)
@@ -389,7 +406,9 @@ class FunctionDescriptorResolver(
             LexicalScopeKind.CONSTRUCTOR_HEADER
         )
         val constructor = constructorDescriptor.initialize(
-            resolveValueParameters(constructorDescriptor, parameterScope, valueParameters, trace, null),
+            resolveValueParameters(
+                constructorDescriptor, parameterScope, valueParameters, trace, null, inferenceSession
+            ),
             resolveVisibilityFromModifiers(
                 modifierList,
                 DescriptorUtils.getDefaultConstructorVisibility(classDescriptor, languageVersionSettings.supportsFeature(LanguageFeature.AllowSealedInheritorsInDifferentFilesOfSamePackage))
@@ -407,7 +426,8 @@ class FunctionDescriptorResolver(
         parameterScope: LexicalWritableScope,
         valueParameters: List<KtParameter>,
         trace: BindingTrace,
-        expectedParameterTypes: List<KotlinType>?
+        expectedParameterTypes: List<KotlinType>?,
+        inferenceSession: InferenceSession?
     ): List<ValueParameterDescriptor> {
         val result = ArrayList<ValueParameterDescriptor>()
 
@@ -451,7 +471,7 @@ class FunctionDescriptorResolver(
             }
 
             val valueParameterDescriptor = descriptorResolver.resolveValueParameterDescriptor(
-                parameterScope, functionDescriptor, valueParameter, i, type, trace, Annotations.EMPTY
+                parameterScope, functionDescriptor, valueParameter, i, type, trace, Annotations.EMPTY, inferenceSession
             )
 
             // Do not report NAME_SHADOWING for lambda destructured parameters as they may be not fully resolved at this time

compiler/frontend/src/org/jetbrains/kotlin/resolve/VariableTypeAndInitializerResolver.kt

@@ -153,7 +153,7 @@ class VariableTypeAndInitializerResolver(
         )
 
         val getterReturnType = delegatedPropertyResolver.getGetValueMethodReturnType(
-            variableDescriptor, delegateExpression, type, trace, scopeForInitializer, dataFlowInfo
+            variableDescriptor, delegateExpression, type, trace, scopeForInitializer, dataFlowInfo, inferenceSession
         )
 
         val delegatedType = getterReturnType?.let { approximateType(it, local) }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallResolver.java

@@ -297,12 +297,14 @@ public class CallResolver {
             @NotNull Call call,
             @NotNull KotlinType expectedType,
             @NotNull DataFlowInfo dataFlowInfo,
-            boolean isAnnotationContext
+            boolean isAnnotationContext,
+            @Nullable InferenceSession inferenceSession
     ) {
         return resolveFunctionCall(
                 BasicCallResolutionContext.create(
                         trace, scope, call, expectedType, dataFlowInfo, ContextDependency.INDEPENDENT, CheckArgumentTypesMode.CHECK_VALUE_ARGUMENTS,
-                        isAnnotationContext, languageVersionSettings, dataFlowValueFactory, InferenceSession.Companion.getDefault()
+                        isAnnotationContext, languageVersionSettings, dataFlowValueFactory,
+                        inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault()
                 )
         );
     }
@@ -423,7 +425,8 @@ public class CallResolver {
     public OverloadResolutionResults<ConstructorDescriptor> resolveConstructorDelegationCall(
             @NotNull BindingTrace trace, @NotNull LexicalScope scope, @NotNull DataFlowInfo dataFlowInfo,
             @NotNull ClassConstructorDescriptor constructorDescriptor,
-            @NotNull KtConstructorDelegationCall call
+            @NotNull KtConstructorDelegationCall call,
+            @Nullable InferenceSession inferenceSession
     ) {
         // Method returns `null` when there is nothing to resolve in trivial cases like `null` call expression or
         // when super call should be conventional enum constructor and super call should be empty
@@ -436,7 +439,7 @@ public class CallResolver {
                 false,
                 languageVersionSettings,
                 dataFlowValueFactory,
-                InferenceSession.Companion.getDefault());
+                inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault());
 
         KtConstructorDelegationReferenceExpression calleeExpression = call.getCalleeExpression();
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/BuilderInferenceSession.kt

@@ -66,7 +66,7 @@ class BuilderInferenceSession(
 
     private val commonCalls = arrayListOf<PSICompletedCallInfo>()
 
-    // Simple calls are calls which might not have gone through type inference, but may contain unsubstituted postponed variables inside their types.
+    // These calls come from the old type inference
     private val oldCallableReferenceCalls = arrayListOf<KtExpression>()
 
     private var hasInapplicableCall = false

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/LazyClassContext.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.SupertypeLoopChecker
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 import org.jetbrains.kotlin.resolve.lazy.declarations.DeclarationProviderFactory
 import org.jetbrains.kotlin.resolve.sam.SamConversionResolver
@@ -30,6 +31,7 @@ import org.jetbrains.kotlin.types.checker.NewKotlinTypeChecker
 
 interface LazyClassContext {
     val declarationScopeProvider: DeclarationScopeProvider
+    val inferenceSession: InferenceSession?
 
     val storageManager: StorageManager
     val trace: BindingTrace

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/ResolveSession.java

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.*;
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession;
 import org.jetbrains.kotlin.resolve.checkers.PlatformDiagnosticSuppressor;
 import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension;
 import org.jetbrains.kotlin.resolve.lazy.declarations.DeclarationProviderFactory;
@@ -525,4 +526,11 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
     public SealedClassInheritorsProvider getSealedClassInheritorsProvider() {
         return sealedClassInheritorsProvider;
     }
+
+    // TODO: get inference session, otherwise, for instance, builder inference session is lost here
+    @Nullable
+    @Override
+    public InferenceSession getInferenceSession() {
+        return null;
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/AbstractLazyMemberScope.kt

@@ -117,7 +117,8 @@ protected constructor(
                     getScopeForMemberDeclarationResolution(functionDeclaration),
                     functionDeclaration,
                     trace,
-                    c.declarationScopeProvider.getOuterDataFlowInfoForDeclaration(functionDeclaration)
+                    c.declarationScopeProvider.getOuterDataFlowInfoForDeclaration(functionDeclaration),
+                    c.inferenceSession
                 )
             )
         }

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyClassMemberScope.kt

@@ -505,7 +505,8 @@ open class LazyClassMemberScope(
 
         if (DescriptorUtils.canHaveDeclaredConstructors(thisDescriptor) || hasPrimaryConstructor) {
             val constructor = c.functionDescriptorResolver.resolvePrimaryConstructorDescriptor(
-                thisDescriptor.scopeForConstructorHeaderResolution, thisDescriptor, classOrObject, trace, c.languageVersionSettings
+                thisDescriptor.scopeForConstructorHeaderResolution, thisDescriptor,
+                classOrObject, trace, c.languageVersionSettings, c.inferenceSession
             )
             constructor ?: return null
             setDeferredReturnType(constructor)
@@ -522,7 +523,8 @@ open class LazyClassMemberScope(
 
         return classOrObject.secondaryConstructors.map { constructor ->
             val descriptor = c.functionDescriptorResolver.resolveSecondaryConstructorDescriptor(
-                thisDescriptor.scopeForConstructorHeaderResolution, thisDescriptor, constructor, trace, c.languageVersionSettings
+                thisDescriptor.scopeForConstructorHeaderResolution, thisDescriptor,
+                constructor, trace, c.languageVersionSettings, c.inferenceSession
             )
             setDeferredReturnType(descriptor)
             descriptor

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingContext.java

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.types.expressions;
 
 import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.config.LanguageVersionSettings;
 import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
@@ -41,10 +42,12 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
             @NotNull KotlinType expectedType,
             @NotNull LanguageVersionSettings languageVersionSettings,
             @NotNull DataFlowValueFactory dataFlowValueFactory,
-            @NotNull InferenceSession inferenceSession
+            @Nullable InferenceSession inferenceSession
     ) {
-        return newContext(trace, scope, dataFlowInfo, expectedType, ContextDependency.INDEPENDENT, StatementFilter.NONE,
-                          languageVersionSettings, dataFlowValueFactory, inferenceSession);
+        return newContext(
+                trace, scope, dataFlowInfo, expectedType, ContextDependency.INDEPENDENT, StatementFilter.NONE, languageVersionSettings,
+                dataFlowValueFactory, inferenceSession != null ? inferenceSession : InferenceSession.Companion.getDefault()
+        );
     }
 
     @NotNull

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java

@@ -159,7 +159,7 @@ public class ExpressionTypingServices {
                 trace,
                 functionInnerScope, dataFlowInfo, expectedReturnType != null ? expectedReturnType : NO_EXPECTED_TYPE,
                 getLanguageVersionSettings(), expressionTypingComponents.dataFlowValueFactory,
-                localContext == null ? InferenceSession.Companion.getDefault() : localContext.inferenceSession
+                localContext != null ? localContext.inferenceSession : InferenceSession.Companion.getDefault()
         );
 
         checkFunctionReturnType(function, context);
@@ -221,7 +221,8 @@ public class ExpressionTypingServices {
             @NotNull LexicalScope outerScope,
             @NotNull DataFlowInfo dataFlowInfo,
             @NotNull KtDeclarationWithBody function,
-            @NotNull FunctionDescriptor functionDescriptor
+            @NotNull FunctionDescriptor functionDescriptor,
+            @Nullable InferenceSession inferenceSession
     ) {
         KtExpression bodyExpression = function.getBodyExpression();
         assert bodyExpression != null;
@@ -230,7 +231,7 @@ public class ExpressionTypingServices {
 
         ExpressionTypingContext context = ExpressionTypingContext.newContext(
                 trace, functionInnerScope, dataFlowInfo, NO_EXPECTED_TYPE, getLanguageVersionSettings(),
-                expressionTypingComponents.dataFlowValueFactory
+                expressionTypingComponents.dataFlowValueFactory, inferenceSession
         );
 
         KotlinResolutionCallbacksImpl.LambdaInfo lambdaInfo = getNewInferenceLambdaInfo(context, function);

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt

@@ -77,7 +77,7 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         val functionDescriptor: SimpleFunctionDescriptor
         if (isDeclaration) {
             functionDescriptor = components.functionDescriptorResolver.resolveFunctionDescriptor(
-                context.scope.ownerDescriptor, context.scope, function, context.trace, context.dataFlowInfo
+                context.scope.ownerDescriptor, context.scope, function, context.trace, context.dataFlowInfo, context.inferenceSession
             )
             assert(statementScope != null) {
                 "statementScope must be not null for function: " + function.name + " at location " + PsiDiagnosticUtils.atLocation(
@@ -88,7 +88,7 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         } else {
             functionDescriptor = components.functionDescriptorResolver.resolveFunctionExpressionDescriptor(
                 context.scope.ownerDescriptor, context.scope, function,
-                context.trace, context.dataFlowInfo, context.expectedType
+                context.trace, context.dataFlowInfo, context.expectedType, context.inferenceSession
             )
         }
         // Necessary for local functions
@@ -105,7 +105,8 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         }
 
         components.valueParameterResolver.resolveValueParameters(
-            function.valueParameters, functionDescriptor.valueParameters, functionInnerScope, context.dataFlowInfo, context.trace
+            function.valueParameters, functionDescriptor.valueParameters, functionInnerScope,
+            context.dataFlowInfo, context.trace, context.inferenceSession
         )
 
         components.modifiersChecker.withTrace(context.trace).checkModifiersForLocalDeclaration(function, functionDescriptor)
@@ -211,7 +212,7 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         }
         components.functionDescriptorResolver.initializeFunctionDescriptorAndExplicitReturnType(
             context.scope.ownerDescriptor, context.scope, functionLiteral,
-            functionDescriptor, context.trace, context.expectedType, context.dataFlowInfo
+            functionDescriptor, context.trace, context.expectedType, context.dataFlowInfo, context.inferenceSession
         )
         for (parameterDescriptor in functionDescriptor.valueParameters) {
             ForceResolveUtil.forceResolveAllContents(parameterDescriptor.annotations)

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/LocalClassifierAnalyzer.kt

@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.debugText.getDebugText
 import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 import org.jetbrains.kotlin.resolve.lazy.*
@@ -158,6 +159,7 @@ class LocalClassDescriptorHolder(
             classDescriptor = LazyClassDescriptor(
                 object : LazyClassContext {
                     override val declarationScopeProvider = declarationScopeProvider
+                    override val inferenceSession = expressionTypingContext.inferenceSession
                     override val storageManager = this@LocalClassDescriptorHolder.storageManager
                     override val trace = expressionTypingContext.trace
                     override val moduleDescriptor = this@LocalClassDescriptorHolder.moduleDescriptor

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ValueParameterResolver.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.DescriptorResolver
 import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
@@ -44,14 +45,15 @@ class ValueParameterResolver(
         valueParameterDescriptors: List<ValueParameterDescriptor>,
         declaringScope: LexicalScope,
         dataFlowInfo: DataFlowInfo,
-        trace: BindingTrace
+        trace: BindingTrace,
+        inferenceSession: InferenceSession?
     ) {
         val scopeForDefaultValue =
             LexicalScopeImpl(declaringScope, declaringScope.ownerDescriptor, false, null, LexicalScopeKind.DEFAULT_VALUE)
 
         val contextForDefaultValue = ExpressionTypingContext.newContext(
             trace, scopeForDefaultValue, dataFlowInfo, TypeUtils.NO_EXPECTED_TYPE,
-            languageVersionSettings, dataFlowValueFactory
+            languageVersionSettings, dataFlowValueFactory, inferenceSession
         )
 
         for ((descriptor, parameter) in valueParameterDescriptors.zip(valueParameters)) {

compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt

@@ -0,0 +1,54 @@
+// WITH_RUNTIME
+// DONT_TARGET_EXACT_BACKEND: WASM
+
+@OptIn(ExperimentalStdlibApi::class)
+fun foo1() {
+    buildList {
+        object {
+            fun foo() = add("")
+        }
+    }
+}
+
+@OptIn(ExperimentalStdlibApi::class)
+fun foo2() {
+    buildList {
+        class A {
+            fun foo() = add("")
+        }
+    }
+}
+
+@OptIn(ExperimentalStdlibApi::class)
+fun foo3() {
+    buildList {
+        object {
+            var x: Int
+                get() = 1
+                set(value) {
+                    add(value)
+                }
+        }
+    }
+}
+
+@OptIn(ExperimentalStdlibApi::class)
+fun foo4() {
+    buildList {
+        class A {
+            var x: Int
+                get() = 1
+                set(value) {
+                    add(value)
+                }
+        }
+    }
+}
+
+fun box(): String {
+    foo1()
+    foo2()
+    foo3()
+    foo4()
+    return "OK"
+}

compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt

@@ -0,0 +1,109 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -EXPERIMENTAL_IS_NOT_ENABLED
+// WITH_RUNTIME
+
+import kotlin.experimental.ExperimentalTypeInference
+
+@OptIn(ExperimentalTypeInference::class)
+class A1<T> {
+    fun <BT1> builder1(@BuilderInference configure: A2<BT1>.() -> Unit) {}
+}
+
+@OptIn(ExperimentalTypeInference::class)
+class A2<A2_BT1> {
+    fun <BT2> builder2(@BuilderInference configure: A3<A2_BT1, BT2>.() -> Unit) {}
+}
+
+@OptIn(ExperimentalTypeInference::class)
+class A3<A3_BT1, A3_BT2> {
+    fun <BT3> builder3(@BuilderInference configure: A4<A3_BT1, A3_BT2, BT3>.() -> Unit) {}
+}
+
+class A4<A3_BT1, A3_BT2, A3_BT3> {
+    fun resolver(x: A3_BT3) {}
+}
+
+fun foo1(x: A1<String>) {
+    x.builder1<String> {
+        val x by lazy {
+            builder2<String> {
+                builder3 {
+                    resolver("")
+                }
+            }
+        }
+    }
+}
+
+fun foo2(x: A1<String>) {
+    x.builder1<String> {
+        builder2<String> {
+            val x by lazy {
+                builder3 {
+                    resolver("")
+                }
+            }
+        }
+    }
+}
+
+fun foo3(x: A1<String>) {
+    x.builder1<String> {
+        builder2<String> {
+            class A {
+                fun foo() = builder3 {
+                    resolver("")
+                }
+            }
+        }
+    }
+}
+
+fun foo4(x: A1<String>) {
+    x.builder1<String> {
+        class A {
+            fun foo() = builder2<String> {
+                builder3 {
+                    resolver("")
+                }
+            }
+        }
+    }
+}
+
+fun foo5(x: A1<String>) {
+    x.builder1<String> {
+        builder2<String> {
+            class A {
+                fun foo() {
+                    builder3 {
+                        resolver("")
+                    }
+                }
+            }
+        }
+    }
+}
+
+fun foo6(x: A1<String>) {
+    x.builder1<String> {
+        class A {
+            fun foo() {
+                builder2<String> {
+                    builder3 {
+                        resolver("")
+                    }
+                }
+            }
+        }
+    }
+}
+
+fun box(): String {
+    foo1(A1())
+    foo2(A1())
+    foo3(A1())
+    foo4(A1())
+    foo5(A1())
+    foo6(A1())
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -16969,6 +16969,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true);
             }
 
+            @Test
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @Test
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
@@ -16987,18 +16993,6 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @Test
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @Test
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
@@ -17058,6 +17052,24 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @Test
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @Test
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @Test
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -16969,6 +16969,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
             }
 
+            @Test
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @Test
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
@@ -16987,18 +16993,6 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @Test
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @Test
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
@@ -17058,6 +17052,24 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @Test
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @Test
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @Test
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -14052,6 +14052,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true);
             }
 
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
@@ -14067,16 +14072,6 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");
@@ -14126,6 +14121,21 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

compiler/tests/org/jetbrains/kotlin/types/DefaultModalityModifiersTest.java

@@ -144,7 +144,7 @@ public class DefaultModalityModifiersTest extends KotlinTestWithEnvironment {
             KtNamedFunction function = (KtNamedFunction) declarations.get(0);
             SimpleFunctionDescriptor functionDescriptor =
                     functionDescriptorResolver.resolveFunctionDescriptor(classDescriptor, scope, function,
-                                                                         DummyTraces.DUMMY_TRACE, DataFlowInfoFactory.EMPTY);
+                                                                         DummyTraces.DUMMY_TRACE, DataFlowInfoFactory.EMPTY, null);
 
             assertEquals(expectedFunctionModality, functionDescriptor.getModality());
         }

compiler/tests/org/jetbrains/kotlin/types/KotlinOverloadTest.java

@@ -171,7 +171,7 @@ public class KotlinOverloadTest extends KotlinTestWithEnvironment {
         KtNamedFunction function = KtPsiFactoryKt.KtPsiFactory(getProject()).createFunction(funDecl);
         LexicalScope scope = TypeTestUtilsKt.builtInPackageAsLexicalScope(module);
         return functionDescriptorResolver.resolveFunctionDescriptor(
-                module, scope, function, DummyTraces.DUMMY_TRACE, DataFlowInfoFactory.EMPTY
+                module, scope, function, DummyTraces.DUMMY_TRACE, DataFlowInfoFactory.EMPTY, null
         );
     }
 }

compiler/tests/org/jetbrains/kotlin/types/KotlinOverridingTest.java

@@ -172,7 +172,7 @@ public class KotlinOverridingTest extends KotlinTestWithEnvironment {
         KtNamedFunction function = KtPsiFactoryKt.KtPsiFactory(getProject()).createFunction(funDecl);
         LexicalScope scope = TypeTestUtilsKt.builtInPackageAsLexicalScope(module);
         return functionDescriptorResolver.resolveFunctionDescriptor(
-                module, scope, function, DummyTraces.DUMMY_TRACE, DataFlowInfoFactory.EMPTY
+                module, scope, function, DummyTraces.DUMMY_TRACE, DataFlowInfoFactory.EMPTY, null
         );
     }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/ResolveElementCache.kt

@@ -601,7 +601,8 @@ class ResolveElementCache(
             descriptor,
             descriptor.unsubstitutedPrimaryConstructor,
             descriptor.scopeForConstructorHeaderResolution,
-            descriptor.scopeForMemberDeclarationResolution
+            descriptor.scopeForMemberDeclarationResolution,
+            resolveSession.inferenceSession
         )
 
         return trace
@@ -725,7 +726,8 @@ class ResolveElementCache(
                 trace,
                 primaryConstructor,
                 constructorDescriptor,
-                scope
+                scope,
+                resolveSession.inferenceSession
             )
 
             forceResolveAnnotationsInside(primaryConstructor)
@@ -755,7 +757,9 @@ class ResolveElementCache(
         val classOrObjectDescriptor = resolveSession.resolveToDescriptor(anonymousInitializer.containingDeclaration) as LazyClassDescriptor
 
         val bodyResolver = createBodyResolver(resolveSession, trace, file, statementFilter)
-        bodyResolver.resolveAnonymousInitializer(DataFlowInfo.EMPTY, anonymousInitializer, classOrObjectDescriptor)
+        bodyResolver.resolveAnonymousInitializer(
+            DataFlowInfo.EMPTY, anonymousInitializer, classOrObjectDescriptor, resolveSession.inferenceSession
+        )
 
         forceResolveAnnotationsInside(anonymousInitializer)
 

idea/src/org/jetbrains/kotlin/idea/inspections/RedundantSamConstructorInspection.kt

@@ -128,7 +128,8 @@ class RedundantSamConstructorInspection : AbstractKotlinInspection() {
             val qualifiedExpression = parentCall.getQualifiedExpressionForSelectorOrThis()
             val expectedType = context[BindingContext.EXPECTED_EXPRESSION_TYPE, qualifiedExpression] ?: TypeUtils.NO_EXPECTED_TYPE
 
-            val resolutionResults = callResolver.resolveFunctionCall(BindingTraceContext(), scope, newCall, expectedType, dataFlow, false)
+            val resolutionResults =
+                callResolver.resolveFunctionCall(BindingTraceContext(), scope, newCall, expectedType, dataFlow, false, null)
 
             if (!resolutionResults.isSuccess) return false
 

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -12341,6 +12341,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS_IR_ES6, true);
             }
 
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
@@ -12356,16 +12361,6 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");
@@ -12410,6 +12405,21 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -11762,6 +11762,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS_IR, true);
             }
 
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
@@ -11777,16 +11782,6 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");
@@ -11831,6 +11826,21 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -11827,6 +11827,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS, true);
             }
 
+            @TestMetadata("builderCallAsReturnTypeInLocalClass.kt")
+            public void testBuilderCallAsReturnTypeInLocalClass() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/builderCallAsReturnTypeInLocalClass.kt");
+            }
+
             @TestMetadata("callableReferenceAndCoercionToUnit.kt")
             public void testCallableReferenceAndCoercionToUnit() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
@@ -11842,16 +11847,6 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");
@@ -11896,6 +11891,21 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             public void testSubstitutelambdaExtensionReceiverType() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
             }
+
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }
 

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -6288,16 +6288,6 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
-            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
-            }
-
-            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
-            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
-                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
-            }
-
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");
@@ -6307,6 +6297,21 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
             public void testSubstituteStubTypeIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteStubTypeIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("topDownCompletionBreakedByNonBuilderInferenceSession.kt")
+            public void testTopDownCompletionBreakedByNonBuilderInferenceSession() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionBreakedByNonBuilderInferenceSession.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithThreeBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("topDownCompletionWithTwoBuilderInferenceCalls.kt")
+            public void testTopDownCompletionWithTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/topDownCompletionWithTwoBuilderInferenceCalls.kt");
+            }
         }
     }