[FIR IDE] Fix IDE return type calculators

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildUnitExpression
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
@@ -881,5 +882,5 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         get() = initializer?.typeRef is FirResolvedTypeRef
 
     private val FirFunction.bodyResolved: Boolean
-        get() = body?.typeRef is FirResolvedTypeRef
+        get() = body !is FirLazyBlock && body?.typeRef is FirResolvedTypeRef
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/fir/PersistentCheckerContextFactory.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
 
 internal object PersistentCheckerContextFactory {
     fun createEmptyPersistenceCheckerContext(sessionHolder: SessionHolder): PersistentCheckerContext {
@@ -18,7 +18,7 @@ internal object PersistentCheckerContextFactory {
             sessionHolder.session,
             ScopeSession(),
             ImplicitBodyResolveComputationSession(),
-            ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+            ::FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
         )
         return PersistentCheckerContext(sessionHolder, returnTypeCalculator)
     }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator.kt

@@ -7,32 +7,32 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirConstructor
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDesignatedBodyResolveTransformerForReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
+import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator
 
-fun FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator(
+fun FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator(
     designation: Iterator<FirElement>,
     session: FirSession,
     scopeSession: ScopeSession,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
     returnTypeCalculator: ReturnTypeCalculator,
     outerBodyResolveContext: BodyResolveContext?,
-): FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl {
+): FirDesignatedBodyResolveTransformerForReturnTypeCalculator {
 
     val designationList = mutableListOf<FirElement>()
     for (element in designation) {
         designationList.add(element)
     }
+    require(designationList.isNotEmpty()) { "Designation should not be empty" }
 
     return FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
         designationList,
@@ -44,8 +44,8 @@ fun FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator(
     )
 }
 
-class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
-    designation: List<FirElement>,
+private class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
+    private val designation: List<FirElement>,
     session: FirSession,
     scopeSession: ScopeSession,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
@@ -59,26 +59,34 @@ class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
     returnTypeCalculator,
     outerBodyResolveContext
 ) {
-    private val declarationDesignation = designation.filterIsInstance<FirDeclaration>()
+    private val targetDeclaration = designation.last()
+
+    private inline fun <D : FirCallableDeclaration> D.processCallable(body: (FirDeclarationDesignation) -> Unit) {
+        if (this !== targetDeclaration) return
+        ensureResolved(FirResolvePhase.TYPES)
+        if (returnTypeRef !is FirImplicitTypeRef) return
+        val declarationList = designation.filterIsInstance<FirDeclaration>()
+        check(declarationList.isNotEmpty()) { "Invalid empty declaration designation" }
+        body(FirDeclarationDesignation(declarationList.dropLast(1), this))
+    }
 
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,
         data: ResolutionMode
     ): FirSimpleFunction {
-        val firDesignation = FirDeclarationDesignation(declarationDesignation, simpleFunction)
-        FirLazyBodiesCalculator.calculateLazyBodiesForFunction(firDesignation)
+        simpleFunction.processCallable {
+            FirLazyBodiesCalculator.calculateLazyBodiesForFunction(it)
+        }
         return super.transformSimpleFunction(simpleFunction, data)
     }
 
-    override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): FirConstructor {
-        val firDesignation = FirDeclarationDesignation(declarationDesignation, constructor)
-        FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(firDesignation)
-        return super.transformConstructor(constructor, data)
-    }
-
-    override fun transformProperty(property: FirProperty, data: ResolutionMode): FirProperty {
-        val firDesignation = FirDeclarationDesignation(declarationDesignation, property)
-        FirLazyBodiesCalculator.calculateLazyBodyForProperty(firDesignation)
+    override fun transformProperty(
+        property: FirProperty,
+        data: ResolutionMode
+    ): FirProperty {
+        property.processCallable {
+            FirLazyBodiesCalculator.calculateLazyBodyForProperty(it)
+        }
         return super.transformProperty(property, data)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirIdeEnsureBasedTransformerForReturnTypeCalculator.kt

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2010-2021 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.idea.fir.low.level.api.element.builder
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
+import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
+
+fun FirIdeEnsureBasedTransformerForReturnTypeCalculator(
+    designation: Iterator<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+    returnTypeCalculator: ReturnTypeCalculator,
+    outerBodyResolveContext: BodyResolveContext?,
+): FirDesignatedBodyResolveTransformerForReturnTypeCalculator {
+
+    val designationList = mutableListOf<FirElement>()
+    for (element in designation) {
+        designationList.add(element)
+    }
+    require(designationList.isNotEmpty()) { "Designation should not be empty" }
+
+    return FirIdeEnsureBasedTransformerForReturnTypeCalculatorImpl(
+        designationList,
+        session,
+        scopeSession,
+        implicitBodyResolveComputationSession,
+        returnTypeCalculator,
+        outerBodyResolveContext
+    )
+}
+
+
+private class FirIdeEnsureBasedTransformerForReturnTypeCalculatorImpl(
+    designation: List<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+    returnTypeCalculator: ReturnTypeCalculator,
+    outerBodyResolveContext: BodyResolveContext?,
+) : FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
+    designation.iterator(),
+    session,
+    scopeSession,
+    implicitBodyResolveComputationSession,
+    returnTypeCalculator,
+    outerBodyResolveContext
+) {
+    private val targetDeclaration = designation.last()
+
+    private fun <T : FirCallableDeclaration> T.ensuredReturnType() {
+        if (this !== targetDeclaration) return
+        ensureResolved(FirResolvePhase.TYPES)
+        if (returnTypeRef !is FirImplicitTypeRef) return
+        ensureResolved(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
+    }
+
+    override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: ResolutionMode): FirSimpleFunction {
+        simpleFunction.ensuredReturnType()
+        return super.transformSimpleFunction(simpleFunction, data)
+    }
+
+
+    override fun transformProperty(property: FirProperty, data: ResolutionMode): FirProperty {
+        property.ensuredReturnType()
+        return super.transformProperty(property, data)
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/LockProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 
 import com.google.common.collect.MapMaker
+import com.google.common.util.concurrent.CycleDetectingLockFactory
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.PrivateForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.lockWithPCECheck
@@ -16,8 +17,16 @@ import kotlin.concurrent.withLock
 internal class LockProvider<KEY> {
     private val locks: ConcurrentMap<KEY, ReentrantLock> = MapMaker().weakKeys().makeMap()
 
+    @Suppress("UnstableApiUsage")
+    private val lockFactory = CycleDetectingLockFactory.newInstance(CycleDetectingLockFactory.Policies.THROW)
+
     @Suppress("NOTHING_TO_INLINE")
-    private inline fun getLockFor(key: KEY) = locks.getOrPut(key) { ReentrantLock() }
+    private inline fun getLockFor(key: KEY) = locks.getOrPut(key) {
+        val file = key as FirFile
+        val name = "${file.packageDirective.packageFqName.asString()}.${file.name}"
+        @Suppress("UnstableApiUsage")
+        lockFactory.newReentrantLock(name)
+    }
 
     @OptIn(PrivateForInline::class)
     inline fun <R> withWriteLock(key: KEY, action: () -> R): R =

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.fir.expressions.FirAnnotationResolveStatus
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
-import org.jetbrains.kotlin.fir.symbols.impl.FirAccessorSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.tryCollectDesignationWithFile

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedBodyResolveTransformerForIDE.kt

@@ -16,7 +16,8 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyRe
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeEnsureBasedTransformerForReturnTypeCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.updatePhaseDeep
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhase
 
@@ -38,7 +39,7 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         session,
         scopeSession,
         ImplicitBodyResolveComputationSession(),
-        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+        ::FirIdeEnsureBasedTransformerForReturnTypeCalculator
     ),
     firTowerDataContextCollector = towerDataContextCollector,
     firProviderInterceptor = firProviderInterceptor,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedImplicitTypesTransformerForIDE.kt

@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTy
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.updatePhaseDeep
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhase
 
@@ -40,7 +40,7 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
         session,
         scopeSession,
         implicitBodyResolveComputationSession,
-        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+        ::FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirFileAnnotationsResolveTransformer.kt

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataCo
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
 
 internal class FirFileAnnotationsResolveTransformer(
     private val firFile: FirFile,
@@ -31,7 +31,7 @@ internal class FirFileAnnotationsResolveTransformer(
         session,
         scopeSession,
         implicitBodyResolveComputationSession,
-        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+        ::FirIdeDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator
     ),
     firTowerDataContextCollector = firTowerDataContextCollector
 ), FirLazyTransformerForIDE {