FIR: Implement FirApplyInferredDeclarationTypesTransformer

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.FirSimpleDiagnostic
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirErrorTypeRefImpl
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
 
@@ -47,6 +49,41 @@ class FirImplicitTypeBodyResolveTransformerAdapter2 : FirTransformer<Nothing?>()
     }
 }
 
+// TODO: This class is unused because it's effectively unneeded while we have mutable trees
+// Once we decide to switch to persistent trees this, it should be applied to each file
+// If the decision is to stay with mutable trees, this class should be removed
+private class FirApplyInferredDeclarationTypesTransformer(
+    private val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession
+) : FirDefaultTransformer<Nothing?>() {
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+        return element.compose()
+    }
+
+    override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        file.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+
+        return (file.transformChildren(this, null) as FirFile).compose()
+    }
+
+    override fun transformRegularClass(regularClass: FirRegularClass, data: Nothing?): CompositeTransformResult<FirStatement> {
+        return (regularClass.transformChildren(this, null) as FirRegularClass).compose()
+    }
+
+    override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        if (simpleFunction.returnTypeRef !is FirImplicitTypeRef) return simpleFunction.compose()
+
+        val transformedDeclaration = implicitBodyResolveComputationSession.getComputedResult(simpleFunction.symbol).transformedDeclaration
+        return transformedDeclaration.compose()
+    }
+
+    override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        if (property.returnTypeRef !is FirImplicitTypeRef) return property.compose()
+
+        val transformedDeclaration = implicitBodyResolveComputationSession.getComputedResult(property.symbol).transformedDeclaration
+        return transformedDeclaration.compose()
+    }
+}
+
 private open class MyFirBodyResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
@@ -235,6 +272,15 @@ private class ImplicitBodyResolveComputationSession {
 
         implicitBodyResolveStatusMap[symbol] = ImplicitBodyResolveComputationStatus.Computed(returnTypeRef, transformedDeclaration)
     }
+
+    fun getComputedResult(symbol: FirCallableSymbol<*>): ImplicitBodyResolveComputationStatus.Computed {
+        val status = getStatus(symbol)
+        require(status is ImplicitBodyResolveComputationStatus.Computed) {
+            "Unexpected status $status for ${symbol.fir.render()}"
+        }
+
+        return status
+    }
 }
 
 private sealed class ImplicitBodyResolveComputationStatus {