Get nested builder inference sessions properly, across delegated inference sessions too

^KT-48445 Fixed

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptorWithAccessors
 import org.jetbrains.kotlin.resolve.calls.components.*
+import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSession
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
@@ -37,6 +38,18 @@ class DelegatedPropertyInferenceSession(
 ) : ManyCandidatesResolver<FunctionDescriptor>(
     psiCallResolver, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents, builtIns
 ) {
+    init {
+        if (parentSession is ManyCandidatesResolver<*>) {
+            parentSession.addNestedInferenceSession(this)
+        }
+    }
+
+    fun getNestedBuilderInferenceSessions(): List<BuilderInferenceSession> {
+        val builderInferenceSessions = nestedInferenceSessions.filterIsInstance<BuilderInferenceSession>()
+        val delegatedPropertyInferenceSessions = nestedInferenceSessions.filterIsInstance<DelegatedPropertyInferenceSession>()
+
+        return builderInferenceSessions + delegatedPropertyInferenceSessions.map { it.getNestedBuilderInferenceSessions() }.flatten()
+    }
 
     override fun prepareForCompletion(commonSystem: NewConstraintSystem, resolvedCallsInfo: List<PSIPartialCallInfo>) {
         val csBuilder = commonSystem.getBuilder()

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.resolve.calls.inference
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.impl.*
@@ -60,14 +59,12 @@ class BuilderInferenceSession(
 ) : ManyCandidatesResolver<CallableDescriptor>(
     psiCallResolver, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents, builtIns
 ) {
-    private var nestedBuilderInferenceSessions: MutableSet<BuilderInferenceSession> = mutableSetOf()
-
     private lateinit var lambda: ResolvedLambdaAtom
     private val commonSystem = NewConstraintSystemImpl(callComponents.constraintInjector, builtIns, callComponents.kotlinTypeRefiner)
 
     init {
-        if (topLevelCallContext.inferenceSession is BuilderInferenceSession) {
-            topLevelCallContext.inferenceSession.nestedBuilderInferenceSessions.add(this)
+        if (topLevelCallContext.inferenceSession is ManyCandidatesResolver<*>) {
+            topLevelCallContext.inferenceSession.addNestedInferenceSession(this)
         }
         stubsForPostponedVariables.keys.forEach(commonSystem::registerVariable)
     }
@@ -266,6 +263,16 @@ class BuilderInferenceSession(
         return commonSystem.fixedTypeVariables.cast() // TODO: SUB
     }
 
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun getNestedBuilderInferenceSessions() = buildList {
+        for (nestedSession in nestedInferenceSessions) {
+            when (nestedSession) {
+                is BuilderInferenceSession -> add(nestedSession)
+                is DelegatedPropertyInferenceSession -> addAll(nestedSession.getNestedBuilderInferenceSessions())
+            }
+        }
+    }
+
     /*
      * We update calls in top-down way:
      * - updating calls within top-level builder inference call
@@ -279,6 +286,8 @@ class BuilderInferenceSession(
             commonSystem.errors
         )
 
+        val nestedBuilderInferenceSessions = getNestedBuilderInferenceSessions()
+
         for (nestedSession in nestedBuilderInferenceSessions) {
             // TODO: exclude injected variables
             nestedSession.updateAllCalls(

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

@@ -30,6 +30,11 @@ abstract class ManyCandidatesResolver<D : CallableDescriptor>(
     protected val partiallyResolvedCallsInfo = arrayListOf<PSIPartialCallInfo>()
     private val errorCallsInfo = arrayListOf<PSIErrorCallInfo<D>>()
     private val completedCalls = hashSetOf<ResolvedAtom>()
+    protected val nestedInferenceSessions = hashSetOf<ManyCandidatesResolver<*>>()
+
+    fun addNestedInferenceSession(inferenceSession: ManyCandidatesResolver<*>) {
+        nestedInferenceSessions.add(inferenceSession)
+    }
 
     open fun prepareForCompletion(commonSystem: NewConstraintSystem, resolvedCallsInfo: List<PSIPartialCallInfo>) {
         // do nothing

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

@@ -14,39 +14,39 @@ internal class TowerDataElementsForName() {
     }
 }
 
-//internal class TowerDataElementsForName2() {
-//    @OptIn(ExperimentalStdlibApi::class)
-//    val reversedFilteredLocalScopes = buildList {
-//        val reversedFilteredLocalScopes by lazy(LazyThreadSafetyMode.NONE) {
-//            @OptIn(ExperimentalStdlibApi::class)
-//            buildList {
-//                for (i in lastIndex downTo 0) {
-//                    add("")
-//                }
-//            }
-//        }
-//        add(reversedFilteredLocalScopes)
-//    }
-//}
-//
-//internal class TowerDataElementsForName3() {
-//    val reversedFilteredLocalScopes by lazy(LazyThreadSafetyMode.NONE) {
-//        @OptIn(ExperimentalStdlibApi::class)
-//        buildList l1@ {
-//            for (i in lastIndex downTo 0) {
-//                val reversedFilteredLocalScopes by lazy(LazyThreadSafetyMode.NONE) {
-//                    @OptIn(ExperimentalStdlibApi::class)
-//                    buildList {
-//                        for (i in lastIndex downTo 0) {
-//                            add("")
-//                            this@l1.add("")
-//                        }
-//                    }
-//                }
-//            }
-//        }
-//    }
-//}
+internal class TowerDataElementsForName2() {
+    @OptIn(ExperimentalStdlibApi::class)
+    val reversedFilteredLocalScopes = buildList {
+        val reversedFilteredLocalScopes by lazy(LazyThreadSafetyMode.NONE) {
+            @OptIn(ExperimentalStdlibApi::class)
+            buildList {
+                for (i in lastIndex downTo 0) {
+                    add("")
+                }
+            }
+        }
+        add(reversedFilteredLocalScopes)
+    }
+}
+
+internal class TowerDataElementsForName3() {
+    val reversedFilteredLocalScopes by lazy(LazyThreadSafetyMode.NONE) {
+        @OptIn(ExperimentalStdlibApi::class)
+        buildList l1@ {
+            for (i in lastIndex downTo 0) {
+                val reversedFilteredLocalScopes by lazy(LazyThreadSafetyMode.NONE) {
+                    @OptIn(ExperimentalStdlibApi::class)
+                    buildList {
+                        for (i in lastIndex downTo 0) {
+                            add("")
+                            this@l1.add("")
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
 
 // mute due to KT-48633
 //internal class TowerDataElementsForName4() {
@@ -68,8 +68,8 @@ internal class TowerDataElementsForName() {
 
 fun box(): String {
     val x1 = TowerDataElementsForName().reversedFilteredLocalScopes
-//    val x2 = TowerDataElementsForName2().reversedFilteredLocalScopes
-//    val x3 = TowerDataElementsForName3().reversedFilteredLocalScopes
+    val x2 = TowerDataElementsForName2().reversedFilteredLocalScopes
+    val x3 = TowerDataElementsForName3().reversedFilteredLocalScopes
 //    val x4 = TowerDataElementsForName4().reversedFilteredLocalScopes
     return "OK"
 }