[LL FIR] extract base logic of lazy resolution tests

^KT-63042

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyDeclarationResolveOverAllPhasesTest.kt

@@ -0,0 +1,81 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirResolveMultiDesignationCollector
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhaseRecursively
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.assertions
+
+/**
+ * This test iterates over all [FirResolvePhase] for the selected declaration and dump output after each phase
+ */
+abstract class AbstractFirLazyDeclarationResolveOverAllPhasesTest : AbstractFirLazyDeclarationResolveTestCase() {
+    abstract fun checkSession(firSession: LLFirResolveSession)
+
+    protected fun doLazyResolveTest(
+        ktFile: KtFile,
+        testServices: TestServices,
+        renderAllFiles: Boolean,
+        resolverProvider: (LLFirResolveSession) -> Pair<FirElementWithResolveState, ((FirResolvePhase) -> Unit)>,
+    ) {
+        val resultBuilder = StringBuilder()
+        val renderer = lazyResolveRenderer(resultBuilder)
+
+        resolveWithClearCaches(ktFile) { firResolveSession ->
+            checkSession(firResolveSession)
+
+            val (elementToResolve, resolver) = resolverProvider(firResolveSession)
+            val filesToRender = if (renderAllFiles) {
+                testServices.allKtFiles().map(firResolveSession::getOrBuildFirFile)
+            } else {
+                val firFile = firResolveSession.getOrBuildFirFile(ktFile)
+                val designations = LLFirResolveMultiDesignationCollector.getDesignationsToResolve(elementToResolve)
+                listOf(firFile).plus(designations.map { it.firFile }).distinct()
+            }
+
+            val shouldRenderDeclaration = elementToResolve !in filesToRender
+            for (currentPhase in FirResolvePhase.entries) {
+                if (currentPhase == FirResolvePhase.SEALED_CLASS_INHERITORS) continue
+                resolver(currentPhase)
+
+                if (resultBuilder.isNotEmpty()) {
+                    resultBuilder.appendLine()
+                }
+
+                resultBuilder.append("${currentPhase.name}:")
+                if (shouldRenderDeclaration) {
+                    resultBuilder.append("\nTARGET: ")
+                    renderer.renderElementAsString(elementToResolve)
+                }
+
+                for (file in filesToRender) {
+                    resultBuilder.appendLine()
+                    renderer.renderElementAsString(file)
+                }
+            }
+        }
+
+        resolveWithClearCaches(ktFile) { llSession ->
+            checkSession(llSession)
+            val firFile = llSession.getOrBuildFirFile(ktFile)
+            firFile.lazyResolveToPhaseRecursively(FirResolvePhase.BODY_RESOLVE)
+            if (resultBuilder.isNotEmpty()) {
+                resultBuilder.appendLine()
+            }
+
+            resultBuilder.append("FILE RAW TO BODY:\n")
+            renderer.renderElementAsString(firFile)
+        }
+
+        testServices.assertions.assertEqualsToTestDataFileSibling(resultBuilder.toString())
+    }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyDeclarationResolveTest.kt

@@ -6,70 +6,31 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirOutOfContentRootTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirScriptTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
-import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
-import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.directives.ConfigurationDirectives.WITH_STDLIB
 import org.jetbrains.kotlin.test.directives.JvmEnvironmentConfigurationDirectives.NO_RUNTIME
-import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 
-abstract class AbstractFirLazyDeclarationResolveTest : AbstractFirLazyDeclarationResolveTestCase() {
+abstract class AbstractFirLazyDeclarationResolveTest : AbstractFirLazyDeclarationResolveOverAllPhasesTest() {
     override fun checkSession(firSession: LLFirResolveSession) {
         require(firSession.isSourceSession)
     }
 
     override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
-        doLazyResolveTest(ktFile, testServices, moduleStructure, renderAllFiles = true) { firResolveSession ->
-            when {
-                Directives.RESOLVE_FILE_ANNOTATIONS in moduleStructure.allDirectives -> {
-                    val annotationContainer = firResolveSession.getOrBuildFirFile(ktFile).annotationsContainer!!
-                    val session = annotationContainer.moduleData.session as LLFirResolvableModuleSession
-                    annotationContainer to fun(phase: FirResolvePhase) {
-                        session.moduleComponents.firModuleLazyDeclarationResolver.lazyResolve(
-                            annotationContainer,
-                            session.getScopeSession(),
-                            phase,
-                        )
-                    }
-                }
-                Directives.RESOLVE_FILE in moduleStructure.allDirectives -> {
-                    val session = firResolveSession.useSiteFirSession as LLFirResolvableModuleSession
-                    val file = session.moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
-                    file to fun(phase: FirResolvePhase) {
-                        file.lazyResolveToPhase(phase)
-                    }
-                }
-                else -> {
-                    val ktDeclaration = if (Directives.RESOLVE_SCRIPT in moduleStructure.allDirectives) {
-                        ktFile.script!!
-                    } else {
-                        testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtDeclaration>(ktFile)
-                    }
-
-                    val declarationSymbol = ktDeclaration.resolveToFirSymbol(firResolveSession)
-                    val firDeclaration = chooseMemberDeclarationIfNeeded(declarationSymbol, moduleStructure, firResolveSession)
-                    firDeclaration.fir to fun(phase: FirResolvePhase) {
-                        firDeclaration.lazyResolveToPhase(phase)
-                    }
-                }
-            }
+        doLazyResolveTest(ktFile, testServices, renderAllFiles = true) { firResolveSession ->
+            findFirDeclarationToResolve(ktFile, moduleStructure, testServices, firResolveSession)
         }
     }
 
     override fun configureTest(builder: TestConfigurationBuilder) {
         super.configureTest(builder)
         with(builder) {
-            forTestsNotMatching("analysis/low-level-api-fir/testData/lazyResolve/noRuntime/*" ) {
+            forTestsNotMatching("analysis/low-level-api-fir/testData/lazyResolve/noRuntime/*") {
                 defaultDirectives {
                     +WITH_STDLIB
                 }
@@ -80,16 +41,8 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractFirLazyDeclaratio
                     +NO_RUNTIME
                 }
             }
-
-            useDirectives(Directives)
         }
     }
-
-    private object Directives : SimpleDirectivesContainer() {
-        val RESOLVE_FILE_ANNOTATIONS by directive("Resolve file annotations instead of declaration at caret")
-        val RESOLVE_SCRIPT by directive("Resolve script instead of declaration at caret")
-        val RESOLVE_FILE by directive("Resolve file instead of declaration at caret")
-    }
 }
 
 abstract class AbstractFirSourceLazyDeclarationResolveTest : AbstractFirLazyDeclarationResolveTest() {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyDeclarationResolveTestCase.kt

@@ -6,10 +6,11 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirResolveMultiDesignationCollector
-import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiLastModuleFirstFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder.findElementIn
-import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
+import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
@@ -28,7 +29,8 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
-import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhaseRecursively
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
@@ -36,78 +38,48 @@ import org.jetbrains.kotlin.test.directives.model.ValueDirective
 import org.jetbrains.kotlin.test.directives.model.singleOrZeroValue
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
-import org.jetbrains.kotlin.test.services.assertions
 
 /**
  * Test that we do not resolve declarations we do not need & do not build bodies for them
  */
-abstract class AbstractFirLazyDeclarationResolveTestCase : AbstractLowLevelApiSingleFileTest() {
-    abstract fun checkSession(firSession: LLFirResolveSession)
-
-    protected fun doLazyResolveTest(
+abstract class AbstractFirLazyDeclarationResolveTestCase : AbstractLowLevelApiLastModuleFirstFileTest() {
+    protected fun findFirDeclarationToResolve(
         ktFile: KtFile,
-        testServices: TestServices,
         moduleStructure: TestModuleStructure,
-        renderAllFiles: Boolean,
-        resolverProvider: (LLFirResolveSession) -> Pair<FirElementWithResolveState, ((FirResolvePhase) -> Unit)>,
-    ) {
-        val resultBuilder = StringBuilder()
-        val renderer = lazyResolveRenderer(resultBuilder)
-
-        resolveWithClearCaches(ktFile) { firResolveSession ->
-            checkSession(firResolveSession)
-
-            val (elementToResolve, resolver) = resolverProvider(firResolveSession)
-            val filesToRender = if (renderAllFiles) {
-                moduleStructure.modules
-                    .flatMap(testServices.ktModuleProvider::getModuleFiles)
-                    .mapNotNull { file -> (file as? KtFile)?.let { firResolveSession.getOrBuildFirFile(it) } }
+        testServices: TestServices,
+        firResolveSession: LLFirResolveSession,
+    ): Pair<FirElementWithResolveState, ((FirResolvePhase) -> Unit)> = when {
+        Directives.RESOLVE_FILE_ANNOTATIONS in moduleStructure.allDirectives -> {
+            val annotationContainer = firResolveSession.getOrBuildFirFile(ktFile).annotationsContainer!!
+            val session = annotationContainer.moduleData.session as LLFirResolvableModuleSession
+            annotationContainer to fun(phase: FirResolvePhase) {
+                session.moduleComponents.firModuleLazyDeclarationResolver.lazyResolve(
+                    annotationContainer,
+                    session.getScopeSession(),
+                    phase,
+                )
+            }
+        }
+        Directives.RESOLVE_FILE in moduleStructure.allDirectives -> {
+            val session = firResolveSession.useSiteFirSession as LLFirResolvableModuleSession
+            val file = session.moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
+            file to fun(phase: FirResolvePhase) {
+                file.lazyResolveToPhase(phase)
+            }
+        }
+        else -> {
+            val ktDeclaration = if (Directives.RESOLVE_SCRIPT in moduleStructure.allDirectives) {
+                ktFile.script!!
             } else {
-                val firFile = firResolveSession.getOrBuildFirFile(ktFile)
-                val designations = LLFirResolveMultiDesignationCollector.getDesignationsToResolve(elementToResolve)
-                listOf(firFile).plus(designations.map { it.firFile }).distinct()
+                testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtDeclaration>(ktFile)
             }
 
-            val shouldRenderDeclaration = elementToResolve !is FirFile && filesToRender.all { file ->
-                findElementIn<FirElementWithResolveState>(file) {
-                    it == elementToResolve
-                } == null
-            }
-
-            for (currentPhase in FirResolvePhase.entries) {
-                if (currentPhase == FirResolvePhase.SEALED_CLASS_INHERITORS) continue
-                resolver(currentPhase)
-
-                if (resultBuilder.isNotEmpty()) {
-                    resultBuilder.appendLine()
-                }
-
-                resultBuilder.append("${currentPhase.name}:")
-                if (shouldRenderDeclaration) {
-                    resultBuilder.append("\nTARGET: ")
-                    renderer.renderElementAsString(elementToResolve)
-                }
-
-                for (file in filesToRender) {
-                    resultBuilder.appendLine()
-                    renderer.renderElementAsString(file)
-                }
+            val declarationSymbol = ktDeclaration.resolveToFirSymbol(firResolveSession)
+            val firDeclaration = chooseMemberDeclarationIfNeeded(declarationSymbol, moduleStructure, firResolveSession)
+            firDeclaration.fir to fun(phase: FirResolvePhase) {
+                firDeclaration.lazyResolveToPhase(phase)
             }
         }
-
-        resolveWithClearCaches(ktFile) { llSession ->
-            checkSession(llSession)
-            val firFile = llSession.getOrBuildFirFile(ktFile)
-            firFile.lazyResolveToPhaseRecursively(FirResolvePhase.BODY_RESOLVE)
-            if (resultBuilder.isNotEmpty()) {
-                resultBuilder.appendLine()
-            }
-
-            resultBuilder.append("FILE RAW TO BODY:\n")
-            renderer.renderElementAsString(firFile)
-        }
-
-        testServices.assertions.assertEqualsToTestDataFileSibling(resultBuilder.toString())
     }
 
     protected fun chooseMemberDeclarationIfNeeded(
@@ -215,6 +187,10 @@ abstract class AbstractFirLazyDeclarationResolveTestCase : AbstractLowLevelApiSi
         }
 
         val IS_GETTER by valueDirective("Choose getter/setter in the case of property", parser = String::toBooleanStrict)
+
+        val RESOLVE_FILE_ANNOTATIONS by directive("Resolve file annotations instead of declaration at caret")
+        val RESOLVE_SCRIPT by directive("Resolve script instead of declaration at caret")
+        val RESOLVE_FILE by directive("Resolve file instead of declaration at caret")
     }
 }
 
@@ -225,3 +201,13 @@ internal fun lazyResolveRenderer(builder: StringBuilder): FirRenderer = FirRende
     errorExpressionRenderer = FirErrorExpressionExtendedRenderer(),
     fileAnnotationsContainerRenderer = FirFileAnnotationsContainerRenderer(),
 )
+
+internal operator fun List<FirFile>.contains(element: FirElementWithResolveState): Boolean = if (element is FirFile) {
+    element in this
+} else {
+    any { file ->
+        findElementIn<FirElementWithResolveState>(file) {
+            it == element
+        } != null
+    }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractStdLibSourcesLazyDeclarationResolveTest.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.test.directives.model.singleValue
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 
-abstract class AbstractStdLibSourcesLazyDeclarationResolveTest : AbstractFirLazyDeclarationResolveTestCase() {
+abstract class AbstractStdLibSourcesLazyDeclarationResolveTest : AbstractFirLazyDeclarationResolveOverAllPhasesTest() {
     override fun checkSession(firSession: LLFirResolveSession) {
         requireIsInstance<KtLibrarySourceModule>(firSession.useSiteKtModule)
     }
@@ -45,7 +45,7 @@ abstract class AbstractStdLibSourcesLazyDeclarationResolveTest : AbstractFirLazy
         val classDeclaration = findRegularClass(classId, module, resolveSession).findPsi() as KtClassOrObject
         val file = classDeclaration.containingFile as KtFile
 
-        doLazyResolveTest(file, testServices, moduleStructure, renderAllFiles = false) { firSession ->
+        doLazyResolveTest(file, testServices, renderAllFiles = false) { firSession ->
             val regularClass = findRegularClass(classId, module, firSession)
             val declarationToResolve = chooseMemberDeclarationIfNeeded(regularClass, moduleStructure, firSession)
             declarationToResolve.fir to fun(phase: FirResolvePhase) {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/firTestUtils.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConf
 import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.KtLibrarySourceModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
@@ -18,10 +19,12 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.services.PreAnalysisHandler
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.moduleStructure
 
 
 internal fun FirElement.renderWithClassName(): String =
@@ -89,3 +92,7 @@ fun Collection<FirFile>.getDeclarationsToResolve() = flatMap { it.collectAllElem
             declaration is FirValueParameter && declaration.containingFunctionSymbol is FirAnonymousFunctionSymbol ||
             declaration is FirProperty && declaration.isLocal
 }
+
+fun TestServices.allKtFiles(): List<KtFile> {
+    return moduleStructure.modules.flatMap(ktModuleProvider::getModuleFiles).mapNotNull { it as? KtFile }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/AbstractLowLevelApiLastModuleFirstFileTest.kt

@@ -0,0 +1,22 @@
+/*
+ * 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.test.base
+
+import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
+
+abstract class AbstractLowLevelApiLastModuleFirstFileTest : AbstractAnalysisApiBasedTest() {
+    final override fun doTestByModuleStructure(moduleStructure: TestModuleStructure, testServices: TestServices) {
+        val lastModule = moduleStructure.modules.last()
+        val firstKtFileFile = testServices.ktModuleProvider.getModuleFiles(lastModule).firstNotNullOf { it as? KtFile }
+        doTestByFileStructure(firstKtFileFile, moduleStructure, testServices)
+    }
+
+    abstract fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices)
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/AbstractLowLevelApiSingleFileTest.kt

@@ -13,8 +13,8 @@ import org.jetbrains.kotlin.test.services.TestServices
 
 abstract class AbstractLowLevelApiSingleFileTest : AbstractAnalysisApiBasedTest() {
     final override fun doTestByModuleStructure(moduleStructure: TestModuleStructure, testServices: TestServices) {
-        val singleModule = moduleStructure.modules.last()
-        val singleFile = testServices.ktModuleProvider.getModuleFiles(singleModule).filterIsInstance<KtFile>().first()
+        val singleModule = moduleStructure.modules.single()
+        val singleFile = testServices.ktModuleProvider.getModuleFiles(singleModule).filterIsInstance<KtFile>().single()
         doTestByFileStructure(singleFile, moduleStructure, testServices)
     }