~ partial resolve
diff --git a/analysis/low-level-api-fir/resources/META-INF/analysis-api/low-level-api-fir.xml b/analysis/low-level-api-fir/resources/META-INF/analysis-api/low-level-api-fir.xml
index db46e22..3405749 100644
--- a/analysis/low-level-api-fir/resources/META-INF/analysis-api/low-level-api-fir.xml
+++ b/analysis/low-level-api-fir/resources/META-INF/analysis-api/low-level-api-fir.xml
@@ -22,6 +22,13 @@
restartRequired="true"
/>
+ <registryKey
+ defaultValue="true"
+ description="Perform partial body analysis in Kotlin. Speeds up occasional references resolution, e.g. in Find Usages."
+ key="kotlin.analysis.partialBodyAnalysis"
+ restartRequired="true"
+ />
+
<projectService
serviceInterface="org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.LLFirElementByPsiElementChooser"
serviceImplementation="org.jetbrains.kotlin.analysis.low.level.api.fir.services.LLRealFirElementByPsiElementChooser"
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirPartialBodyResolveTarget.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirPartialBodyResolveTarget.kt
new file mode 100644
index 0000000..b9d5a98
--- /dev/null
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirPartialBodyResolveTarget.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2024 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.api.FirDesignation
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isPartialBodyResolvable
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.psi.KtElement
+
+internal class LLFirPartialBodyResolveTarget(
+ designation: FirDesignation,
+ val request: LLPartialBodyResolveRequest
+) : LLFirResolveTarget(designation) {
+ override fun visitTargetElement(element: FirElementWithResolveState, visitor: LLFirResolveTargetVisitor) {
+ visitor.performAction(element)
+ }
+}
+
+internal class LLPartialBodyResolveRequest(
+ val target: FirDeclaration,
+ val totalPsiStatementCount: Int,
+ val targetPsiStatementCount: Int,
+ val stopElement: KtElement?
+) {
+ init {
+ require(target.isPartialBodyResolvable)
+ }
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirSingleResolveTarget.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirSingleResolveTarget.kt
index 91d0ea6..383f4c2 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirSingleResolveTarget.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLFirSingleResolveTarget.kt
@@ -9,6 +9,11 @@
import org.jetbrains.kotlin.fir.FirElementWithResolveState
import org.jetbrains.kotlin.fir.declarations.*
+private object PartialBodyResolveStateKey : FirDeclarationDataKey()
+
+internal var FirDeclaration.partialBodyResolveState: LLPartialBodyResolveState?
+ by FirDeclarationDataRegistry.data(PartialBodyResolveStateKey)
+
/**
* [LLFirResolveTarget] representing single target to resolve. The [target] can be any of [FirElementWithResolveState]
*/
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLPartialBodyResolveState.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLPartialBodyResolveState.kt
new file mode 100644
index 0000000..eb46fd2
--- /dev/null
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/targets/LLPartialBodyResolveState.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2010-2024 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.declarations.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
+
+/**
+ * Represents the incomplete body resolve state.
+ * The attribute is only present if the function body was analyzed partially.
+ */
+internal data class LLPartialBodyResolveState(
+ val totalPsiStatementCount: Int,
+ val analyzedPsiStatementCount: Int,
+ val analyzedFirStatementCount: Int,
+ val performedAnalysesCount: Int,
+ val analysisStateSnapshot: LLPartialBodyResolveSnapshot?
+) {
+ val isFullyAnalyzed: Boolean
+ get() = totalPsiStatementCount == analyzedPsiStatementCount
+
+ override fun toString(): String {
+ return "$analyzedFirStatementCount($analyzedPsiStatementCount/$totalPsiStatementCount) #$performedAnalysesCount"
+ }
+}
+
+internal class LLPartialBodyResolveSnapshot(
+ val towerDataContext: FirTowerDataContext,
+ val dataFlowAnalyzerContext: DataFlowAnalyzerContext
+)
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
index d4f61a4..6af2975 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
@@ -11,6 +11,7 @@
import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FirElementsRecorder
import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.declarationCanBeLazilyResolved
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isPartialBodyResolvable
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.parentsWithSelfCodeFragmentAware
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.requireTypeIntersectionWith
import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
@@ -121,7 +122,18 @@
val structureElement = fileStructure.getStructureElementFor(element, nonLocalContainer)
val mappings = structureElement.mappings
- return mappings.getFir(psi)
+
+ val firElement = mappings.getFir(psi)
+
+ if (firElement is FirElementWithResolveState) {
+ // Partially resolvable declarations might have unresolved bodies in the mapping.
+ // Here we forcibly resolve them to obey to the 'getOrBuildFirFor()' contract.
+ if (firElement.isPartialBodyResolvable && firElement.resolvePhase < FirResolvePhase.BODY_RESOLVE) {
+ firElement.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+ }
+ }
+
+ return firElement
}
private inline fun <T : KtElement, E : PsiElement> getFirForNonBodyElement(
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt
index 9922142..6a5dcb8 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt
@@ -6,7 +6,9 @@
package org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder
import com.intellij.openapi.util.registry.Registry
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.partialBodyResolveState
import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirLazyResolveContractChecker
+import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.PartialBodyAnalysisSuspendedException
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.lockWithPCECheck
@@ -132,10 +134,15 @@
try {
action()
} catch (e: Throwable) {
- exceptionOccurred = true
+ if (e !is PartialBodyAnalysisSuspendedException) {
+ exceptionOccurred = true
+ }
throw e
} finally {
- val newPhase = if (updatePhase && !exceptionOccurred) toPhase else stateSnapshot.resolvePhase
+ val newPhase = when {
+ !updatePhase || exceptionOccurred -> stateSnapshot.resolvePhase
+ else -> computeNewPhase(stateSnapshot, toPhase)
+ }
unlock(toPhase = newPhase)
}
@@ -149,6 +156,19 @@
}
}
+ private fun FirElementWithResolveState.computeNewPhase(stateSnapshot: FirResolveState, toPhase: FirResolvePhase): FirResolvePhase {
+ if (this is FirDeclaration && toPhase == FirResolvePhase.BODY_RESOLVE) {
+ val state = partialBodyResolveState
+ if (state != null && !state.isFullyAnalyzed) {
+ // We only update the phase to BODY_RESOLVE if all statements are resolved.
+ // Otherwise, we set (BODY_RESOLVE - 1), so the next BODY_RESOLVE phase run can finish the analysis.
+ return stateSnapshot.resolvePhase
+ }
+ }
+
+ return toPhase
+ }
+
private fun waitOnBarrier(
stateSnapshot: FirInProcessOfResolvingToPhaseStateWithBarrier,
): Boolean {
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileElementFactory.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileElementFactory.kt
index fa03bf1c..ed7f713 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileElementFactory.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileElementFactory.kt
@@ -35,11 +35,13 @@
}
else -> {
- firDeclaration.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
if (firDeclaration is FirPrimaryConstructor) {
+ firDeclaration.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
firDeclaration.valueParameters.forEach { parameter ->
parameter.correspondingProperty?.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
}
+ } else {
+ firDeclaration.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE.previous)
}
DeclarationStructureElement(firFile, firDeclaration, moduleComponents)
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt
index 06e4615..59ed0aa 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt
@@ -5,17 +5,40 @@
package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
+import com.intellij.openapi.util.registry.Registry
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiErrorElement
+import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.persistentHashMapOf
+import kotlinx.collections.immutable.persistentMapOf
import org.jetbrains.kotlin.KtFakeSourceElementKind
import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBodyResolveRequest
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBodyResolveState
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.partialBodyResolveState
import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.*
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirResolveDesignationCollector
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirResolvableSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.body
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isPartialBodyResolvable
import org.jetbrains.kotlin.fir.FirElement
import org.jetbrains.kotlin.fir.correspondingProperty
import org.jetbrains.kotlin.fir.declarations.*
import org.jetbrains.kotlin.fir.declarations.impl.FirPrimaryConstructor
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.realPsi
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
import org.jetbrains.kotlin.fir.visitors.FirVisitor
import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
+import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
+import kotlin.collections.indexOf
/**
* Collects [KT -> FIR][KtToFirMapping] mapping and [diagnostics][FileStructureElementDiagnostics] for [declaration].
@@ -28,8 +51,9 @@
internal sealed class FileStructureElement(
val declaration: FirDeclaration,
val diagnostics: FileStructureElementDiagnostics,
+ elementProvider: FirElementProvider = EagerFirElementProvider(declaration)
) {
- val mappings: KtToFirMapping = KtToFirMapping(declaration)
+ val mappings: KtToFirMapping = KtToFirMapping(elementProvider)
companion object {
fun recorderFor(fir: FirDeclaration): FirElementsRecorder = when (fir) {
@@ -41,14 +65,175 @@
}
}
-internal class KtToFirMapping(firElement: FirDeclaration) {
+private typealias FirElementProvider = (KtElement) -> FirElement?
+
+private class EagerFirElementProvider(firDeclaration: FirDeclaration) : FirElementProvider {
private val mapping = FirElementsRecorder.recordElementsFrom(
- firElement = firElement,
- recorder = FileStructureElement.recorderFor(firElement),
+ firElement = firDeclaration,
+ recorder = FileStructureElement.recorderFor(firDeclaration),
)
- fun getElement(ktElement: KtElement): FirElement? {
- return mapping[ktElement]
+ override fun invoke(element: KtElement): FirElement? {
+ return mapping[element]
+ }
+}
+
+private class PartialBodyFirElementProvider(
+ private val declaration: FirDeclaration,
+ private val psiDeclaration: KtDeclaration,
+ private val psiBlock: KtBlockExpression,
+ private val psiStatements: List<KtExpression>,
+ private val session: LLFirResolvableModuleSession
+) : FirElementProvider {
+ /**
+ * Contains the latest known partial body resolution state.
+ *
+ * Initially, the [lastState] is empty, even though the declaration itself may already be partially resolved.
+ * On querying the mapping (by calling [invoke]), the actual resolved state is synchronized with the [lastState],
+ * and all missing elements are added to [bodyMappings].
+ */
+ @Volatile
+ private var lastState = LLPartialBodyResolveState(
+ totalPsiStatementCount = psiStatements.size,
+ analyzedPsiStatementCount = 0,
+ analyzedFirStatementCount = 0,
+ performedAnalysesCount = 0,
+ analysisStateSnapshot = null
+ )
+
+ /**
+ * Contains mappings for non-body elements.
+ */
+ private val signatureMappings: Map<KtElement, FirElement>
+
+ /**
+ * Contains collected mappings.
+ * Initially, only signature mappings are registered (the body is ignored).
+ * After consequent partial body analysis, elements from analyzed statements are appended.
+ */
+ @Volatile
+ private var bodyMappings: PersistentMap<KtElement, FirElement> = persistentMapOf()
+
+ // The body block cannot be cached on the element provider construction, as the body might be lazy at that point
+ private val bodyBlock: FirBlock
+ get() = declaration.body ?: error("Partial body element provider supports only declarations with bodies")
+
+ private val lockProvider: LLFirLockProvider
+ get() = session.moduleComponents.globalResolveComponents.lockProvider
+
+ init {
+ signatureMappings = persistentHashMapOf<KtElement, FirElement>()
+ .builder()
+ .also { declaration.accept(DeclarationStructureElement.SignatureRecorder(bodyBlock), it) }
+ .build()
+ }
+
+ override fun invoke(psiElement: KtElement): FirElement? {
+ val psiStatement = findTopmostStatement(psiElement)
+
+ if (psiStatement == null) {
+ // Didn't find a containing topmost statement. It's either a signature element or some unrelated one
+ return signatureMappings[psiElement]
+ }
+
+ val psiStatementIndex = psiStatements.indexOf(psiStatement)
+
+ checkWithAttachment(psiStatementIndex >= 0, { "The topmost statement was not found" }) {
+ withPsiEntry("statement", psiStatement)
+ withPsiEntry("declaration", psiDeclaration)
+ }
+
+ synchronized(this) {
+ if (lastState.analyzedPsiStatementCount > psiStatementIndex) {
+ // The statement is already analyzed and its children are in collected
+ return bodyMappings[psiElement]
+ }
+ }
+
+ performBodyAnalysis(psiStatementIndex)
+
+ synchronized(this) {
+ val newState = fetchPartialBodyResolveState()
+ if (newState != null) {
+ val lastCount = lastState.analyzedFirStatementCount
+ val newCount = newState.analyzedFirStatementCount
+
+ // There are newly analyzed statements, let's collect them
+ if (lastCount < newCount) {
+ val firBody = bodyBlock
+ require(firBody !is FirLazyBlock)
+
+ val newBodyMappingsBuilder = bodyMappings.builder()
+
+ for (index in lastCount until newCount) {
+ val firStatement = firBody.statements[index]
+ firStatement.accept(DeclarationStructureElement.Recorder, newBodyMappingsBuilder)
+ }
+
+ bodyMappings = newBodyMappingsBuilder.build()
+ lastState = newState
+ }
+ } else {
+ // The body has never been analyzed (otherwise the partial body resolve state should have been present)
+ bodyMappings = bodyMappings
+ .builder()
+ .also { bodyBlock.accept(FileStructureElement.recorderFor(declaration), it) }
+ .build()
+ }
+ }
+
+ return bodyMappings[psiElement]
+ }
+
+ private fun findTopmostStatement(psiElement: KtElement): KtElement? {
+ var previous: PsiElement? = null
+
+ for (current in psiElement.parentsWithSelf) {
+ when (current) {
+ psiBlock -> return previous as? KtElement
+ psiDeclaration -> return null
+ }
+
+ previous = current
+ }
+
+ return null
+ }
+
+ private fun fetchPartialBodyResolveState(): LLPartialBodyResolveState? {
+ var result: LLPartialBodyResolveState? = null
+ var isRun = false
+ lockProvider.withReadLock(declaration, FirResolvePhase.BODY_RESOLVE) {
+ result = declaration.partialBodyResolveState
+ isRun = true // 'withReadLock' does not call the lambda if the declaration already resolved
+ }
+ return if (isRun) result else declaration.partialBodyResolveState
+ }
+
+ private fun performBodyAnalysis(psiStatementIndex: Int) {
+ val psiStatementLimit = psiStatementIndex + 1
+ if (psiStatementLimit < psiStatements.size) {
+ val request = LLPartialBodyResolveRequest(
+ target = declaration,
+ totalPsiStatementCount = psiStatements.size,
+ targetPsiStatementCount = psiStatementLimit,
+ stopElement = psiStatements[psiStatementLimit]
+ )
+
+ val target = LLFirResolveDesignationCollector.getDesignationToResolveForPartialBody(request)
+ if (target != null) {
+ session.moduleComponents.firModuleLazyDeclarationResolver.lazyResolveTarget(target, FirResolvePhase.BODY_RESOLVE)
+ return
+ }
+ }
+
+ declaration.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+ }
+}
+
+internal class KtToFirMapping(private val elementProvider: FirElementProvider) {
+ private fun getElement(ktElement: KtElement): FirElement? {
+ return elementProvider(ktElement)
}
fun getFir(element: KtElement): FirElement? {
@@ -221,8 +406,59 @@
moduleComponents = moduleComponents,
)
),
+ elementProvider = createFirElementProvider(declaration),
) {
- object Recorder : FirElementsRecorder() {
+ private companion object {
+ private val IS_PARTIAL_RESOLVE_ENABLED = Registry.`is`("kotlin.analysis.partialBodyAnalysis", true)
+
+ private fun createFirElementProvider(declaration: FirDeclaration): FirElementProvider {
+ if (IS_PARTIAL_RESOLVE_ENABLED) {
+ val bodyBlock = declaration.body
+ if (declaration.isPartialBodyResolvable && bodyBlock != null && declaration.resolvePhase < FirResolvePhase.BODY_RESOLVE) {
+ require(declaration.resolvePhase == FirResolvePhase.BODY_RESOLVE.previous)
+
+ val isPartiallyResolvable = when (bodyBlock) {
+ is FirSingleExpressionBlock -> false
+ is FirEmptyExpressionBlock -> false
+ is FirLazyBlock -> true // Optimistic (however, below we also check the PSI statement count)
+ else -> bodyBlock.statements.size > 1
+ }
+
+ val session = declaration.llFirResolvableSession
+ val psiDeclaration = declaration.realPsi as? KtDeclaration
+ val psiBodyBlock = psiDeclaration?.bodyBlock
+ val psiStatements = psiBodyBlock?.statements
+
+ if (isPartiallyResolvable && session != null && psiStatements != null && psiStatements.size > 1) {
+ return PartialBodyFirElementProvider(declaration, psiDeclaration, psiBodyBlock, psiStatements, session)
+ }
+ }
+ }
+
+ declaration.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+ return EagerFirElementProvider(declaration)
+ }
+
+ private val KtDeclaration.bodyBlock: KtBlockExpression?
+ get() = when (this) {
+ is KtAnonymousInitializer -> body as? KtBlockExpression
+ is KtDeclarationWithBody -> bodyBlockExpression
+ else -> null
+ }
+ }
+
+ object Recorder : AbstractRecorder()
+
+ class SignatureRecorder(val bodyElement: FirElement) : AbstractRecorder() {
+ override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {
+ if (element === bodyElement) {
+ return
+ }
+ super.visitElement(element, data)
+ }
+ }
+
+ abstract class AbstractRecorder : FirElementsRecorder() {
override fun visitConstructor(constructor: FirConstructor, data: MutableMap<KtElement, FirElement>) {
super.visitConstructor(constructor, data)
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/inBlockModification.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/inBlockModification.kt
index f1bc902..8e61c78 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/inBlockModification.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/inBlockModification.kt
@@ -5,7 +5,9 @@
package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.partialBodyResolveState
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isPartialBodyResolvable
import org.jetbrains.kotlin.fir.declarations.*
import org.jetbrains.kotlin.fir.expressions.FirBlock
import org.jetbrains.kotlin.fir.expressions.FirExpression
@@ -180,6 +182,13 @@
get() = this is FirFunction && body?.statements?.firstOrNull() is FirContractCallBlock
private fun FirDeclaration.decreasePhase(newPhase: FirResolvePhase) {
+ if (isPartialBodyResolvable) {
+ val oldPhase = resolvePhase
+ if (oldPhase >= FirResolvePhase.BODY_RESOLVE.previous) {
+ partialBodyResolveState = null
+ }
+ }
+
@OptIn(ResolveStateAccess::class)
resolveState = newPhase.asResolveState()
}
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt
index a926875..86fd342 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt
@@ -10,6 +10,7 @@
import org.jetbrains.kotlin.analysis.low.level.api.fir.projectStructure.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.PartialBodyAnalysisSuspendedException
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.FirElementWithResolveState
@@ -102,6 +103,8 @@
if (toPhase == FirResolvePhase.IMPORTS) return
lazyResolveTargets(target, toPhase)
+ } catch (_: PartialBodyAnalysisSuspendedException) {
+ // Do nothing, partial body resolve is complete
} catch (e: Exception) {
handleExceptionFromResolve(e, target, toPhase)
}
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirResolveDesignationCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirResolveDesignationCollector.kt
index 49c6d5c..8e8d62e 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirResolveDesignationCollector.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirResolveDesignationCollector.kt
@@ -7,8 +7,10 @@
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirClassWithAllCallablesResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirPartialBodyResolveTarget
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeElementResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBodyResolveRequest
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.tryCollectDesignation
@@ -39,6 +41,12 @@
return getDesignationToResolve(target, ::LLFirWholeElementResolveTarget)
}
+ fun getDesignationToResolveForPartialBody(request: LLPartialBodyResolveRequest): LLFirResolveTarget? {
+ return getDesignationToResolve(request.target) {
+ LLFirPartialBodyResolveTarget(it, request)
+ }
+ }
+
private fun getDesignationToResolve(
target: FirElementWithResolveState,
resolveTarget: (FirDesignation) -> LLFirResolveTarget,
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
index 0b4d080..d23ccd8 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
@@ -5,19 +5,28 @@
package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
+import com.intellij.psi.util.descendantsOfType
import kotlinx.collections.immutable.toPersistentList
+import org.jetbrains.kotlin.KtPsiSourceElement
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getModule
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirPartialBodyResolveTarget
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBodyResolveRequest
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBodyResolveSnapshot
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBodyResolveState
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.partialBodyResolveState
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.codeFragmentScopeProvider
import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
import org.jetbrains.kotlin.analysis.low.level.api.fir.projectStructure.llFirModuleData
import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
+import org.jetbrains.kotlin.fir.FirElement
import org.jetbrains.kotlin.fir.FirElementWithResolveState
import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.isLocalMember
import org.jetbrains.kotlin.fir.declarations.*
import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
import org.jetbrains.kotlin.fir.expressions.*
@@ -33,6 +42,7 @@
import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference
import org.jetbrains.kotlin.fir.references.builder.buildExplicitThisReference
import org.jetbrains.kotlin.fir.resolve.FirCodeFragmentContext
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
import org.jetbrains.kotlin.fir.resolve.ScopeSession
import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
import org.jetbrains.kotlin.fir.resolve.codeFragmentContext
@@ -41,7 +51,11 @@
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForClass
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForFile
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForScript
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformerDispatcher
import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirExpressionsResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.writeResultType
import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractsDslNames
import org.jetbrains.kotlin.fir.scopes.DelicateScopeAPI
import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -50,10 +64,12 @@
import org.jetbrains.kotlin.fir.types.isResolved
import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtElement
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
+import java.util.concurrent.CancellationException
internal object LLFirBodyLazyResolver : LLFirLazyResolver(FirResolvePhase.BODY_RESOLVE) {
override fun createTargetResolver(target: LLFirResolveTarget): LLFirTargetResolver = LLFirBodyTargetResolver(target)
@@ -71,6 +87,215 @@
}
}
+internal class PartialBodyAnalysisSuspendedException : CancellationException("Partial body analysis was suspended")
+
+private class FirPartialBodyExpressionResolveTransformer(
+ transformer: FirAbstractBodyResolveTransformerDispatcher,
+ private val target: LLFirResolveTarget
+) : FirExpressionsResolveTransformer(transformer) {
+ private companion object {
+ // After a certain number of partial analyses,
+ // trigger the full analysis so we don't return to the same declaration over and over again
+ private const val MAX_ANALYSES_COUNT = 3
+ }
+
+ private var shouldStop = false
+
+ override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): E {
+ if (shouldStop) {
+ return element
+ }
+
+ return super.transformElement(element, data)
+ }
+
+ override fun transformBlock(block: FirBlock, data: ResolutionMode): FirStatement {
+ val declaration = context.containerIfAny
+
+ val isApplicable = declaration is FirDeclaration
+ && !declaration.isLocalMember
+ && declaration.isPartialBodyResolvable
+ && declaration.body == block
+ && block.statements.size > 1
+
+ if (!isApplicable) {
+ return super.transformBlock(block, data)
+ }
+
+ require(data is ResolutionMode.ContextIndependent)
+
+ val state = declaration.partialBodyResolveState
+
+ if (target is LLFirPartialBodyResolveTarget && (state == null || state.performedAnalysesCount < MAX_ANALYSES_COUNT)) {
+ transformPartially(target.request, block, data, state)
+ } else {
+ transformFully(declaration, block, data, state)
+ }
+
+ return block
+ }
+
+ private fun transformPartially(
+ request: LLPartialBodyResolveRequest,
+ block: FirBlock,
+ data: ResolutionMode,
+ state: LLPartialBodyResolveState?
+ ): FirStatement {
+ val declaration = target.target as FirDeclaration
+
+ if (state == null) {
+ if (request.stopElement == null) {
+ // Without a 'stopElement', we resolve "the rest of the body".
+ // We didn't analyze the body yet, though.
+ // So here we just delegate to the non-partial implementation.
+ require(request.targetPsiStatementCount == request.totalPsiStatementCount)
+ return super.transformBlock(block, data)
+ }
+
+ context.forBlock(session) {
+ dataFlowAnalyzer.enterBlock(block)
+ if (transformStatementsPartially(request, block, data, startIndex = 0, performedAnalysesCount = 0)) {
+ dataFlowAnalyzer.exitBlock(block)
+ }
+ }
+
+ return block
+ }
+
+ if (state.analyzedPsiStatementCount >= request.targetPsiStatementCount) {
+ // Nothing to analyze
+ return block
+ }
+
+ val resolveSnapshot = state.analysisStateSnapshot
+ checkWithAttachment(resolveSnapshot != null, { "Snapshot should be available for a partially analyzed declaration" }) {
+ withFirEntry("target", declaration)
+ withEntry("state", state) { it.toString() }
+ }
+
+ // Run analysis with the tower data context restored from the snapshot
+ context.withTowerDataContext(resolveSnapshot.towerDataContext) {
+ // Also, restore the previous state of the data flow graph from the snapshot
+ context.dataFlowAnalyzerContext.resetFrom(resolveSnapshot.dataFlowAnalyzerContext)
+
+ val isAnalyzedEntirely = transformStatementsPartially(
+ request, block, data,
+ startIndex = state.analyzedFirStatementCount,
+ performedAnalysesCount = state.performedAnalysesCount
+ )
+
+ if (isAnalyzedEntirely) {
+ dataFlowAnalyzer.exitBlock(block)
+ }
+ }
+
+ return block
+ }
+
+ private fun transformStatementsPartially(
+ request: LLPartialBodyResolveRequest,
+ block: FirBlock,
+ data: ResolutionMode,
+ startIndex: Int,
+ performedAnalysesCount: Int,
+ ): Boolean {
+ val declaration = request.target
+
+ val stopElement = request.stopElement
+ val stopElements = stopElement?.descendantsOfType<KtElement>(childrenFirst = false)?.toList().orEmpty()
+
+ var index = 0
+ val iterator = (block.statements as MutableList<FirStatement>).listIterator()
+
+ while (iterator.hasNext()) {
+ val statement = iterator.next()
+
+ if (index >= startIndex) {
+ if (stopElement != null && !shouldTransform(statement, stopElement, stopElements)) {
+ // Here we reached a stop element.
+ // It means all statements up to the target one are now analyzed.
+ // So now we save the current context and suspend further analysis.
+ declaration.partialBodyResolveState = LLPartialBodyResolveState(
+ totalPsiStatementCount = request.totalPsiStatementCount,
+ analyzedPsiStatementCount = request.targetPsiStatementCount,
+ analyzedFirStatementCount = index,
+ performedAnalysesCount = performedAnalysesCount + 1,
+ analysisStateSnapshot = LLPartialBodyResolveSnapshot(
+ towerDataContext = context.towerDataContext.createSnapshot(keepMutable = true),
+ dataFlowAnalyzerContext = context.dataFlowAnalyzerContext
+ )
+ )
+ shouldStop = true
+ throw PartialBodyAnalysisSuspendedException()
+ return false
+ }
+
+ val newStatement = statement.transform<FirStatement, ResolutionMode>(transformer, data)
+ if (statement !== newStatement) {
+ iterator.set(newStatement)
+ }
+ }
+
+ index += 1
+ }
+
+ // Nothing stopped us from analyzing all statements for some reason.
+ // Most likely, we missed the stop element.
+ // Let's still wrap things out – the function is now fully analyzed.
+ block.transformOtherChildren(transformer, data)
+
+ // This makes the compiler think the declaration is fully resolved (see 'FirExpression.isResolved')
+ block.writeResultType(session)
+
+ declaration.partialBodyResolveState = LLPartialBodyResolveState(
+ totalPsiStatementCount = request.totalPsiStatementCount,
+ analyzedPsiStatementCount = request.totalPsiStatementCount,
+ analyzedFirStatementCount = index,
+ performedAnalysesCount = 1,
+ analysisStateSnapshot = null
+ )
+
+ return true
+ }
+
+ private fun transformFully(
+ declaration: FirDeclaration,
+ block: FirBlock,
+ data: ResolutionMode,
+ currentState: LLPartialBodyResolveState?
+ ): FirStatement {
+ if (currentState == null) {
+ // The declaration body is not resolved at all, and a full resolution is requested.
+ // So, here we delegate straight to the non-partial implementation.
+ return super.transformBlock(block, data)
+ }
+
+ val request = LLPartialBodyResolveRequest(
+ target = declaration,
+ totalPsiStatementCount = currentState.totalPsiStatementCount,
+ targetPsiStatementCount = currentState.totalPsiStatementCount,
+ stopElement = null
+ )
+
+ // Otherwise, use the partial resolve to finish the ongoing resolution
+ return transformPartially(request, block, data, currentState)
+ }
+
+ private fun shouldTransform(element: FirElement, stopElement: KtElement, stopElements: List<KtElement>): Boolean {
+ val source = element.source
+ if (source is KtPsiSourceElement) {
+ if (source.psi == stopElement) {
+ return false
+ } else if (source.psi in stopElements) {
+ println("found in children!")
+ return false
+ }
+ }
+
+ return true
+ }
+}
+
/**
* This resolver is responsible for [BODY_RESOLVE][FirResolvePhase.BODY_RESOLVE] phase.
*
@@ -93,17 +318,28 @@
* @see FirBodyResolveTransformer
* @see FirResolvePhase.BODY_RESOLVE
*/
-private class LLFirBodyTargetResolver(target: LLFirResolveTarget) : LLFirAbstractBodyTargetResolver(
+private class LLFirBodyTargetResolver(private val target: LLFirResolveTarget) : LLFirAbstractBodyTargetResolver(
target,
FirResolvePhase.BODY_RESOLVE,
) {
- override val transformer = object : FirBodyResolveTransformer(
+ override val transformer = BodyTransformerDispatcher()
+
+ inner class BodyTransformerDispatcher : FirAbstractBodyResolveTransformerDispatcher(
resolveTargetSession,
phase = resolverPhase,
implicitTypeOnly = false,
- scopeSession = resolveTargetScopeSession,
+ scopeSession = resolveTargetSession.getScopeSession(),
returnTypeCalculator = createReturnTypeCalculator(),
+ expandTypeAliases = true
) {
+ override val expressionsTransformer: FirExpressionsResolveTransformer = if (target.target.isPartialBodyResolvable) {
+ FirPartialBodyExpressionResolveTransformer(this, target)
+ } else {
+ FirExpressionsResolveTransformer(this)
+ }
+
+ override val declarationsTransformer: FirDeclarationsResolveTransformer = FirDeclarationsResolveTransformer(this)
+
override val preserveCFGForClasses: Boolean get() = false
override val buildCfgForScripts: Boolean get() = false
override val buildCfgForFiles: Boolean get() = false
@@ -357,7 +593,15 @@
}
override fun rawResolve(target: FirElementWithResolveState) {
- super.rawResolve(target)
+ try {
+ super.rawResolve(target)
+ } catch (e: Throwable) {
+ if (e is PartialBodyAnalysisSuspendedException) {
+ // We successfully analyzed some part of the body so need to keep track of it
+ LLFirDeclarationModificationService.bodyResolved(target, resolverPhase)
+ }
+ throw e
+ }
LLFirDeclarationModificationService.bodyResolved(target, resolverPhase)
}
@@ -368,7 +612,20 @@
builder.add(FirCodeFragment::block, FirCodeFragment::replaceBlock, ::blockGuard)
}
- val ANONYMOUS_INITIALIZER: StateKeeper<FirAnonymousInitializer, FirDesignation> = stateKeeper { builder, _, _ ->
+ val PARTIAL_BODY_RESOLVABLE: StateKeeper<FirDeclaration, FirDesignation> = stateKeeper { builder, declaration, context ->
+ builder.add(
+ provider = FirDeclaration::partialBodyResolveState,
+ mutator = { _, _ ->
+ // On exceptions, remove the partial body resolve state so the function will be analyzed from the beginning
+ null
+ }
+ )
+ }
+
+ val ANONYMOUS_INITIALIZER: StateKeeper<FirAnonymousInitializer, FirDesignation> = stateKeeper { builder, initializer, designation ->
+ builder.add(PARTIAL_BODY_RESOLVABLE, designation)
+ preserveResolvedStatements(builder, initializer)
+
builder.add(FirAnonymousInitializer::body, FirAnonymousInitializer::replaceBody, ::blockGuard)
builder.add(FirAnonymousInitializer::controlFlowGraphReference, FirAnonymousInitializer::replaceControlFlowGraphReference)
}
@@ -385,7 +642,8 @@
builder.add(FirFunction::returnTypeRef, FirFunction::replaceReturnTypeRef)
if (!isCallableWithSpecialBody(function)) {
- preserveContractBlock(builder, function)
+ builder.add(PARTIAL_BODY_RESOLVABLE, designation)
+ preserveResolvedStatements(builder, function)
builder.add(FirFunction::body, FirFunction::replaceBody, ::blockGuard)
builder.entityList(function.valueParameters, VALUE_PARAMETER, designation)
@@ -439,7 +697,20 @@
}
}
-private fun StateKeeperScope<FirFunction, FirDesignation>.preserveContractBlock(builder: StateKeeperBuilder, function: FirFunction) {
+@Suppress("CONTEXT_RECEIVERS_DEPRECATED")
+private fun StateKeeperScope<FirAnonymousInitializer, FirDesignation>.preserveResolvedStatements(
+ builder: StateKeeperBuilder,
+ initializer: FirAnonymousInitializer
+) {
+ preservePartialBodyResolveResult(builder, initializer, FirAnonymousInitializer::body)
+}
+
+@Suppress("CONTEXT_RECEIVERS_DEPRECATED")
+private fun StateKeeperScope<FirFunction, FirDesignation>.preserveResolvedStatements(builder: StateKeeperBuilder, function: FirFunction) {
+ if (preservePartialBodyResolveResult(builder, function, FirFunction::body)) {
+ return
+ }
+
val oldBody = function.body
if (oldBody == null || oldBody is FirLazyBlock) {
return
@@ -477,6 +748,37 @@
}
}
+@Suppress("CONTEXT_RECEIVERS_DEPRECATED")
+private fun <T : FirDeclaration> StateKeeperScope<T, FirDesignation>.preservePartialBodyResolveResult(
+ builder: StateKeeperBuilder,
+ declaration: T,
+ bodySupplier: (T) -> FirBlock?,
+): Boolean {
+ val oldBody = bodySupplier(declaration)
+ if (oldBody == null || oldBody is FirLazyBlock) {
+ return false
+ }
+
+ val state = declaration.partialBodyResolveState
+ if (state == null) {
+ return false
+ }
+
+ builder.postProcess {
+ val newBody = bodySupplier(declaration)
+ if (newBody != null && newBody.statements.isNotEmpty()) {
+ require(oldBody.statements.size == newBody.statements.size) { "Bodies do not match" }
+
+ val newBodyStatements = newBody.statements as MutableList<FirStatement>
+ for (index in 0..<state.analyzedFirStatementCount) {
+ newBodyStatements[index] = oldBody.statements[index]
+ }
+ }
+ }
+
+ return true
+}
+
private val FirFunction.isCertainlyResolved: Boolean
get() {
if (this is FirPropertyAccessor) {
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyResolver.kt
index 156271e..3f74ef7 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyResolver.kt
@@ -5,6 +5,7 @@
package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirPartialBodyResolveTarget
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
@@ -32,7 +33,10 @@
)
resolver.resolveDesignation()
- target.forEachTarget(::checkIsResolved)
+
+ if (target !is LLFirPartialBodyResolveTarget) {
+ target.forEachTarget(::checkIsResolved)
+ }
checkCanceled()
}
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/StateKeeper.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/StateKeeper.kt
index b6cd62b..9d577aa 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/StateKeeper.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/StateKeeper.kt
@@ -272,6 +272,8 @@
prepareTarget(target)
preservedState.postProcess()
return action()
+ } catch (e: PartialBodyAnalysisSuspendedException) {
+ throw e
} catch (e: Throwable) {
preservedState?.restore()
throw e
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt
index 024676b..2fc03e8 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt
@@ -17,7 +17,9 @@
import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirProvider
import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirBlock
import org.jetbrains.kotlin.fir.psi
import org.jetbrains.kotlin.fir.realPsi
import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
@@ -232,6 +234,20 @@
}
}
+internal val FirElementWithResolveState.isPartialBodyResolvable: Boolean
+ get() = when (this) {
+ is FirFunction -> this !is FirErrorFunction
+ is FirAnonymousInitializer -> true
+ else -> false
+ }
+
+internal val FirElementWithResolveState.body: FirBlock?
+ get() = when (this) {
+ is FirFunction -> body
+ is FirAnonymousInitializer -> body
+ else -> null
+ }
+
/**
* Some "local" declarations are not local from the lazy resolution perspective.
*/
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.kt b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.kt
new file mode 100644
index 0000000..dc63561
--- /dev/null
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.kt
@@ -0,0 +1,7 @@
+fun test() {
+ consume(<expr>1 + 2</expr>)
+ consume(<expr_1>"foo"</expr_1>)
+ consume('b')
+}
+
+fun consume(a: Int) {}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.txt b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.txt
new file mode 100644
index 0000000..d59699c
--- /dev/null
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.txt
@@ -0,0 +1,39 @@
+KT element: KtValueArgument
+KT element text:
+1 + 2
+FIR element: FirIntegerLiteralOperatorCallImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+Int(1).R|kotlin/Int.plus|(Int(2))
+
+FIR FILE:
+FILE: [ResolvedTo(IMPORTS)] simple.kt
+ public final [ResolvedTo(ANNOTATION_ARGUMENTS)] [PartialBodyResolveStateKey=1(1/3) #1] fun test(): R|kotlin/Unit| {
+ R|/consume|(Int(1).R|kotlin/Int.plus|(Int(2)))
+ consume#(String(foo))
+ consume#(Char(b))
+ }
+ public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] a: R|kotlin/Int|): R|kotlin/Unit| {
+ }
+
+=====
+
+KT element: KtValueArgument
+KT element text:
+"foo"
+FIR element: FirLiteralExpressionImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+String(foo)
+
+FIR FILE:
+FILE: [ResolvedTo(IMPORTS)] simple.kt
+ public final [ResolvedTo(ANNOTATION_ARGUMENTS)] [PartialBodyResolveStateKey=2(2/3) #1] fun test(): R|kotlin/Unit| {
+ R|/consume|(Int(1).R|kotlin/Int.plus|(Int(2)))
+ R|/consume<Inapplicable(INAPPLICABLE): /consume>#|(String(foo))
+ consume#(Char(b))
+ }
+ public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] a: R|kotlin/Int|): R|kotlin/Unit| {
+ }
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.kt b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.kt
new file mode 100644
index 0000000..865cfec
--- /dev/null
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.kt
@@ -0,0 +1,7 @@
+fun test(obj: Any) {
+ if (<expr>obj</expr> !is String) return
+ consume(<expr_1>obj</expr_1>)
+ obj.toString()
+}
+
+fun <T : CharSequence> consume(obj: T) {}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.txt b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.txt
new file mode 100644
index 0000000..086f84c
--- /dev/null
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.txt
@@ -0,0 +1,48 @@
+KT element: KtNameReferenceExpression
+KT element text:
+obj
+FIR element: FirPropertyAccessExpressionImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+R|<local>/obj|
+
+FIR FILE:
+FILE: [ResolvedTo(IMPORTS)] simpleDataFlow.kt
+ public final [ResolvedTo(ANNOTATION_ARGUMENTS)] [PartialBodyResolveStateKey=1(1/3) #1] fun test([ResolvedTo(ANNOTATION_ARGUMENTS)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+ when () {
+ (R|<local>/obj| !is R|kotlin/String|) -> {
+ ^test Unit
+ }
+ }
+
+ consume#(obj#)
+ obj#.toString#()
+ }
+ public? final? [ResolvedTo(RAW_FIR)] fun <[ResolvedTo(RAW_FIR)] T : CharSequence> consume([ResolvedTo(RAW_FIR)] obj: T): R|kotlin/Unit| { LAZY_BLOCK }
+
+=====
+
+KT element: KtValueArgument
+KT element text:
+obj
+FIR element: FirSmartCastExpressionImpl
+FIR source kind: SmartCastExpression
+
+FIR element rendered:
+R|<local>/obj|
+
+FIR FILE:
+FILE: [ResolvedTo(IMPORTS)] simpleDataFlow.kt
+ public final [ResolvedTo(ANNOTATION_ARGUMENTS)] [PartialBodyResolveStateKey=2(2/3) #1] fun test([ResolvedTo(ANNOTATION_ARGUMENTS)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+ when () {
+ (R|<local>/obj| !is R|kotlin/String|) -> {
+ ^test Unit
+ }
+ }
+
+ R|/consume|<R|kotlin/String|>(R|<local>/obj|)
+ obj#.toString#()
+ }
+ public final [ResolvedTo(CONTRACTS)] fun <[ResolvedTo(CONTRACTS)] T : R|kotlin/CharSequence|> consume([ResolvedTo(CONTRACTS)] obj: R|T|): R|kotlin/Unit| {
+ }
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractGetOrBuildFirTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractGetOrBuildFirTest.kt
index 6aab2f5..1d6e351 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractGetOrBuildFirTest.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractGetOrBuildFirTest.kt
@@ -27,18 +27,43 @@
abstract class AbstractGetOrBuildFirTest : AbstractAnalysisApiBasedTest() {
override fun doTestByMainFile(mainFile: KtFile, mainModule: KtTestModule, testServices: TestServices) {
- val selectedElement = testServices.expressionMarkerProvider
- .getTopmostSelectedElementOfTypeByDirective(mainFile, mainModule) as KtElement
+ val results = resolveWithClearCaches(mainFile) { session ->
+ val elementsToAnalyze = sequence<KtElement> {
+ val expressionMarkerProvider = testServices.expressionMarkerProvider
- val actual = resolveWithClearCaches(mainFile) { session ->
- renderActualFir(
- fir = selectedElement.getOrBuildFir(session),
- ktElement = selectedElement,
- renderingOptions = testServices.firRenderingOptions,
- firFile = mainFile.getOrBuildFirFile(session),
- )
+ val firstCandidate = expressionMarkerProvider
+ .getTopmostSelectedElementOfTypeByDirective(mainFile, mainModule, qualifier = "")
+
+ yield(firstCandidate as KtElement)
+
+ var index = 1
+ while (true) {
+ val candidate = expressionMarkerProvider
+ .getTopmostSelectedElementOfTypeByDirectiveOrNull(mainFile, mainModule, qualifier = "$index")
+ ?: break
+ yield(candidate as KtElement)
+ index += 1
+ }
+ }.toList()
+
+ val renderingOptions = testServices.firRenderingOptions
+ .copy(renderKtText = elementsToAnalyze.size > 1)
+
+ val results = mutableListOf<String>()
+
+ for (element in elementsToAnalyze) {
+ results += renderActualFir(
+ fir = element.getOrBuildFir(session),
+ ktElement = element,
+ renderingOptions = renderingOptions,
+ firFile = mainFile.getOrBuildFirFile(session),
+ )
+ }
+
+ return@resolveWithClearCaches results
}
+ val actual = results.joinToString(separator = "\n\n=====\n\n")
testServices.assertions.assertEqualsToTestDataFileSibling(actual)
}
}
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/OutOfContentRootGetOrBuildFirTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/OutOfContentRootGetOrBuildFirTestGenerated.java
index a3869d6..153b6c4 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/OutOfContentRootGetOrBuildFirTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/OutOfContentRootGetOrBuildFirTestGenerated.java
@@ -1217,6 +1217,28 @@
}
@Nested
+ @TestMetadata("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis")
+ @TestDataPath("$PROJECT_ROOT")
+ public class PartialBodyAnalysis {
+ @Test
+ public void testAllFilesPresentInPartialBodyAnalysis() {
+ KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis"), Pattern.compile("^(.+)\\.(kt)$"), null, true);
+ }
+
+ @Test
+ @TestMetadata("simple.kt")
+ public void testSimple() {
+ runTest("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.kt");
+ }
+
+ @Test
+ @TestMetadata("simpleDataFlow.kt")
+ public void testSimpleDataFlow() {
+ runTest("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.kt");
+ }
+ }
+
+ @Nested
@TestMetadata("analysis/low-level-api-fir/testData/getOrBuildFir/qualifiedExpressions")
@TestDataPath("$PROJECT_ROOT")
public class QualifiedExpressions {
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptGetOrBuildFirTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptGetOrBuildFirTestGenerated.java
index 7098015..939913f 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptGetOrBuildFirTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptGetOrBuildFirTestGenerated.java
@@ -743,6 +743,16 @@
}
@Nested
+ @TestMetadata("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis")
+ @TestDataPath("$PROJECT_ROOT")
+ public class PartialBodyAnalysis {
+ @Test
+ public void testAllFilesPresentInPartialBodyAnalysis() {
+ KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis"), Pattern.compile("^(.+)\\.(kts)$"), null, true);
+ }
+ }
+
+ @Nested
@TestMetadata("analysis/low-level-api-fir/testData/getOrBuildFir/qualifiedExpressions")
@TestDataPath("$PROJECT_ROOT")
public class QualifiedExpressions {
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/SourceGetOrBuildFirTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/SourceGetOrBuildFirTestGenerated.java
index d170d84..4a3c432 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/SourceGetOrBuildFirTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/SourceGetOrBuildFirTestGenerated.java
@@ -1217,6 +1217,28 @@
}
@Nested
+ @TestMetadata("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis")
+ @TestDataPath("$PROJECT_ROOT")
+ public class PartialBodyAnalysis {
+ @Test
+ public void testAllFilesPresentInPartialBodyAnalysis() {
+ KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis"), Pattern.compile("^(.+)\\.(kt)$"), null, true);
+ }
+
+ @Test
+ @TestMetadata("simple.kt")
+ public void testSimple() {
+ runTest("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simple.kt");
+ }
+
+ @Test
+ @TestMetadata("simpleDataFlow.kt")
+ public void testSimpleDataFlow() {
+ runTest("analysis/low-level-api-fir/testData/getOrBuildFir/partialBodyAnalysis/simpleDataFlow.kt");
+ }
+ }
+
+ @Nested
@TestMetadata("analysis/low-level-api-fir/testData/getOrBuildFir/qualifiedExpressions")
@TestDataPath("$PROJECT_ROOT")
public class QualifiedExpressions {
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/services/FirRenderingOptions.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/services/FirRenderingOptions.kt
index d3dbeaa..90f3f7b 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/services/FirRenderingOptions.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/services/FirRenderingOptions.kt
@@ -8,7 +8,7 @@
import org.jetbrains.kotlin.test.services.TestService
import org.jetbrains.kotlin.test.services.TestServices
-internal class FirRenderingOptions(
+internal data class FirRenderingOptions(
val renderKtText: Boolean = false,
val renderKtFileName: Boolean = false,
val renderContainerSource: Boolean = false,
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
index 3424c16..29a367f8 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
@@ -43,10 +43,13 @@
import org.jetbrains.kotlin.util.OperatorNameConventions
class DataFlowAnalyzerContext(private val session: FirSession) {
- val graphBuilder: ControlFlowGraphBuilder = ControlFlowGraphBuilder()
- internal val variableAssignmentAnalyzer: FirLocalVariableAssignmentAnalyzer = FirLocalVariableAssignmentAnalyzer()
+ internal var graphBuilder: ControlFlowGraphBuilder = ControlFlowGraphBuilder()
+ private set
- var variableStorage: VariableStorage = VariableStorage(session)
+ internal var variableAssignmentAnalyzer: FirLocalVariableAssignmentAnalyzer = FirLocalVariableAssignmentAnalyzer()
+ private set
+
+ internal var variableStorage: VariableStorage = VariableStorage(session)
private set
private var assignmentCounter = 0
@@ -55,6 +58,15 @@
return assignmentCounter++
}
+ fun resetFrom(source: DataFlowAnalyzerContext) {
+ reset()
+
+ graphBuilder = source.graphBuilder
+ variableAssignmentAnalyzer = source.variableAssignmentAnalyzer
+ variableStorage = source.variableStorage
+ assignmentCounter = source.assignmentCounter
+ }
+
fun reset() {
graphBuilder.reset()
variableAssignmentAnalyzer.reset()
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
index cbc79db..cf81c62 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
@@ -137,6 +137,17 @@
}
}
+ @OptIn(PrivateForInline::class)
+ inline fun <T> withTowerDataContext(newContext: FirTowerDataContext, f: () -> T): T {
+ val initialContext = towerDataContext
+ return try {
+ replaceTowerDataContext(newContext)
+ f()
+ } finally {
+ replaceTowerDataContext(initialContext)
+ }
+ }
+
private inline fun <R> withLambdaBeingAnalyzedInDependentContext(lambda: FirAnonymousFunctionSymbol, l: () -> R): R {
anonymousFunctionsAnalyzedInDependentContext.add(lambda)
return try {
