Propagate inference session into declaration analyzers

It prevents missing inference session for local declaration (local functions, local classes or objects)

^KT-44801 Fixed

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext;
 
 import java.util.Collection;
 import java.util.Map;
@@ -58,4 +59,7 @@ public interface BodiesResolveContext {
 
     @NotNull
     TopDownAnalysisMode getTopDownAnalysisMode();
+
+    @Nullable
+    ExpressionTypingContext getLocalContext();
 }

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

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil;
 import org.jetbrains.kotlin.resolve.multiplatform.ExpectedActualResolver;
 import org.jetbrains.kotlin.resolve.scopes.*;
 import org.jetbrains.kotlin.types.*;
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices;
 import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
 import org.jetbrains.kotlin.types.expressions.ValueParameterResolver;
@@ -134,7 +135,7 @@ public class BodyResolver {
         for (Map.Entry<KtSecondaryConstructor, ClassConstructorDescriptor> entry : c.getSecondaryConstructors().entrySet()) {
             LexicalScope declaringScope = c.getDeclaringScope(entry.getKey());
             assert declaringScope != null : "Declaring scope should be registered before body resolve";
-            resolveSecondaryConstructorBody(c.getOuterDataFlowInfo(), trace, entry.getKey(), entry.getValue(), declaringScope);
+            resolveSecondaryConstructorBody(c.getOuterDataFlowInfo(), trace, entry.getKey(), entry.getValue(), declaringScope, c.getLocalContext());
         }
         if (c.getSecondaryConstructors().isEmpty()) return;
         Set<ConstructorDescriptor> visitedConstructors = new HashSet<>();
@@ -148,18 +149,22 @@ public class BodyResolver {
             @NotNull BindingTrace trace,
             @NotNull KtSecondaryConstructor constructor,
             @NotNull ClassConstructorDescriptor descriptor,
-            @NotNull LexicalScope declaringScope
+            @NotNull LexicalScope declaringScope,
+            @Nullable ExpressionTypingContext localContext
     ) {
         ForceResolveUtil.forceResolveAllContents(descriptor.getAnnotations());
 
-        resolveFunctionBody(outerDataFlowInfo, trace, constructor, descriptor, declaringScope,
-                            headerInnerScope -> resolveSecondaryConstructorDelegationCall(
-                                    outerDataFlowInfo, trace, headerInnerScope, constructor, descriptor
-                            ),
-                            scope -> new LexicalScopeImpl(
-                                    scope, descriptor, scope.isOwnerDescriptorAccessibleByLabel(), scope.getImplicitReceiver(),
-                                    LexicalScopeKind.CONSTRUCTOR_HEADER
-                            ));
+        resolveFunctionBody(
+                outerDataFlowInfo, trace, constructor, descriptor, declaringScope,
+                headerInnerScope -> resolveSecondaryConstructorDelegationCall(
+                        outerDataFlowInfo, trace, headerInnerScope, constructor, descriptor
+                ),
+                scope -> new LexicalScopeImpl(
+                        scope, descriptor, scope.isOwnerDescriptorAccessibleByLabel(), scope.getImplicitReceiver(),
+                        LexicalScopeKind.CONSTRUCTOR_HEADER
+                ),
+                localContext
+        );
     }
 
     @Nullable
@@ -835,7 +840,7 @@ public class BodyResolver {
         if (getterDescriptor != null) {
             if (getter != null) {
                 LexicalScope accessorScope = makeScopeForPropertyAccessor(c, getter, propertyDescriptor);
-                resolveFunctionBody(c.getOuterDataFlowInfo(), fieldAccessTrackingTrace, getter, getterDescriptor, accessorScope);
+                resolveFunctionBody(c.getOuterDataFlowInfo(), fieldAccessTrackingTrace, getter, getterDescriptor, accessorScope, c.getLocalContext());
             }
 
             if (getter != null || forceResolveAnnotations) {
@@ -849,7 +854,7 @@ public class BodyResolver {
         if (setterDescriptor != null) {
             if (setter != null) {
                 LexicalScope accessorScope = makeScopeForPropertyAccessor(c, setter, propertyDescriptor);
-                resolveFunctionBody(c.getOuterDataFlowInfo(), fieldAccessTrackingTrace, setter, setterDescriptor, accessorScope);
+                resolveFunctionBody(c.getOuterDataFlowInfo(), fieldAccessTrackingTrace, setter, setterDescriptor, accessorScope, c.getLocalContext());
             }
 
             if (setter != null || forceResolveAnnotations) {
@@ -921,7 +926,7 @@ public class BodyResolver {
                 bodyResolveCache.resolveFunctionBody(declaration).addOwnDataTo(trace, true);
             }
             else {
-                resolveFunctionBody(c.getOuterDataFlowInfo(), trace, declaration, entry.getValue(), scope);
+                resolveFunctionBody(c.getOuterDataFlowInfo(), trace, declaration, entry.getValue(), scope, c.getLocalContext());
             }
         }
     }
@@ -931,11 +936,12 @@ public class BodyResolver {
             @NotNull BindingTrace trace,
             @NotNull KtDeclarationWithBody function,
             @NotNull FunctionDescriptor functionDescriptor,
-            @NotNull LexicalScope declaringScope
+            @NotNull LexicalScope declaringScope,
+            @Nullable ExpressionTypingContext localContext
     ) {
         computeDeferredType(functionDescriptor.getReturnType());
 
-        resolveFunctionBody(outerDataFlowInfo, trace, function, functionDescriptor, declaringScope, null, null);
+        resolveFunctionBody(outerDataFlowInfo, trace, function, functionDescriptor, declaringScope, null, null, localContext);
 
         assert functionDescriptor.getReturnType() != null;
     }
@@ -948,7 +954,8 @@ public class BodyResolver {
             @NotNull LexicalScope scope,
             @Nullable Function1<LexicalScope, DataFlowInfo> beforeBlockBody,
             // Creates wrapper scope for header resolution if necessary (see resolveSecondaryConstructorBody)
-            @Nullable Function1<LexicalScope, LexicalScope> headerScopeFactory
+            @Nullable Function1<LexicalScope, LexicalScope> headerScopeFactory,
+            @Nullable ExpressionTypingContext localContext
     ) {
         ProgressManager.checkCanceled();
 
@@ -989,7 +996,8 @@ public class BodyResolver {
 
         if (function.hasBody()) {
             expressionTypingServices.checkFunctionReturnType(
-                    innerScope, function, functionDescriptor, dataFlowInfo != null ? dataFlowInfo : outerDataFlowInfo, null, trace);
+                    innerScope, function, functionDescriptor, dataFlowInfo != null ? dataFlowInfo : outerDataFlowInfo, null, trace, localContext
+            );
         }
 
         assert functionDescriptor.getReturnType() != null;

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

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.resolve.checkers.checkClassifierUsages
 import org.jetbrains.kotlin.resolve.deprecation.DeprecationResolver
 import org.jetbrains.kotlin.resolve.lazy.*
 import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyClassDescriptor
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext
 import java.util.*
 
 class LazyTopDownAnalyzer(
@@ -57,9 +58,10 @@ class LazyTopDownAnalyzer(
     fun analyzeDeclarations(
         topDownAnalysisMode: TopDownAnalysisMode,
         declarations: Collection<PsiElement>,
-        outerDataFlowInfo: DataFlowInfo = DataFlowInfo.EMPTY
+        outerDataFlowInfo: DataFlowInfo = DataFlowInfo.EMPTY,
+        localContext: ExpressionTypingContext? = null
     ): TopDownAnalysisContext {
-        val c = TopDownAnalysisContext(topDownAnalysisMode, outerDataFlowInfo, declarationScopeProvider)
+        val c = TopDownAnalysisContext(topDownAnalysisMode, outerDataFlowInfo, declarationScopeProvider, localContext)
 
         val topLevelFqNames = HashMultimap.create<FqName, KtElement>()
 

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

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.lazy.DeclarationScopeProvider;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext;
 
 import java.io.PrintStream;
 import java.util.*;
@@ -49,6 +50,7 @@ public class TopDownAnalysisContext implements BodiesResolveContext {
 
     private final TopDownAnalysisMode topDownAnalysisMode;
     private final DeclarationScopeProvider declarationScopeProvider;
+    private final ExpressionTypingContext localContext;
 
     private StringBuilder debugOutput;
 
@@ -60,6 +62,25 @@ public class TopDownAnalysisContext implements BodiesResolveContext {
         this.topDownAnalysisMode = topDownAnalysisMode;
         this.outerDataFlowInfo = outerDataFlowInfo;
         this.declarationScopeProvider = declarationScopeProvider;
+        this.localContext = null;
+    }
+
+    public TopDownAnalysisContext(
+            @NotNull TopDownAnalysisMode topDownAnalysisMode,
+            @NotNull DataFlowInfo outerDataFlowInfo,
+            @NotNull DeclarationScopeProvider declarationScopeProvider,
+            @Nullable ExpressionTypingContext localContext
+    ) {
+        this.topDownAnalysisMode = topDownAnalysisMode;
+        this.outerDataFlowInfo = outerDataFlowInfo;
+        this.declarationScopeProvider = declarationScopeProvider;
+        this.localContext = localContext;
+    }
+
+    @Override
+    @Nullable
+    public ExpressionTypingContext getLocalContext() {
+        return localContext;
     }
 
     @Override

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

@@ -145,7 +145,8 @@ public class ExpressionTypingServices {
             @NotNull FunctionDescriptor functionDescriptor,
             @NotNull DataFlowInfo dataFlowInfo,
             @Nullable KotlinType expectedReturnType,
-            BindingTrace trace
+            BindingTrace trace,
+            @Nullable ExpressionTypingContext localContext
     ) {
         if (expectedReturnType == null) {
             expectedReturnType = functionDescriptor.getReturnType();
@@ -153,11 +154,15 @@ public class ExpressionTypingServices {
                 expectedReturnType = NO_EXPECTED_TYPE;
             }
         }
-        checkFunctionReturnType(function, ExpressionTypingContext.newContext(
+
+        ExpressionTypingContext context = ExpressionTypingContext.newContext(
                 trace,
-                functionInnerScope, dataFlowInfo, expectedReturnType != null ? expectedReturnType : NO_EXPECTED_TYPE, getLanguageVersionSettings(),
-                expressionTypingComponents.dataFlowValueFactory
-        ));
+                functionInnerScope, dataFlowInfo, expectedReturnType != null ? expectedReturnType : NO_EXPECTED_TYPE,
+                getLanguageVersionSettings(), expressionTypingComponents.dataFlowValueFactory,
+                localContext == null ? InferenceSession.Companion.getDefault() : localContext.inferenceSession
+        );
+
+        checkFunctionReturnType(function, context);
     }
 
     /*package*/ void checkFunctionReturnType(KtDeclarationWithBody function, ExpressionTypingContext context) {

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

@@ -100,7 +100,7 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
             ForceResolveUtil.forceResolveAllContents(functionDescriptor.returnType)
         } else {
             components.expressionTypingServices.checkFunctionReturnType(
-                functionInnerScope, function, functionDescriptor, context.dataFlowInfo, null, context.trace
+                functionInnerScope, function, functionDescriptor, context.dataFlowInfo, null, context.trace, context
             )
         }
 

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

@@ -116,7 +116,8 @@ class LocalClassifierAnalyzer(
         container.get<LazyTopDownAnalyzer>().analyzeDeclarations(
             TopDownAnalysisMode.LocalDeclarations,
             listOf(classOrObject),
-            context.dataFlowInfo
+            context.dataFlowInfo,
+            localContext = context
         )
     }
 }