FIR: Extract LocalClassesNavigationInfo computation

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -7,7 +7,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
-import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
@@ -27,7 +26,6 @@ import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
-import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
 
@@ -71,9 +69,8 @@ class FirImplicitTypeBodyResolveTransformerAdapterForLocalClasses(
     }
 
     override fun <F : FirClass<F>> transformClass(klass: FirClass<F>, data: Nothing?): CompositeTransformResult<FirStatement> {
-        val (designationMap, targetedClasses) = CollectingVisitor().run {
+        val (designationMap, targetedClasses) = klass.collectLocalClassesNavigationInfo().run {
-            visitClass(klass)
+            designationMap to parentForClass.keys + klass
-            resultingMap to targetedClasses
         }
 
         val implicitBodyResolveComputationSession =
@@ -102,47 +99,6 @@ class FirImplicitTypeBodyResolveTransformerAdapterForLocalClasses(
     }
 }
 
-private class CollectingVisitor : FirDefaultVisitorVoid() {
-    val resultingMap = mutableMapOf<FirCallableMemberDeclaration<*>, List<FirClass<*>>>()
-    val targetedClasses = mutableSetOf<FirClass<*>>()
-    private var currentPath = persistentListOf<FirClass<*>>()
-
-    override fun visitElement(element: FirElement) {}
-
-    override fun visitRegularClass(regularClass: FirRegularClass) {
-        visitClass(regularClass)
-    }
-
-    override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) {
-        visitClass(anonymousObject)
-    }
-
-    override fun <F : FirClass<F>> visitClass(klass: FirClass<F>) {
-        targetedClasses.add(klass)
-        val prev = currentPath
-        currentPath = currentPath.add(klass)
-
-        klass.acceptChildren(this)
-
-        currentPath = prev
-    }
-
-    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) {
-        visitCallableMemberDeclaration(simpleFunction)
-    }
-
-    override fun visitProperty(property: FirProperty) {
-        visitCallableMemberDeclaration(property)
-    }
-
-    override fun <F : FirCallableMemberDeclaration<F>> visitCallableMemberDeclaration(
-        callableMemberDeclaration: FirCallableMemberDeclaration<F>
-    ) {
-        if (callableMemberDeclaration.returnTypeRef !is FirImplicitTypeRef) return
-        resultingMap[callableMemberDeclaration] = currentPath
-    }
-}
-
 fun createReturnTypeCalculatorForIDE(session: FirSession, scopeSession: ScopeSession): ReturnTypeCalculator =
     ReturnTypeCalculatorWithJump(session, scopeSession, ImplicitBodyResolveComputationSession())
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/LocalClassesNavigation.kt

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2020 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.resolve.transformers.body.resolve
+
+import kotlinx.collections.immutable.persistentListOf
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
+import org.jetbrains.kotlin.utils.keysToMap
+
+class LocalClassesNavigationInfo(
+    val parentForClass: Map<FirClass<*>, FirClass<*>?>,
+    private val parentClassForFunction: Map<FirCallableMemberDeclaration<*>, FirClass<*>>
+) {
+    val designationMap by lazy {
+        parentClassForFunction.keys.keysToMap {
+            pathForCallable(it)
+        }
+    }
+
+    private fun pathForCallable(callableMemberDeclaration: FirCallableMemberDeclaration<*>): List<FirClass<*>> {
+        val result = mutableListOf<FirClass<*>>()
+        var current = parentClassForFunction[callableMemberDeclaration]
+
+        while (current != null) {
+            result += current
+            current = parentForClass[current]
+        }
+
+        return result.asReversed()
+    }
+}
+
+fun FirClass<*>.collectLocalClassesNavigationInfo(): LocalClassesNavigationInfo =
+    NavigationInfoVisitor().run {
+        this@collectLocalClassesNavigationInfo.accept(this@run)
+
+        LocalClassesNavigationInfo(parentForClass, resultingMap)
+    }
+
+private class NavigationInfoVisitor : FirDefaultVisitorVoid() {
+    val resultingMap = mutableMapOf<FirCallableMemberDeclaration<*>, FirClass<*>>()
+    val parentForClass = mutableMapOf<FirClass<*>, FirClass<*>?>()
+    private var currentPath = persistentListOf<FirClass<*>>()
+
+    override fun visitElement(element: FirElement) {}
+
+    override fun visitRegularClass(regularClass: FirRegularClass) {
+        visitClass(regularClass)
+    }
+
+    override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) {
+        visitClass(anonymousObject)
+    }
+
+    override fun <F : FirClass<F>> visitClass(klass: FirClass<F>) {
+        parentForClass[klass] = currentPath.lastOrNull()
+        val prev = currentPath
+        currentPath = currentPath.add(klass)
+
+        klass.acceptChildren(this)
+
+        currentPath = prev
+    }
+
+    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) {
+        visitCallableMemberDeclaration(simpleFunction)
+    }
+
+    override fun visitProperty(property: FirProperty) {
+        visitCallableMemberDeclaration(property)
+    }
+
+    override fun <F : FirCallableMemberDeclaration<F>> visitCallableMemberDeclaration(
+        callableMemberDeclaration: FirCallableMemberDeclaration<F>
+    ) {
+        if (callableMemberDeclaration.returnTypeRef !is FirImplicitTypeRef) return
+        resultingMap[callableMemberDeclaration] = currentPath.last()
+    }
+}