Return "lazy" kind of nested classifier scope for special cases

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAbstractTreeTransformerWithSuperTypes.kt

@@ -10,15 +10,11 @@ import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.isCompanion
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.lookupSuperTypes
 import org.jetbrains.kotlin.fir.scopes.impl.FirCompositeScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirMemberTypeParameterScope
-import org.jetbrains.kotlin.fir.scopes.impl.FirNestedClassifierScope
 import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
-import org.jetbrains.kotlin.name.ClassId
 
 abstract class FirAbstractTreeTransformerWithSuperTypes(
     phase: FirResolvePhase,
@@ -37,11 +33,6 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
         return result
     }
 
-    private fun nestedClassifierScope(classId: ClassId): FirNestedClassifierScope? {
-        val classSymbol = session.firSymbolProvider.getClassLikeSymbolByFqName(classId) as? FirClassSymbol ?: return null
-        return nestedClassifierScope(classSymbol.fir)
-    }
-
     protected fun resolveNestedClassesSupertypes(
         regularClass: FirRegularClass,
         data: Nothing?
@@ -49,8 +40,8 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
         return withScopeCleanup {
             // ? Is it Ok to use original file session here ?
             lookupSuperTypes(regularClass, lookupInterfaces = false, deep = true, useSiteSession = session)
-                .asReversed().mapNotNullTo(towerScope.scopes) {
-                    nestedClassifierScope(it.lookupTag.classId)
+                .asReversed().mapTo(towerScope.scopes) {
+                    nestedClassifierScope(it.lookupTag.classId, session)
                 }
             val companionObjects = regularClass.declarations.filterIsInstance<FirRegularClass>().filter { it.isCompanion }
             for (companionObject in companionObjects) {

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.scopes.impl
 
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.FirCallableMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
@@ -16,6 +17,7 @@ import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction.NEXT
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction.STOP
 import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 class FirClassDeclaredMemberScopeProvider : FirSessionComponent {
@@ -50,6 +52,10 @@ fun nestedClassifierScope(klass: FirRegularClass): FirNestedClassifierScope {
         .nestedClassifierScope(klass)
 }
 
+fun nestedClassifierScope(classId: ClassId, session: FirSession): FirLazyNestedClassifierScope {
+    return FirLazyNestedClassifierScope(classId, session)
+}
+
 class FirClassDeclaredMemberScope(klass: FirRegularClass) : FirScope() {
     private val nestedClassifierScope = nestedClassifierScope(klass)
 

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.scopes.impl
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
+import org.jetbrains.kotlin.fir.scopes.FirPosition
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.ProcessorAction
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+
+class FirLazyNestedClassifierScope(
+    val classId: ClassId,
+    session: FirSession
+) : FirScope() {
+
+    private val symbolProvider = session.firSymbolProvider
+
+    override fun processClassifiersByName(
+        name: Name,
+        position: FirPosition,
+        processor: (FirClassifierSymbol<*>) -> ProcessorAction
+    ): ProcessorAction {
+        val child = classId.createNestedClassId(name)
+        val symbol = symbolProvider.getClassLikeSymbolByFqName(child) ?: return ProcessorAction.NONE
+
+        return processor(symbol)
+    }
+}