[LL FIR] rework LLFirResolveTarget

* Move visiting logic to LLFirTargetResolver to
avoid code duplication
* Introduce more common LLFirWholeElementResolveTarget

^KT-61296

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt

@@ -7,8 +7,8 @@ package org.jetbrains.kotlin.analysis.api.fir.components
 
 import com.intellij.openapi.progress.ProgressManager
 import org.jetbrains.kotlin.analysis.api.compile.CodeFragmentCapturedValue
-import org.jetbrains.kotlin.analysis.api.components.KtCompilerFacility
 import org.jetbrains.kotlin.analysis.api.components.KtCompilationResult
+import org.jetbrains.kotlin.analysis.api.components.KtCompilerFacility
 import org.jetbrains.kotlin.analysis.api.components.KtCompilerTarget
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnostic
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
@@ -18,7 +18,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilt
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDiagnosticsForFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeElementResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.resolve
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CodeFragmentCapturedId
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CodeFragmentCapturedValueAnalyzer
@@ -42,8 +42,8 @@ import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.backend.jvm.*
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.fir.pipeline.signatureComposerForJvmFir2Ir
 import org.jetbrains.kotlin.fir.psi
@@ -66,8 +66,8 @@ import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbolInternals
 import org.jetbrains.kotlin.ir.types.IrSimpleType
-import org.jetbrains.kotlin.ir.util.classId
 import org.jetbrains.kotlin.ir.util.StubGeneratorExtensions
+import org.jetbrains.kotlin.ir.util.classId
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
@@ -291,7 +291,7 @@ internal class KtFirCompilerFacility(
 
     private fun getFullyResolvedFirFile(file: KtFile): FirFile {
         val firFile = file.getOrBuildFirFile(firResolveSession)
-        LLFirWholeFileResolveTarget(firFile).resolve(FirResolvePhase.BODY_RESOLVE)
+        LLFirWholeElementResolveTarget(firFile).resolve(FirResolvePhase.BODY_RESOLVE)
         return firFile
     }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirClassWithAllCallablesResolveTarget.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDeclaration
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
@@ -15,18 +16,19 @@ import org.jetbrains.kotlin.fir.declarations.FirRegularClass
  */
 class LLFirClassWithAllCallablesResolveTarget(
     firFile: FirFile,
-    classPath: List<FirRegularClass>,
+    containerClasses: List<FirRegularClass>,
     target: FirRegularClass,
-) : LLFirResolveTarget(firFile, classPath, target) {
-    override fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
-        action(target)
-        forEachCallable(action)
-    }
-
-    inline fun forEachCallable(action: (FirCallableDeclaration) -> Unit) {
-        for (member in (target as FirRegularClass).declarations) {
-            if (member is FirCallableDeclaration) {
-                action(member)
+) : LLFirResolveTarget(firFile, containerClasses, target) {
+    override fun visitTargetElement(
+        element: FirElementWithResolveState,
+        visitor: LLFirResolveTargetVisitor,
+    ) {
+        visitor.performAction(element)
+        visitor.withRegularClass(element as FirRegularClass) {
+            element.forEachDeclaration {
+                if (it is FirCallableDeclaration) {
+                    visitor.performAction(it)
+                }
             }
         }
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirClassWithAllMembersResolveTarget.kt

@@ -5,8 +5,8 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDeclaration
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 
@@ -15,15 +15,16 @@ import org.jetbrains.kotlin.fir.declarations.FirRegularClass
  */
 class LLFirClassWithAllMembersResolveTarget(
     firFile: FirFile,
-    classPath: List<FirRegularClass>,
+    containerClasses: List<FirRegularClass>,
     target: FirRegularClass,
-) : LLFirResolveTarget(firFile, classPath, target) {
-    override fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
-        action(target)
-        forEachMember(action)
-    }
-
-    inline fun forEachMember(action: (FirDeclaration) -> Unit) {
-        (target as FirRegularClass).declarations.forEach(action)
+) : LLFirResolveTarget(firFile, containerClasses, target) {
+    override fun visitTargetElement(
+        element: FirElementWithResolveState,
+        visitor: LLFirResolveTargetVisitor,
+    ) {
+        visitor.performAction(element)
+        visitor.withRegularClass(element as FirRegularClass) {
+            element.forEachDeclaration(visitor::performAction)
+        }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirClassWithSpecificMembersResolveTarget.kt

@@ -12,17 +12,18 @@ import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 
 class LLFirClassWithSpecificMembersResolveTarget(
     firFile: FirFile,
-    classPath: List<FirRegularClass>,
+    containerClasses: List<FirRegularClass>,
     target: FirRegularClass,
     val members: List<FirDeclaration>,
-) : LLFirResolveTarget(firFile, classPath, target) {
-    override fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
-        action(target)
-        forEachMember(action)
-    }
-
-    fun forEachMember(action: (FirDeclaration) -> Unit) {
-        members.forEach(action)
+) : LLFirResolveTarget(firFile, containerClasses, target) {
+    override fun visitTargetElement(
+        element: FirElementWithResolveState,
+        visitor: LLFirResolveTargetVisitor,
+    ) {
+        visitor.performAction(element)
+        visitor.withRegularClass(element as FirRegularClass) {
+            members.forEach(visitor::performAction)
+        }
     }
 
     override fun toStringAdditionalSuffix(): String = members.joinToString(prefix = "[", postfix = "]") { it.symbol.toString() }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirResolveTarget.kt

@@ -32,7 +32,7 @@ sealed class LLFirResolveTarget(
     /**
      * The list of [FirRegularClass] which are the required to go from file to target declarations in the top-down order.
      */
-    classPath: List<FirRegularClass>,
+    containerClasses: List<FirRegularClass>,
 
     /**
      * A dedicated main element.
@@ -44,12 +44,59 @@ sealed class LLFirResolveTarget(
      *
      * If resolve target is [FirRegularClass] or [FirScript] itself, it's not included into the [path].
      */
-    val path: List<FirDeclaration> = pathWithScript(firFile, classPath, target)
+    val path: List<FirDeclaration> = pathWithScript(firFile, containerClasses, target)
+
+    /**
+     * Visit [path], [target] and optionally its subgraph.
+     * Each nested declaration will be wrapped with corresponding [withFile], [withClass] and [withScript] recursively.
+     */
+    fun visit(visitor: LLFirResolveTargetVisitor) {
+        if (target is FirFile) {
+            visitor.performAction(target)
+        }
+
+        visitor.withFile(firFile) {
+            val pathIterator = path.iterator()
+            goToTarget(pathIterator, visitor)
+        }
+    }
+
+    private fun goToTarget(
+        pathIterator: Iterator<FirDeclaration>,
+        visitor: LLFirResolveTargetVisitor,
+    ) {
+        if (pathIterator.hasNext()) {
+            when (val declaration = pathIterator.next()) {
+                is FirRegularClass -> visitor.withRegularClass(declaration) { goToTarget(pathIterator, visitor) }
+                is FirScript -> visitor.withScript(declaration) { goToTarget(pathIterator, visitor) }
+                else -> errorWithFirSpecificEntries(
+                    "Unexpected declaration in path: ${declaration::class.simpleName}",
+                    fir = declaration,
+                )
+            }
+        } else {
+            visitTargetElement(target, visitor)
+        }
+    }
+
+    /**
+     * [element] with [FirFile] will be processed before.
+     */
+    protected abstract fun visitTargetElement(
+        element: FirElementWithResolveState,
+        visitor: LLFirResolveTargetVisitor,
+    )
 
     /**
      * Executions the [action] for each target that this [LLFirResolveTarget] represents.
      */
-    abstract fun forEachTarget(action: (FirElementWithResolveState) -> Unit)
+    fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
+        visit(object : LLFirResolveTargetVisitor {
+            override fun performAction(element: FirElementWithResolveState) {
+                action(element)
+            }
+        })
+    }
 
     override fun toString(): String = buildString {
         append(this@LLFirResolveTarget::class.simpleName)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirResolveTargetVisitor.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2023 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.analysis.low.level.api.fir.api.targets
+
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirScript
+
+/**
+ * This interface describes how to process nested declarations.
+ *
+ * @see LLFirResolveTarget
+ */
+interface LLFirResolveTargetVisitor {
+    /**
+     * Access to [FirFile] declaration will be performed inside [action].
+     */
+    fun withFile(firFile: FirFile, action: () -> Unit): Unit = action()
+
+    /**
+     * Access to elements inside [FirRegularClass] will be performed inside [action].
+     * Will be called for each nested [FirRegularClass] on the path.
+     */
+    fun withRegularClass(firClass: FirRegularClass, action: () -> Unit): Unit = action()
+
+    /**
+     * Access to elements inside [FirScript] will be performed inside [action].
+     */
+    fun withScript(firScript: FirScript, action: () -> Unit): Unit = action()
+
+    /**
+     * This method will be performed on some target element depends on [LLFirResolveTarget] implementation.
+     */
+    fun performAction(element: FirElementWithResolveState)
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirSingleResolveTarget.kt

@@ -13,10 +13,15 @@ import org.jetbrains.kotlin.fir.declarations.*
  */
 class LLFirSingleResolveTarget(
     firFile: FirFile,
-    classPath: List<FirRegularClass>,
+    containerClasses: List<FirRegularClass>,
     target: FirElementWithResolveState,
-) : LLFirResolveTarget(firFile, classPath, target) {
-    override fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
-        action(target)
+) : LLFirResolveTarget(firFile, containerClasses, target) {
+    override fun visitTargetElement(
+        element: FirElementWithResolveState,
+        visitor: LLFirResolveTargetVisitor,
+    ) {
+        if (element !is FirFile) {
+            visitor.performAction(element)
+        }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirWholeClassResolveTarget.kt

@@ -1,41 +0,0 @@
-/*
- * Copyright 2010-2023 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.analysis.low.level.api.fir.api.targets
-
-import org.jetbrains.kotlin.fir.FirElementWithResolveState
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-
-/**
- * [LLFirResolveTarget] representing all declarations in a class (including nested classes).
- * All of them are going to be resolved.
- */
-class LLFirWholeClassResolveTarget(
-    firFile: FirFile,
-    classPath: List<FirRegularClass>,
-    target: FirRegularClass,
-) : LLFirResolveTarget(firFile, classPath, target) {
-    override fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
-        action(target)
-
-        fun goInside(target: FirElementWithResolveState) {
-            action(target)
-            if (target is FirRegularClass) {
-                target.declarations.forEach(::goInside)
-            }
-        }
-
-        forEachDeclaration(::goInside)
-    }
-
-    inline fun forEachDeclaration(action: (FirElementWithResolveState) -> Unit) {
-        for (member in (target as FirRegularClass).declarations) {
-            action(member)
-        }
-    }
-
-    override fun toStringAdditionalSuffix(): String = "*"
-}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirWholeElementResolveTarget.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2010-2023 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.analysis.low.level.api.fir.api.targets
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isDeclarationContainer
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirScript
+
+/**
+ * [LLFirResolveTarget] representing all declarations in [target] recursively.
+ * All of them are going to be resolved.
+ */
+class LLFirWholeElementResolveTarget(
+    firFile: FirFile,
+    containerClasses: List<FirRegularClass>,
+    target: FirElementWithResolveState,
+) : LLFirResolveTarget(firFile, containerClasses, target) {
+    constructor(firFile: FirFile) : this(firFile, emptyList(), firFile)
+
+    override fun visitTargetElement(
+        element: FirElementWithResolveState,
+        visitor: LLFirResolveTargetVisitor,
+    ) {
+        if (element !is FirFile) {
+            visitor.performAction(element)
+        }
+
+        when {
+            element !is FirDeclaration || !element.isDeclarationContainer -> {}
+
+            element is FirRegularClass -> visitor.withRegularClass(element) {
+                element.forEachDeclaration {
+                    visitTargetElement(it, visitor)
+                }
+            }
+
+            element is FirFile -> {
+                element.forEachDeclaration {
+                    visitTargetElement(it, visitor)
+                }
+            }
+
+            element is FirScript -> visitor.withScript(element) {
+                element.forEachDeclaration {
+                    visitTargetElement(it, visitor)
+                }
+            }
+
+            else -> errorWithFirSpecificEntries("Unexpected declaration: ${element::class.simpleName}", fir = element)
+        }
+    }
+
+    override fun toStringAdditionalSuffix(): String = "*"
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirWholeFileResolveTarget.kt

@@ -1,39 +0,0 @@
-/*
- * Copyright 2010-2023 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.analysis.low.level.api.fir.api.targets
-
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDeclaration
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isDeclarationContainer
-import org.jetbrains.kotlin.fir.FirElementWithResolveState
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-
-/**
- * [LLFirResolveTarget] representing all declarations in file. All of them are going to be resolved.
- */
-class LLFirWholeFileResolveTarget(firFile: FirFile) : LLFirResolveTarget(firFile, emptyList(), firFile) {
-    override fun forEachTarget(action: (FirElementWithResolveState) -> Unit) {
-        fun goInside(target: FirElementWithResolveState) {
-            action(target)
-            if (target is FirDeclaration && target.isDeclarationContainer) {
-                target.forEachDeclaration(::goInside)
-            }
-        }
-
-        action(firFile)
-        forEachTopLevelDeclaration(::goInside)
-    }
-
-    inline fun forEachTopLevelDeclaration(action: (FirElementWithResolveState) -> Unit) {
-        firFile.annotationsContainer?.let { action(it) }
-
-        for (member in firFile.declarations) {
-            action(member)
-        }
-    }
-
-    override fun toStringAdditionalSuffix(): String = "*"
-}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt

@@ -9,22 +9,21 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirSingleResolveTarget
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeClassResolveTarget
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeElementResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirLazyResolverRunner
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.withOnAirDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
-import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
-import org.jetbrains.kotlin.utils.exceptions.rethrowExceptionWithDetails
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
+import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.utils.exceptions.rethrowExceptionWithDetails
 
 internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirModuleResolveComponents) {
     /**
@@ -44,7 +43,7 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
         val fromPhase = target.resolvePhase
 
         /**
-         * Currently [lazyResolve] on file means [LLFirWholeFileResolveTarget], but also [FirFile] itself
+         * Currently [lazyResolve] on file means [LLFirWholeElementResolveTarget], but also [FirFile] itself
          * has [resolvePhase] which does not match with the entire file resolution state.
          * This additional [FirFile] condition can be dropped after KT-61296
          */
@@ -138,9 +137,9 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
         resolveFileToImportsWithLock(designation.firFile)
 
         val target = when {
-            designation.target is FirFile -> LLFirWholeFileResolveTarget(designation.firFile)
+            designation.target is FirFile -> LLFirWholeElementResolveTarget(designation.firFile)
             resolvePhase == FirResolvePhase.BODY_RESOLVE && designation.target is FirRegularClass -> {
-                LLFirWholeClassResolveTarget(designation.firFile, designation.path, designation.target)
+                LLFirWholeElementResolveTarget(designation.firFile, designation.path, designation.target)
             }
 
             else -> LLFirSingleResolveTarget(designation.firFile, designation.path, designation.target)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirResolveMultiDesignationCollector.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDesignationWithFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirClassWithAllCallablesResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeElementResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.asResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.tryCollectDesignationWithFile
@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticPropertyAcces
 internal object LLFirResolveMultiDesignationCollector {
     fun getDesignationsToResolve(target: FirElementWithResolveState): List<LLFirResolveTarget> {
         return when (target) {
-            is FirFile -> listOf(LLFirWholeFileResolveTarget(target))
+            is FirFile -> listOf(LLFirWholeElementResolveTarget(target))
             else -> getMainDesignationToResolve(target)?.withAnnotationContainer()
         } ?: emptyList()
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTargetResolver.kt

@@ -8,11 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDeclaration
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isDeclarationContainer
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -23,7 +19,7 @@ internal abstract class LLFirTargetResolver(
     protected val lockProvider: LLFirLockProvider,
     protected val resolverPhase: FirResolvePhase,
     private val isJumpingPhase: Boolean = false,
-) {
+) : LLFirResolveTargetVisitor {
     private val _nestedClassesStack = mutableListOf<FirRegularClass>()
 
     val nestedClassesStack: List<FirRegularClass> get() = _nestedClassesStack.toList()
@@ -39,11 +35,11 @@ internal abstract class LLFirTargetResolver(
         }
     }
 
-    protected open fun withFile(firFile: FirFile, action: () -> Unit) {
+    override fun withFile(firFile: FirFile, action: () -> Unit) {
         action()
     }
 
-    protected open fun withScript(firScript: FirScript, action: () -> Unit) {
+    override fun withScript(firScript: FirScript, action: () -> Unit) {
         action()
     }
 
@@ -53,7 +49,7 @@ internal abstract class LLFirTargetResolver(
     }
 
     @Suppress("DEPRECATION_ERROR")
-    protected fun withRegularClass(firClass: FirRegularClass, action: () -> Unit) {
+    final override fun withRegularClass(firClass: FirRegularClass, action: () -> Unit) {
         _nestedClassesStack += firClass
         withRegularClassImpl(firClass, action)
         check(_nestedClassesStack.removeLast() === firClass)
@@ -70,79 +66,11 @@ internal abstract class LLFirTargetResolver(
 
     fun resolveDesignation() {
         checkResolveConsistency()
-        if (resolveTarget is LLFirSingleResolveTarget && resolveTarget.target is FirFile) {
-            performResolve(resolveTarget.firFile)
-            return
-        }
-        if (resolveTarget is LLFirWholeFileResolveTarget) {
-            performResolve(resolveTarget.firFile)
-        }
-        withFile(resolveTarget.firFile) {
-            goToTargets(resolveTarget.path.iterator())
-        }
+        resolveTarget.visit(this)
     }
 
-    private fun goToTargets(path: Iterator<FirDeclaration>) {
-        if (path.hasNext()) {
-            when (val firDeclaration = path.next()) {
-                is FirRegularClass -> withRegularClass(firDeclaration) { goToTargets(path) }
-                is FirScript -> withScript(firDeclaration) { goToTargets(path) }
-                else -> errorWithFirSpecificEntries(
-                    "Unexpected declaration in path: ${firDeclaration::class.simpleName}",
-                    fir = firDeclaration,
-                )
-            }
-        } else {
-            when (resolveTarget) {
-                is LLFirClassWithSpecificMembersResolveTarget -> {
-                    performResolve(resolveTarget.target)
-                    withRegularClass(resolveTarget.target as FirRegularClass) {
-                        resolveTarget.forEachMember(::performResolve)
-                    }
-                }
-                is LLFirClassWithAllMembersResolveTarget -> {
-                    performResolve(resolveTarget.target)
-                    withRegularClass(resolveTarget.target as FirRegularClass) {
-                        resolveTarget.forEachMember(::performResolve)
-                    }
-                }
-                is LLFirClassWithAllCallablesResolveTarget -> {
-                    performResolve(resolveTarget.target)
-                    withRegularClass(resolveTarget.target as FirRegularClass) {
-                        resolveTarget.forEachCallable(::performResolve)
-                    }
-                }
-                is LLFirSingleResolveTarget -> {
-                    performResolve(resolveTarget.target)
-                }
-                is LLFirWholeFileResolveTarget -> {
-                    resolveTarget.forEachTopLevelDeclaration(::resolveTargetWithNestedDeclarations)
-                }
-                is LLFirWholeClassResolveTarget -> {
-                    performResolve(resolveTarget.target)
-                    withRegularClass(resolveTarget.target as FirRegularClass) {
-                        resolveTarget.forEachDeclaration(::resolveTargetWithNestedDeclarations)
-                    }
-                }
-            }
-        }
-    }
-
-    private fun resolveTargetWithNestedDeclarations(target: FirElementWithResolveState) {
-        performResolve(target)
-        when {
-            target !is FirDeclaration || !target.isDeclarationContainer -> {}
-
-            target is FirRegularClass -> withRegularClass(target) {
-                target.forEachDeclaration(::resolveTargetWithNestedDeclarations)
-            }
-
-            target is FirScript -> withScript(target) {
-                target.forEachDeclaration(::resolveTargetWithNestedDeclarations)
-            }
-
-            else -> errorWithFirSpecificEntries("Unexpected declaration: ${target::class.simpleName}", fir = target)
-        }
+    final override fun performAction(element: FirElementWithResolveState) {
+        performResolve(element)
     }
 
     protected fun performResolve(target: FirElementWithResolveState) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt

@@ -219,12 +219,17 @@ internal inline fun FirRegularClass.forEachDeclaration(action: (FirDeclaration)
     declarations.forEach(action)
 }
 
-internal val FirDeclaration.isDeclarationContainer: Boolean get() = this is FirRegularClass || this is FirScript
+internal inline fun FirFile.forEachDeclaration(action: (FirDeclaration) -> Unit) {
+    declarations.forEach(action)
+}
+
+internal val FirDeclaration.isDeclarationContainer: Boolean get() = this is FirRegularClass || this is FirScript || this is FirFile
 
 internal inline fun FirDeclaration.forEachDeclaration(action: (FirDeclaration) -> Unit) {
     when (this) {
         is FirRegularClass -> forEachDeclaration(action)
         is FirScript -> forEachDeclaration(action)
+        is FirFile -> forEachDeclaration(action)
         else -> errorWithFirSpecificEntries("Unsupported declarations container", fir = this)
     }
 }