diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt new file mode 100644 index 00000000000..78da3bc8296 --- /dev/null +++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt @@ -0,0 +1,453 @@ +/* + * 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.util + +import com.intellij.openapi.progress.ProgressManager +import com.intellij.psi.PsiElement +import kotlinx.collections.immutable.PersistentMap +import kotlinx.collections.immutable.persistentHashMapOf +import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget +import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.ContextKind +import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.Context +import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.FilterResponse +import org.jetbrains.kotlin.fir.FirElement +import org.jetbrains.kotlin.fir.FirSession +import org.jetbrains.kotlin.fir.declarations.* +import org.jetbrains.kotlin.fir.declarations.utils.isLocal +import org.jetbrains.kotlin.fir.expressions.* +import org.jetbrains.kotlin.fir.psi +import org.jetbrains.kotlin.fir.resolve.ResolutionMode +import org.jetbrains.kotlin.fir.resolve.SessionHolder +import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext +import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve +import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext +import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol +import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase +import org.jetbrains.kotlin.fir.types.ConeKotlinType +import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid +import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid + +internal object ContextCollector { + enum class ContextKind { + /** Represents the context of the declaration itself. */ + SELF, + + /** Represents the context inside the body (of a class, function, block, etc.) */ + BODY + } + + class Context( + val towerDataContext: FirTowerDataContext, + val smartCasts: Map, Set>, + ) + + enum class FilterResponse { + /** Store context for the element and continue the traversal. */ + CONTINUE, + + /** Store context for the element and stop traversal. */ + STOP, + + /** Skip the element and continue the traversal. */ + SKIP + } + + fun process(target: LLFirResolveTarget, holder: SessionHolder, filter: (PsiElement) -> FilterResponse): ContextProvider { + val pathIterator = (listOf(target.firFile) + target.path).iterator() + + val visitor = ContextCollectorVisitor(holder, filter) { + if (pathIterator.hasNext()) pathIterator.next() else null + } + + target.firFile.accept(visitor) + visitor.processTarget(target) + + return ContextProvider { element, kind -> visitor[element, kind] } + } +} + +internal fun interface ContextProvider { + operator fun get(element: PsiElement, kind: ContextKind): Context? +} + +private class ContextCollectorVisitor( + private val holder: SessionHolder, + private val filter: (PsiElement) -> FilterResponse, + private val interceptor: () -> FirElement? +) : FirDefaultVisitorVoid() { + operator fun get(element: PsiElement, kind: ContextKind): Context? { + return result[element to kind] + } + + private val session: FirSession + get() = holder.session + + private var isActive = true + + private val context = BodyResolveContext( + returnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve.Default, + dataFlowAnalyzerContext = DataFlowAnalyzerContext(session) + ) + + private var smartCasts: PersistentMap, Set> = persistentHashMapOf() + + private val result = HashMap, Context>() + + fun processTarget(target: LLFirResolveTarget) { + target.firFile.accept(this) + } + + override fun visitElement(element: FirElement) { + dumpContext(element.psi, ContextKind.SELF) + + onActive { + element.acceptChildren(this) + } + } + + private fun dumpContext(psi: PsiElement?, kind: ContextKind) { + ProgressManager.checkCanceled() + + if (psi == null) { + return + } + + val key = Pair(psi, kind) + if (key in result) { + return + } + + val response = filter(psi) + if (response != FilterResponse.SKIP) { + result[key] = Context(context.towerDataContext, smartCasts) + } + + if (response == FilterResponse.STOP) { + isActive = false + } + } + + override fun visitFile(file: FirFile) { + file.lazyResolveToPhase(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) + + context.withFile(file, holder) { + withInterceptor { + super.visitFile(file) + } + } + } + + override fun visitAnnotationCall(annotationCall: FirAnnotationCall) { + context.forAnnotation { + super.visitAnnotationCall(annotationCall) + } + } + + override fun visitRegularClass(regularClass: FirRegularClass) = withProcessor { + dumpContext(regularClass.psi, ContextKind.SELF) + + regularClass.lazyResolveToPhase(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) + + processSignatureAnnotations(regularClass) + + context.withContainingClass(regularClass) { + processList(regularClass.typeParameters) + + context.withRegularClass(regularClass, holder) { + dumpContext(regularClass.psi, ContextKind.BODY) + + onActive { + withInterceptor { + processChildren(regularClass) + } + } + } + } + + if (regularClass.isLocal) { + context.storeClassIfNotNested(regularClass, session) + } + } + + override fun visitConstructor(constructor: FirConstructor) = withProcessor { + dumpContext(constructor.psi, ContextKind.SELF) + + constructor.lazyResolveToPhase(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) + + processSignatureAnnotations(constructor) + + context.withConstructor(constructor) { + val owningClass = context.containerIfAny as? FirRegularClass + context.forConstructorParameters(constructor, owningClass, holder) { + processList(constructor.valueParameters) + } + + context.forConstructorBody(constructor, session) { + processList(constructor.valueParameters) + + dumpContext(constructor.psi, ContextKind.BODY) + + onActive { + constructor.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE) + process(constructor.body) + } + } + + onActive { + context.forDelegatedConstructorCall(constructor, owningClass = null, holder) { + process(constructor.delegatedConstructor) + } + } + + onActive { + processChildren(constructor) + } + } + } + + override fun visitEnumEntry(enumEntry: FirEnumEntry) { + context.forEnumEntry { + super.visitEnumEntry(enumEntry) + } + } + + override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) = withProcessor { + dumpContext(simpleFunction.psi, ContextKind.SELF) + + simpleFunction.lazyResolveToPhase(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) + + processSignatureAnnotations(simpleFunction) + + context.withSimpleFunction(simpleFunction, session) { + context.forFunctionBody(simpleFunction, holder) { + processList(simpleFunction.valueParameters) + + dumpContext(simpleFunction.psi, ContextKind.BODY) + + onActive { + simpleFunction.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE) + process(simpleFunction.body) + } + } + + onActive { + processChildren(simpleFunction) + } + } + } + + override fun visitProperty(property: FirProperty) = withProcessor { + dumpContext(property.psi, ContextKind.SELF) + + property.lazyResolveToPhase(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) + + processSignatureAnnotations(property) + + onActive { + property.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE) + + context.withProperty(property) { + context.forPropertyInitializer { + onActive { + process(property.initializer) + } + + onActive { + process(property.delegate) + } + + onActive { + process(property.backingField) + } + } + + onActive { + processChildren(property) + } + } + } + + if (property.isLocal) { + context.storeVariable(property, session) + } + } + + override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor) = withProcessor { + dumpContext(propertyAccessor.psi, ContextKind.SELF) + + processSignatureAnnotations(propertyAccessor) + + context.withPropertyAccessor(propertyAccessor.propertySymbol.fir, propertyAccessor, holder) { + dumpContext(propertyAccessor.psi, ContextKind.BODY) + + onActive { + processChildren(propertyAccessor) + } + } + } + + override fun visitValueParameter(valueParameter: FirValueParameter) = withProcessor { + dumpContext(valueParameter.psi, ContextKind.SELF) + + processSignatureAnnotations(valueParameter) + + context.withValueParameter(valueParameter, session) { + dumpContext(valueParameter.psi, ContextKind.SELF) + + onActive { + processChildren(valueParameter) + } + } + } + + override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer) = withProcessor { + dumpContext(anonymousInitializer.psi, ContextKind.SELF) + + anonymousInitializer.lazyResolveToPhase(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) + + processSignatureAnnotations(anonymousInitializer) + + context.withAnonymousInitializer(anonymousInitializer, session) { + dumpContext(anonymousInitializer.psi, ContextKind.BODY) + + onActive { + anonymousInitializer.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE) + processChildren(anonymousInitializer) + } + } + } + + override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) = withProcessor { + dumpContext(anonymousFunction.psi, ContextKind.SELF) + + processSignatureAnnotations(anonymousFunction) + + context.withAnonymousFunction(anonymousFunction, holder, ResolutionMode.ContextIndependent) { + context.forFunctionBody(anonymousFunction, holder) { + for (parameter in anonymousFunction.valueParameters) { + process(parameter) + context.storeVariable(parameter, holder.session) + } + + dumpContext(anonymousFunction.psi, ContextKind.BODY) + + onActive { + process(anonymousFunction.body) + } + } + + onActive { + processChildren(anonymousFunction) + } + } + } + + override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) = withProcessor { + dumpContext(anonymousObject.psi, ContextKind.SELF) + + processSignatureAnnotations(anonymousObject) + + context.withAnonymousObject(anonymousObject, holder) { + dumpContext(anonymousObject.psi, ContextKind.BODY) + + onActive { + processChildren(anonymousObject) + } + } + } + + override fun visitBlock(block: FirBlock) = withProcessor { + dumpContext(block.psi, ContextKind.SELF) + + context.forBlock(session) { + processChildren(block) + + dumpContext(block.psi, ContextKind.BODY) + } + } + + override fun visitSmartCastExpression(smartCastExpression: FirSmartCastExpression) { + if (smartCastExpression.isStable) { + val symbol = smartCastExpression.originalExpression.toResolvedCallableSymbol() + if (symbol != null) { + val previousSmartCasts = smartCasts + try { + smartCasts = smartCasts.put(symbol, smartCastExpression.typesFromSmartCast.toSet()) + super.visitSmartCastExpression(smartCastExpression) + } finally { + smartCasts = previousSmartCasts + } + } + } + + super.visitSmartCastExpression(smartCastExpression) + } + + @ContextCollectorDsl + private fun Processor.processSignatureAnnotations(declaration: FirDeclaration) { + for (annotation in declaration.annotations) { + onActive { + process(annotation) + } + } + } + + private inline fun withProcessor(block: Processor.() -> Unit) { + Processor(this).block() + } + + private class Processor(private val delegate: FirVisitorVoid) { + private val elementsToSkip = HashSet() + + @ContextCollectorDsl + fun process(element: FirElement?) { + if (element != null) { + element.accept(delegate) + elementsToSkip += element + } + } + + @ContextCollectorDsl + fun processList(elements: Collection) { + for (element in elements) { + process(element) + elementsToSkip += element + } + } + + @ContextCollectorDsl + fun processChildren(element: FirElement) { + val visitor = FilteringVisitor(delegate, elementsToSkip) + element.acceptChildren(visitor) + } + } + + private class FilteringVisitor(val delegate: FirVisitorVoid, val elementsToSkip: Set) : FirVisitorVoid() { + override fun visitElement(element: FirElement) { + if (element !in elementsToSkip) { + element.accept(delegate) + } + } + } + + private inline fun withInterceptor(block: () -> Unit) { + val target = interceptor() + if (target != null) { + target.accept(this) + } else { + block() + } + } + + private inline fun onActive(block: () -> Unit) { + if (isActive) { + block() + } + } +} + +@DslMarker +private annotation class ContextCollectorDsl \ No newline at end of file