[FE 1.0] Take care standalone lambdas during updating types in the builder inference

^KT-50520 Fixed

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

@@ -19558,6 +19558,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @Test
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {

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

@@ -190,9 +190,6 @@ class LocalVariableResolver(
             initializeWithDefaultGetterSetter(propertyDescriptor)
             trace.record(BindingContext.VARIABLE, variable, propertyDescriptor)
             result = propertyDescriptor
-            if (inferenceSession is BuilderInferenceSession) {
-                inferenceSession.addLocalVariable(variable)
-            }
         } else {
             val variableDescriptor = resolveLocalVariableDescriptorWithType(scope, variable, null, trace)
             // For a local variable the type must not be deferred
@@ -200,11 +197,11 @@ class LocalVariableResolver(
                 variableDescriptor, scope, variable, dataFlowInfo, inferenceSession, trace, local = true
             )
             variableDescriptor.setOutType(type)
-            if (inferenceSession is BuilderInferenceSession) {
-                inferenceSession.addLocalVariable(variable)
-            }
             result = variableDescriptor
         }
+        if (inferenceSession is BuilderInferenceSession) {
+            inferenceSession.addExpression(variable)
+        }
         variableTypeAndInitializerResolver
             .setConstantForVariableIfNeeded(result, scope, variable, dataFlowInfo, type, inferenceSession, trace)
         // Type annotations also should be resolved

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

@@ -8,15 +8,15 @@ package org.jetbrains.kotlin.resolve.calls.inference
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
-import org.jetbrains.kotlin.descriptors.impl.*
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
+import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
-import org.jetbrains.kotlin.resolve.calls.util.shouldBeSubstituteWithStubTypes
-import org.jetbrains.kotlin.resolve.calls.util.toOldSubstitution
 import org.jetbrains.kotlin.resolve.calls.components.*
 import org.jetbrains.kotlin.resolve.calls.components.candidate.ResolutionCandidate
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
@@ -25,6 +25,8 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.calls.tower.*
 import org.jetbrains.kotlin.resolve.calls.util.FakeCallableDescriptorForObject
+import org.jetbrains.kotlin.resolve.calls.util.shouldBeSubstituteWithStubTypes
+import org.jetbrains.kotlin.resolve.calls.util.toOldSubstitution
 import org.jetbrains.kotlin.resolve.deprecation.DeprecationResolver
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasBuilderInferenceAnnotation
 import org.jetbrains.kotlin.resolve.descriptorUtil.shouldBeSubstituteWithStubTypes
@@ -73,11 +75,7 @@ class BuilderInferenceSession(
     }
 
     private val commonCalls = arrayListOf<PSICompletedCallInfo>()
-
+    private val commonExpressions = arrayListOf<KtExpression>()
-    private val localVariables = arrayListOf<KtVariableDeclaration>()
-
-    // These calls come from the old type inference
-    private val oldDoubleColonExpressionCalls = arrayListOf<KtExpression>()
 
     private var hasInapplicableCall = false
 
@@ -130,10 +128,6 @@ class BuilderInferenceSession(
         }
     }
 
-    fun addOldCallableReferenceCalls(callExpression: KtExpression) {
-        oldDoubleColonExpressionCalls.add(callExpression)
-    }
-
     override fun addCompletedCallInfo(callInfo: CompletedCallInfo) {
         require(callInfo is PSICompletedCallInfo) { "Wrong instance of callInfo: $callInfo" }
 
@@ -158,8 +152,8 @@ class BuilderInferenceSession(
         }
     }
 
-    fun addLocalVariable(variable: KtVariableDeclaration) {
+    fun addExpression(expression: KtExpression) {
-        localVariables.add(variable)
+        commonExpressions.add(expression)
     }
 
     private fun anyReceiverContainStubType(descriptor: CallableDescriptor): Boolean {
@@ -180,7 +174,6 @@ class BuilderInferenceSession(
         return null
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
     private fun findAllParentBuildInferenceSessions() = buildList {
         var currentSession: BuilderInferenceSession? = findParentBuildInferenceSession()
 
@@ -259,7 +252,7 @@ class BuilderInferenceSession(
         kotlinConstraintSystemCompleter.completeConstraintSystem(
             commonSystem.asConstraintSystemCompleterContext(),
             builtIns.unitType,
-            partiallyResolvedCallsInfo.map { it.callResolutionResult.resultCallAtom },
+            commonPartiallyResolvedCalls.map { it.callResolutionResult.resultCallAtom },
             completionMode,
             diagnosticsHolder
         )
@@ -277,7 +270,6 @@ class BuilderInferenceSession(
         return commonSystem.fixedTypeVariables.cast() // TODO: SUB
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
     private fun getNestedBuilderInferenceSessions() = buildList {
         for (nestedSession in nestedInferenceSessions) {
             when (nestedSession) {
@@ -448,13 +440,9 @@ class BuilderInferenceSession(
 
         integrateConstraints(initialStorage, nonFixedToVariablesSubstitutor, false)
 
-        for (call in commonCalls) {
+        for (call in commonCalls + commonPartiallyResolvedCalls) {
             val storage = call.callResolutionResult.constraintSystem.getBuilder().currentStorage()
-            integrateConstraints(storage, nonFixedToVariablesSubstitutor, false)
+            integrateConstraints(storage, nonFixedToVariablesSubstitutor, shouldIntegrateAllConstraints = call is PSIPartialCallInfo)
-        }
-        for (call in partiallyResolvedCallsInfo) {
-            val storage = call.callResolutionResult.constraintSystem.getBuilder().currentStorage()
-            integrateConstraints(storage, nonFixedToVariablesSubstitutor, true)
         }
 
         return commonSystem.notFixedTypeVariables.all { it.value.constraints.isEmpty() }
@@ -470,10 +458,31 @@ class BuilderInferenceSession(
         )
     }
 
-    private fun updateLocalVariable(localVariable: KtVariableDeclaration, substitutor: NewTypeSubstitutor) {
+    private fun updateExpressionDescriptorAndType(expression: KtExpression, substitutor: NewTypeSubstitutor) {
-        val descriptor = trace[BindingContext.VARIABLE, localVariable] as? LocalVariableDescriptor
+        val currentExpressionType = trace.getType(expression)
-        if (descriptor != null && descriptor.type.shouldBeUpdated()) {
+        if (currentExpressionType != null) {
-            descriptor.setOutType(substitutor.safeSubstitute(descriptor.type.unwrap()))
+            trace.recordType(expression, substitutor.safeSubstitute(currentExpressionType.unwrap()))
+        }
+
+        val (currentDescriptorType, updateDescriptorType) = when (expression) {
+            is KtLambdaExpression -> {
+                val descriptor = trace[BindingContext.FUNCTION, expression.functionLiteral] as? AnonymousFunctionDescriptor ?: return
+                val currentType = descriptor.returnType ?: return
+                currentType to descriptor::setReturnType
+            }
+            is KtVariableDeclaration -> {
+                val descriptor = trace[BindingContext.VARIABLE, expression] as? LocalVariableDescriptor ?: return
+                descriptor.type to descriptor::setOutType
+            }
+            is KtDoubleColonExpression -> {
+                completeDoubleColonExpression(expression, substitutor)
+                return
+            }
+            else -> return
+        }
+
+        if (currentDescriptorType.shouldBeUpdated()) {
+            updateDescriptorType(substitutor.safeSubstitute(currentDescriptorType.unwrap()))
         }
     }
 
@@ -513,12 +522,6 @@ class BuilderInferenceSession(
             atomCompleter.substituteFunctionLiteralDescriptor(resolvedAtom = null, descriptor = declarationDescriptor, substitutor)
         }
 
-        val recordedType = trace.getType(expression)
-
-        if (recordedType != null) {
-            trace.recordType(expression, substitutor.safeSubstitute(recordedType.unwrap()))
-        }
-
         val targetExpression = when (expression) {
             is KtCallableReferenceExpression -> expression.callableReference
             is KtClassLiteralExpression -> expression.receiverExpression
@@ -610,25 +613,18 @@ class BuilderInferenceSession(
                 topLevelCallContext.replaceBindingTrace(findTopLevelTrace()).replaceInferenceSession(this)
             )
 
-            for (localVariable in localVariables) {
+            for (expression in commonExpressions) {
-                updateLocalVariable(localVariable, nonFixedTypesToResultSubstitutor)
+                updateExpressionDescriptorAndType(expression, nonFixedTypesToResultSubstitutor)
             }
 
-            for (completedCall in commonCalls) {
+            for (call in commonCalls) {
-                updateCall(completedCall, nonFixedTypesToResultSubstitutor, nonFixedTypesToResult)
+                updateCall(call, nonFixedTypesToResultSubstitutor, nonFixedTypesToResult)
-                reportErrors(completedCall, completedCall.resolvedCall, errors)
+                reportErrors(call, call.resolvedCall, errors)
             }
 
-            for (callInfo in partiallyResolvedCallsInfo) {
+            for (call in commonPartiallyResolvedCalls) {
-                val resolvedCall = completeCall(callInfo, atomCompleter) ?: continue
+                val resolvedCall = completeCall(call, atomCompleter) ?: continue
-                reportErrors(callInfo, resolvedCall, errors)
+                reportErrors(call, resolvedCall, errors)
-            }
-
-            for (call in oldDoubleColonExpressionCalls) {
-                when (call) {
-                    is KtDoubleColonExpression -> completeDoubleColonExpression(call, nonFixedTypesToResultSubstitutor)
-                    else -> throw Exception("Unsupported call expression type")
-                }
             }
 
             atomCompleter.completeAll(lambda)

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/StubTypesBasedInferenceSession.kt

@@ -28,7 +28,7 @@ abstract class StubTypesBasedInferenceSession<D : CallableDescriptor>(
     protected val callComponents: KotlinCallComponents,
     val builtIns: KotlinBuiltIns
 ) : InferenceSession {
-    protected val partiallyResolvedCallsInfo = arrayListOf<PSIPartialCallInfo>()
+    protected val commonPartiallyResolvedCalls = arrayListOf<PSIPartialCallInfo>()
     val errorCallsInfo = arrayListOf<PSIErrorCallInfo<D>>()
     private val completedCalls = hashSetOf<ResolvedAtom>()
     protected val nestedInferenceSessions = hashSetOf<StubTypesBasedInferenceSession<*>>()
@@ -49,7 +49,7 @@ abstract class StubTypesBasedInferenceSession<D : CallableDescriptor>(
         if (callInfo !is PSIPartialCallInfo) {
             throw AssertionError("Call info for $callInfo should be instance of PSIPartialCallInfo")
         }
-        partiallyResolvedCallsInfo.add(callInfo)
+        commonPartiallyResolvedCalls.add(callInfo)
     }
 
     override fun addCompletedCallInfo(callInfo: CompletedCallInfo) {
@@ -65,7 +65,7 @@ abstract class StubTypesBasedInferenceSession<D : CallableDescriptor>(
     }
 
     override fun currentConstraintSystem(): ConstraintStorage {
-        return partiallyResolvedCallsInfo.lastOrNull()?.callResolutionResult?.constraintSystem?.getBuilder()?.currentStorage()
+        return commonPartiallyResolvedCalls.lastOrNull()?.callResolutionResult?.constraintSystem?.getBuilder()?.currentStorage()
             ?: ConstraintStorage.Empty
     }
 
@@ -75,7 +75,7 @@ abstract class StubTypesBasedInferenceSession<D : CallableDescriptor>(
     override fun shouldCompleteResolvedSubAtomsOf(resolvedCallAtom: ResolvedCallAtom) = true
 
     fun resolveCandidates(resolutionCallbacks: KotlinResolutionCallbacks): List<ResolutionResultCallInfo<D>> {
-        val resolvedCallsInfo = partiallyResolvedCallsInfo.toList()
+        val resolvedCallsInfo = commonPartiallyResolvedCalls.toList()
 
         val diagnosticHolder = KotlinDiagnosticsHolder.SimpleHolder()
 

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

@@ -111,7 +111,7 @@ class DoubleColonExpressionResolver(
             val result = resolveDoubleColonLHS(expression, c)
 
             if (c.inferenceSession is BuilderInferenceSession && result?.type?.contains { it is StubTypeForBuilderInference } == true) {
-                c.inferenceSession.addOldCallableReferenceCalls(expression)
+                c.inferenceSession.addExpression(expression)
             }
 
             if (result != null && !result.type.isError) {
@@ -557,7 +557,7 @@ class DoubleColonExpressionResolver(
         val dataFlowInfo = (lhs as? DoubleColonLHS.Expression)?.dataFlowInfo ?: c.dataFlowInfo
 
         if (c.inferenceSession is BuilderInferenceSession && result?.contains { it is StubTypeForBuilderInference } == true) {
-            c.inferenceSession.addOldCallableReferenceCalls(expression)
+            c.inferenceSession.addExpression(expression)
         }
 
         return dataFlowAnalyzer.checkType(createTypeInfo(result, dataFlowInfo), expression, c)

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getAnnotationEntries
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.BindingContext.EXPECTED_RETURN_TYPE
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
+import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSession
 import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableTypeConstructor
 import org.jetbrains.kotlin.resolve.checkers.TrailingCommaChecker
 import org.jetbrains.kotlin.resolve.checkers.UnderscoreChecker
@@ -182,6 +183,10 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
             functionDescriptor.createFunctionType(components.builtIns, suspendFunctionTypeExpected)!!
         )
 
+        if (context.inferenceSession is BuilderInferenceSession) {
+            context.inferenceSession.addExpression(expression)
+        }
+
         if (functionTypeExpected) {
             // all checks were done before
             return createTypeInfo(resultType, context)

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

@@ -0,0 +1,11 @@
+// WITH_STDLIB
+// IGNORE_BACKEND_FIR: JVM_IR
+// FIR status: different behavour with FE 1.0, reported `NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER`
+
+fun box(): String {
+    buildList {
+        val foo = { first() }
+        add(0, foo)
+    }
+    return "OK"
+}

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

@@ -19096,6 +19096,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @Test
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {

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

@@ -19558,6 +19558,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @Test
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {

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

@@ -15868,6 +15868,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");

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

@@ -14856,6 +14856,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @Test
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {

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

@@ -14820,6 +14820,12 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @Test
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {

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

@@ -13188,6 +13188,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @TestMetadata("kt50520.kt")
+            public void testKt50520() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+            }
+
             @TestMetadata("kt51988.kt")
             public void testKt51988() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");

native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java

@@ -15948,6 +15948,12 @@ public class NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTest
                     runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
                 }
 
+                @Test
+                @TestMetadata("kt50520.kt")
+                public void testKt50520() throws Exception {
+                    runTest("compiler/testData/codegen/box/inference/builderInference/kt50520.kt");
+                }
+
                 @Test
                 @TestMetadata("kt51988.kt")
                 public void testKt51988() throws Exception {