FIR: Fix inner type resolution during body transformation

compiler/fir/analysis-tests/testData/resolve/nested/innerTypeFromSuperClassInBody.kt

@@ -0,0 +1,15 @@
+abstract class A<X : CharSequence> {
+    inner class Inner
+    fun foo(x: Inner.() -> Unit) {}
+}
+
+object B : A<String>() {
+
+    fun bar() {
+        val y: Inner.() -> Unit = {}
+        foo(y)
+        baz(y)
+    }
+}
+
+fun baz(x: (A<String>.Inner) -> Unit) {}

compiler/fir/analysis-tests/testData/resolve/nested/innerTypeFromSuperClassInBody.txt

@@ -0,0 +1,34 @@
+FILE: innerTypeFromSuperClassInBody.kt
+    public abstract class A<X : R|kotlin/CharSequence|> : R|kotlin/Any| {
+        public constructor<X : R|kotlin/CharSequence|>(): R|A<X>| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final inner class Inner<X : R|kotlin/CharSequence|> : R|kotlin/Any| {
+            public constructor(): R|A.Inner<X>| {
+                super<R|kotlin/Any|>()
+            }
+
+        }
+
+        public final fun foo(x: R|A.Inner<X>.() -> kotlin/Unit|): R|kotlin/Unit| {
+        }
+
+    }
+    public final object B : R|A<kotlin/String>| {
+        private constructor(): R|B| {
+            super<R|A<kotlin/String>|>()
+        }
+
+        public final fun bar(): R|kotlin/Unit| {
+            lval y: R|A.Inner<kotlin/String>.() -> kotlin/Unit| = fun R|A.Inner<kotlin/String>|.<anonymous>(): R|kotlin/Unit| {
+                ^ Unit
+            }
+
+            this@R|/B|.R|SubstitutionOverride</B.foo: R|kotlin/Unit|>|(R|<local>/y|)
+            R|/baz|(R|<local>/y|)
+        }
+
+    }
+    public final fun baz(x: R|(A.Inner<kotlin/String>) -> kotlin/Unit|): R|kotlin/Unit| {
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -1923,6 +1923,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/nested/inner.kt");
         }
 
+        @TestMetadata("innerTypeFromSuperClassInBody.kt")
+        public void testInnerTypeFromSuperClassInBody() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/nested/innerTypeFromSuperClassInBody.kt");
+        }
+
         @TestMetadata("innerTypes.kt")
         public void testInnerTypes() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/nested/innerTypes.kt");

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -1923,6 +1923,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/nested/inner.kt");
         }
 
+        @TestMetadata("innerTypeFromSuperClassInBody.kt")
+        public void testInnerTypeFromSuperClassInBody() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/nested/innerTypeFromSuperClassInBody.kt");
+        }
+
         @TestMetadata("innerTypes.kt")
         public void testInnerTypes() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/nested/innerTypes.kt");

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -1923,6 +1923,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/nested/inner.kt");
         }
 
+        @TestMetadata("innerTypeFromSuperClassInBody.kt")
+        public void testInnerTypeFromSuperClassInBody() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/nested/innerTypeFromSuperClassInBody.kt");
+        }
+
         @TestMetadata("innerTypes.kt")
         public void testInnerTypes() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/nested/innerTypes.kt");

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
+import org.jetbrains.kotlin.fir.scopes.impl.wrapNestedClassifierScopeWithSubstitutionForSuperType
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -194,9 +195,13 @@ fun SessionHolder.collectTowerDataElementsForClass(owner: FirClass<*>, defaultTy
 
     val superClassesStaticsAndCompanionReceivers = mutableListOf<FirTowerDataElement>()
     for (superType in lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session)) {
-        val superClass = superType.fullyExpandedType(session).lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: continue
+        val expandedType = superType.fullyExpandedType(session)
+        val superClass = expandedType.lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: continue
 
-        superClass.staticScope(this)?.asTowerDataElement(isLocal = false)?.let(superClassesStaticsAndCompanionReceivers::add)
+        superClass.staticScope(this)
+            ?.wrapNestedClassifierScopeWithSubstitutionForSuperType(expandedType, session)
+            ?.asTowerDataElement(isLocal = false)
+            ?.let(superClassesStaticsAndCompanionReceivers::add)
 
         (superClass as? FirRegularClass)?.companionObject?.let { companion ->
             val superCompanionReceiver = ImplicitDispatchReceiverValue(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirNestedClassifierScope.kt

@@ -43,11 +43,9 @@ class FirNestedClassifierScope(val klass: FirClass<*>) : FirScope(), FirContaini
         processor(matchedClass, ConeSubstitutorByMap(substitution))
     }
 
-    fun getClassifierByName(name: Name): FirRegularClassSymbol? = classIndex[name]
-
     override fun getClassifierNames(): Set<Name> = classIndex.keys
 
     override fun getCallableNames(): Set<Name> = emptySet()
 }
 
-fun FirTypeParameterRef.toConeType(): ConeKotlinType = ConeTypeParameterTypeImpl(ConeTypeParameterLookupTag(symbol), isNullable = false)
+fun FirTypeParameterRef.toConeType(): ConeKotlinType = ConeTypeParameterTypeImpl(ConeTypeParameterLookupTag(symbol), isNullable = false)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirNestedClassifierScopeWithSubstitution.kt

@@ -10,16 +10,34 @@ import org.jetbrains.kotlin.fir.declarations.isInner
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.createSubstitutionForSupertype
 import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
+import org.jetbrains.kotlin.fir.scopes.getSingleClassifier
+import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.name.Name
 
-class FirNestedClassifierScopeWithSubstitution(
+private class FirNestedClassifierScopeWithSubstitution(
-    private val scope: FirNestedClassifierScope,
+    private val scope: FirScope,
     private val substitutor: ConeSubstitutor
 ) : FirScope() {
+
+    override fun processFunctionsByName(name: Name, processor: (FirFunctionSymbol<*>) -> Unit) {
+        scope.processFunctionsByName(name, processor)
+    }
+
+    override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
+        scope.processPropertiesByName(name, processor)
+    }
+
+    override fun processDeclaredConstructors(processor: (FirConstructorSymbol) -> Unit) {
+        scope.processDeclaredConstructors(processor)
+    }
+
+    override fun mayContainName(name: Name): Boolean {
+        return scope.mayContainName(name)
+    }
+
     override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
-        val matchedClass = scope.getClassifierByName(name) ?: return
+        val matchedClass = scope.getSingleClassifier(name) as? FirRegularClassSymbol ?: return
         val substitutor = substitutor.takeIf { matchedClass.fir.isInner } ?: ConeSubstitutor.Empty
         processor(matchedClass, substitutor)
     }
@@ -28,9 +46,7 @@ class FirNestedClassifierScopeWithSubstitution(
 fun FirScope.wrapNestedClassifierScopeWithSubstitutionForSuperType(
     superType: ConeClassLikeType,
     session: FirSession
-): FirScope = if (this is FirNestedClassifierScope) {
+): FirScope {
     val substitutor = createSubstitutionForSupertype(superType, session)
-    FirNestedClassifierScopeWithSubstitution(this, substitutor)
+    return FirNestedClassifierScopeWithSubstitution(this, substitutor)
-} else {
+}
-    this
-}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/scopes/FirScope.kt

@@ -35,6 +35,10 @@ abstract class FirScope {
     open fun mayContainName(name: Name) = true
 }
 
+fun FirScope.getSingleClassifier(name: Name): FirClassifierSymbol<*>? = mutableListOf<FirClassifierSymbol<*>>().apply {
+    processClassifiersByName(name, this::add)
+}.singleOrNull()
+
 fun FirScope.getFunctions(name: Name): List<FirFunctionSymbol<*>> = mutableListOf<FirFunctionSymbol<*>>().apply {
     processFunctionsByName(name, this::add)
 }