[LL API] Add 'ContextCollector'
The traditional way to get a 'FirTowerDataContext' in K2 was to reparse the callable body and pass it through all resolution phases. It is terribly inefficient and requires a lot of magic. Currently, 'ContextCollector' is only going to be used to collect a code fragment context. After proper test coverage and a plenty of patches it's going to become a generic replacement for read-only on-air resolution.
This commit is contained in:
Yan Zhulanow
+453
@@ -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<FirBasedSymbol<*>, Set<ConeKotlinType>>,
|
||||||
|
)
|
||||||
|
|
||||||
|
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<FirBasedSymbol<*>, Set<ConeKotlinType>> = persistentHashMapOf()
|
||||||
|
|
||||||
|
private val result = HashMap<Pair<PsiElement, ContextKind>, 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<FirElement>()
|
||||||
|
|
||||||
|
@ContextCollectorDsl
|
||||||
|
fun process(element: FirElement?) {
|
||||||
|
if (element != null) {
|
||||||
|
element.accept(delegate)
|
||||||
|
elementsToSkip += element
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@ContextCollectorDsl
|
||||||
|
fun processList(elements: Collection<FirElement>) {
|
||||||
|
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<FirElement>) : 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
|
||||||
Reference in New Issue
Block a user