FIR IDE: refactor designated lazy transformers

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedBodyResolveTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedContractsResolveTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedImplicitTypesTransformerForIDE
+import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.*
 import org.jetbrains.kotlin.psi.*
 
@@ -158,7 +159,6 @@ internal class FirLazyDeclarationResolver(
 
         val transformer = phase.createLazyTransformer(
             designation,
-            firDeclarationToResolve,
             containerFirFile,
             scopeSession,
             towerDataContextCollector
@@ -171,27 +171,23 @@ internal class FirLazyDeclarationResolver(
 
     private fun FirResolvePhase.createLazyTransformer(
         designation: List<FirDeclaration>,
-        targetDeclaration: FirDeclaration,
         containerFirFile: FirFile,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?
     ) = when (this) {
         FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
-            designation.iterator(),
+            FirDesignation(designation),
-            targetDeclaration,
             containerFirFile.session,
             scopeSession,
         )
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
-            designation.iterator(),
+            FirDesignation(designation),
-            targetDeclaration,
             containerFirFile.session,
             scopeSession,
             towerDataContextCollector,
         )
         FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
-            designation.iterator(),
+            FirDesignation(designation),
-            targetDeclaration,
             containerFirFile.session,
             scopeSession,
             towerDataContextCollector

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
 
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -21,8 +20,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesigna
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 
 internal class FirDesignatedBodyResolveTransformerForIDE(
-    private val designation: Iterator<FirDeclaration>,
+    designation: FirDesignation,
-    targetDeclaration: FirDeclaration,
     session: FirSession,
     scopeSession: ScopeSession,
     private val towerDataContextCollector: FirTowerDataContextCollector? = null
@@ -38,17 +36,14 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
-    private val phaseReplaceOracle = PhaseReplaceOracle(targetDeclaration)
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
-    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+    @Suppress("NAME_SHADOWING")
-        if (designation.hasNext()) phaseReplaceOracle.transformDeclarationInside(designation.next()) {
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> =
-            it.visitNoTransform(this, data)
+        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) { declaration, data ->
-            return declaration.compose()
+            super.transformDeclarationContent(declaration, data)
         }
 
-        return super.transformDeclarationContent(declaration, data)
-    }
-
     override fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {
         towerDataContextCollector?.addDeclarationContext(declaration, context.towerDataContext)
     }
@@ -57,5 +52,7 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         towerDataContextCollector?.addStatementContext(statement, context.towerDataContext)
     }
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = phaseReplaceOracle.needReplacePhase(firDeclaration)
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+        ideDeclarationTransformer.needReplacePhase(firDeclaration)
 }
+

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trasformers/FirDesignatedContractsResolveTransformerForIDE.kt

@@ -5,30 +5,26 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
 
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformer
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
-import org.jetbrains.kotlin.fir.visitors.compose
 
 internal class FirDesignatedContractsResolveTransformerForIDE(
-    private val designation: Iterator<FirDeclaration>,
+    designation: FirDesignation,
-    targetDeclaration: FirDeclaration,
     session: FirSession,
     scopeSession: ScopeSession,
 ) : FirContractResolveTransformer(session, scopeSession) {
-    private val phaseReplaceOracle = PhaseReplaceOracle(targetDeclaration)
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
-    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+    @Suppress("NAME_SHADOWING")
-        if (designation.hasNext()) phaseReplaceOracle.transformDeclarationInside(designation.next()) {
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> =
-            it.visitNoTransform(this, data)
+        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) { declaration, data ->
-            return declaration.compose()
+            super.transformDeclarationContent(declaration, data)
         }
-        return super.transformDeclarationContent(declaration, data)
-    }
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = phaseReplaceOracle.needReplacePhase(firDeclaration)
+
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase(firDeclaration)
 }

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

@@ -20,8 +20,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesigna
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 
 internal class FirDesignatedImplicitTypesTransformerForIDE(
-    private val designation: Iterator<FirDeclaration>,
+    designation: FirDesignation,
-    targetDeclaration: FirDeclaration,
     session: FirSession,
     scopeSession: ScopeSession,
     private val towerDataContextCollector: FirTowerDataContextCollector?,
@@ -39,17 +38,16 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
-    private val phaseReplaceOracle = PhaseReplaceOracle(targetDeclaration)
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
-    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+    @Suppress("NAME_SHADOWING")
-        if (designation.hasNext()) phaseReplaceOracle.transformDeclarationInside(designation.next()) {
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> =
-            it.visitNoTransform(this, data)
+        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) { declaration, data ->
-            return declaration.compose()
+            super.transformDeclarationContent(declaration, data)
         }
-        return super.transformDeclarationContent(declaration, data)
-    }
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = phaseReplaceOracle.needReplacePhase(firDeclaration)
+
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase(firDeclaration)
 
     override fun onBeforeStatementResolution(statement: FirStatement) {
         towerDataContextCollector?.addStatementContext(statement, context.towerDataContext)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trasformers/FirDesignation.kt

@@ -0,0 +1,26 @@
+/*
+ * 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.trasformers
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+
+class FirDesignation(val designation: List<FirDeclaration>) {
+    private var currentIndex = 0
+
+    fun canGoNext(): Boolean = currentIndex <= designation.lastIndex
+
+    fun isTargetDeclaration(): Boolean = currentIndex == designation.lastIndex
+
+    val currentDeclaration get() = designation[currentIndex]
+
+    fun goNext() {
+        if (canGoNext()) {
+            currentIndex++
+        } else {
+            throw IndexOutOfBoundsException()
+        }
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trasformers/IDEDeclarationTransformer.kt

@@ -0,0 +1,61 @@
+/*
+ * 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.trasformers
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.PrivateForInline
+
+internal class IDEDeclarationTransformer(private val designation: FirDesignation) {
+    private var isInsideTargetDeclaration: Boolean = false
+
+    private inline fun <R> insideTargetDeclaration(insideCurrent: Boolean, action: () -> R): R {
+        val oldValue = isInsideTargetDeclaration
+        isInsideTargetDeclaration = insideCurrent
+        try {
+            return action()
+        } finally {
+            isInsideTargetDeclaration = oldValue
+        }
+    }
+
+    inline fun transformDeclarationContent(
+        transformer: FirAbstractBodyResolveTransformer,
+        declaration: FirDeclaration,
+        data: ResolutionMode,
+        transformDeclaration: (FirDeclaration, ResolutionMode) -> CompositeTransformResult<FirDeclaration>
+    ): CompositeTransformResult<FirDeclaration> {
+        return if (designation.canGoNext()) {
+            val declarationToTransform = designation.currentDeclaration
+            val isTargetDeclaration = designation.isTargetDeclaration()
+            designation.goNext()
+            insideTargetDeclaration(isTargetDeclaration) {
+                declarationToTransform.visitNoTransform(transformer, data)
+            }
+            declaration.compose()
+        } else {
+            if (isInsideTargetDeclaration) {
+                transformDeclaration(declaration, data)
+            } else {
+                declaration.compose()
+            }
+        }
+    }
+
+    fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+        isInsideTargetDeclaration
+}
+
+private fun <D> FirElement.visitNoTransform(transformer: FirTransformer<D>, data: D) {
+    val result = this.transform<FirElement, D>(transformer, data)
+    require(result.single === this) { "become ${result.single}: `${result.single.render()}`, was ${this}: `${this.render()}`" }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/trasformers/PhaseReplaceOracle.kt

@@ -1,24 +0,0 @@
-/*
- * 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.idea.fir.low.level.api.trasformers
-
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-
-internal class PhaseReplaceOracle(private val targetDeclaration: FirDeclaration) {
-    private var isInsideCurrentDeclaration = false
-
-    fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
-        isInsideCurrentDeclaration
-
-    inline fun <T> transformDeclarationInside(declaration: FirDeclaration, transform: (FirDeclaration) -> T): T {
-        if (declaration == targetDeclaration) isInsideCurrentDeclaration = true
-        return try {
-            transform(declaration)
-        } finally {
-            if (declaration == targetDeclaration) isInsideCurrentDeclaration = false
-        }
-    }
-}