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 index 1a15b595dc9..8c1cb1c0e32 100644 --- 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 @@ -39,6 +39,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf import org.jetbrains.kotlin.util.PrivateForInline import org.jetbrains.kotlin.utils.yieldIfNotNull +import java.util.ArrayList object ContextCollector { enum class ContextKind { @@ -181,6 +182,8 @@ private class ContextCollectorVisitor( private var isActive = true + private val parents = ArrayList() + private val context = BodyResolveContext( returnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve.Default, dataFlowAnalyzerContext = DataFlowAnalyzerContext(session) @@ -192,7 +195,9 @@ private class ContextCollectorVisitor( dumpContext(element, ContextKind.SELF) onActive { - element.acceptChildren(this) + withParent(element) { + element.acceptChildren(this) + } } } @@ -229,7 +234,7 @@ private class ContextCollectorVisitor( // So here we modify the types inside the 'context', then make a snapshot, and restore the types back. val oldReceiverTypes = mutableListOf>() - val cfgNode = getControlFlowNode(fir) + val cfgNode = getClosestControlFlowNode(fir) if (cfgNode != null) { val flow = cfgNode.flow @@ -268,6 +273,23 @@ private class ContextCollectorVisitor( return Context(towerDataContextSnapshot, smartCasts) } + private fun getClosestControlFlowNode(fir: FirElement): CFGNode<*>? { + val selfNode = getControlFlowNode(fir) + if (selfNode != null) { + return selfNode + } + + // For some specific elements, such as types or references, there is usually no associated 'CFGNode'. + for (parent in parents.asReversed()) { + val parentNode = getControlFlowNode(parent) + if (parentNode != null) { + return parentNode + } + } + + return null + } + private fun getControlFlowNode(fir: FirElement): CFGNode<*>? { for (container in context.containers.asReversed()) { val cfgOwner = container as? FirControlFlowGraphOwner ?: continue @@ -324,7 +346,7 @@ private class ContextCollectorVisitor( } } - override fun visitRegularClass(regularClass: FirRegularClass) = withProcessor { + override fun visitRegularClass(regularClass: FirRegularClass) = withProcessor(regularClass) { dumpContext(regularClass, ContextKind.SELF) processSignatureAnnotations(regularClass) @@ -368,7 +390,7 @@ private class ContextCollectorVisitor( } } - override fun visitConstructor(constructor: FirConstructor) = withProcessor { + override fun visitConstructor(constructor: FirConstructor) = withProcessor(constructor) { dumpContext(constructor, ContextKind.SELF) processSignatureAnnotations(constructor) @@ -421,7 +443,7 @@ private class ContextCollectorVisitor( } } - override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) = withProcessor { + override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) = withProcessor(simpleFunction) { dumpContext(simpleFunction, ContextKind.SELF) processSignatureAnnotations(simpleFunction) @@ -447,7 +469,7 @@ private class ContextCollectorVisitor( } } - override fun visitProperty(property: FirProperty) = withProcessor { + override fun visitProperty(property: FirProperty) = withProcessor(property) { dumpContext(property, ContextKind.SELF) processSignatureAnnotations(property) @@ -494,7 +516,7 @@ private class ContextCollectorVisitor( * It's not accessible from the delegated constructor, it's just added to the * `Foo` class body. */ - override fun visitField(field: FirField) = withProcessor { + override fun visitField(field: FirField) = withProcessor(field) { dumpContext(field, ContextKind.SELF) processSignatureAnnotations(field) @@ -512,7 +534,7 @@ private class ContextCollectorVisitor( } } - override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor) = withProcessor { + override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor) = withProcessor(propertyAccessor) { dumpContext(propertyAccessor, ContextKind.SELF) processSignatureAnnotations(propertyAccessor) @@ -528,7 +550,7 @@ private class ContextCollectorVisitor( } } - override fun visitValueParameter(valueParameter: FirValueParameter) = withProcessor { + override fun visitValueParameter(valueParameter: FirValueParameter) = withProcessor(valueParameter) { dumpContext(valueParameter, ContextKind.SELF) processSignatureAnnotations(valueParameter) @@ -545,7 +567,7 @@ private class ContextCollectorVisitor( } - override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer) = withProcessor { + override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer) = withProcessor(anonymousInitializer) { dumpContext(anonymousInitializer, ContextKind.SELF) processSignatureAnnotations(anonymousInitializer) @@ -562,7 +584,7 @@ private class ContextCollectorVisitor( } } - override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) = withProcessor { + override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) = withProcessor(anonymousFunction) { dumpContext(anonymousFunction, ContextKind.SELF) processSignatureAnnotations(anonymousFunction) @@ -588,7 +610,7 @@ private class ContextCollectorVisitor( } - override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) = withProcessor { + override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) = withProcessor(anonymousObject) { dumpContext(anonymousObject, ContextKind.SELF) processSignatureAnnotations(anonymousObject) @@ -605,7 +627,7 @@ private class ContextCollectorVisitor( } - override fun visitBlock(block: FirBlock) = withProcessor { + override fun visitBlock(block: FirBlock) = withProcessor(block) { dumpContext(block, ContextKind.SELF) onActiveBody { @@ -626,8 +648,10 @@ private class ContextCollectorVisitor( } } - private inline fun withProcessor(block: Processor.() -> Unit) { - Processor(this).block() + private inline fun withProcessor(parent: FirElement, block: Processor.() -> Unit) { + withParent(parent) { + Processor(this).block() + } } private class Processor(private val delegate: FirVisitorVoid) { @@ -678,6 +702,15 @@ private class ContextCollectorVisitor( } } + private inline fun withParent(parent: FirElement, block: () -> Unit) { + parents.add(parent) + try { + block() + } finally { + parents.removeLast() + } + } + private inline fun onActive(block: () -> Unit) { if (isActive) { block() diff --git a/analysis/low-level-api-fir/testData/contextCollector/KT-61728.txt b/analysis/low-level-api-fir/testData/contextCollector/KT-61728.txt index 3db2c14994a..564862dda8e 100644 --- a/analysis/low-level-api-fir/testData/contextCollector/KT-61728.txt +++ b/analysis/low-level-api-fir/testData/contextCollector/KT-61728.txt @@ -25,6 +25,10 @@ Tower Data Context: FirValueParameterSymbol it: R|kotlin/String| Element 10 Scope: FirLocalScope +Smart Casts: + FirValueParameterSymbol param: R|kotlin/String?| + Types: + kotlin/Any FILE: [ResolvedTo(IMPORTS)] KT-61728.kt public final [ResolvedTo(BODY_RESOLVE)] fun foo([ResolvedTo(BODY_RESOLVE)] param: R|kotlin/String?|): R|kotlin/String?| { diff --git a/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.kt b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.kt new file mode 100644 index 00000000000..adb1ce714c4 --- /dev/null +++ b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.kt @@ -0,0 +1,17 @@ +fun test(a: Any) { + if (a !is Foo) { + return + } + + var result: Int = 0 + + loop@ while (true) { + if (!a.process()) { + break@loop + } + } +} + +interface Foo { + fun process(): Boolean +} \ No newline at end of file diff --git a/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.txt b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.txt new file mode 100644 index 00000000000..99bc5e27fe8 --- /dev/null +++ b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.txt @@ -0,0 +1,52 @@ +Tower Data Context: + Element 0 + Scope: FirDefaultStarImportingScope + Element 1 + Scope: FirDefaultSimpleImportingScope + Element 2 + Scope: FirExplicitStarImportingScope + Element 3 + Scope: FirDefaultSimpleImportingScope + Element 4 + Scope: FirDefaultSimpleImportingScope + Element 5 + Scope: FirPackageMemberScope + Element 6 + Scope: FirExplicitSimpleImportingScope + Element 7 + Scope: FirLocalScope + Properties: + FirValueParameterSymbol a: R|kotlin/Any| + Element 8 + Scope: FirLocalScope + Properties: + FirPropertySymbol lvar result: R|kotlin/Int| +Smart Casts: + FirValueParameterSymbol a: R|kotlin/Any| + Types: + Foo + kotlin/Any + +FILE: [ResolvedTo(IMPORTS)] onLabel.kt + public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] a: R|kotlin/Any|): R|kotlin/Unit| { + when () { + (R|/a| !is R|Foo|) -> { + ^test Unit + } + } + + [ResolvedTo(BODY_RESOLVE)] lvar result: R|kotlin/Int| = Int(0) + loop@while(Boolean(true)) { + when () { + R|/a|.R|/Foo.process|().R|kotlin/Boolean.not|() -> { + break@@@[Boolean(true)] + } + } + + } + + } + public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| { + public abstract [ResolvedTo(CONTRACTS)] fun process(): R|kotlin/Boolean| + + } diff --git a/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.kt b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.kt new file mode 100644 index 00000000000..500cb0588e5 --- /dev/null +++ b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.kt @@ -0,0 +1,13 @@ +fun test(a: Any) { + if (a !is Foo) { + return + } + + check(a.process()) +} + +fun check(condition: Boolean) {} + +interface Foo { + fun process(): Boolean +} \ No newline at end of file diff --git a/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.txt b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.txt new file mode 100644 index 00000000000..da3b7d00696 --- /dev/null +++ b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.txt @@ -0,0 +1,43 @@ +Tower Data Context: + Element 0 + Scope: FirDefaultStarImportingScope + Element 1 + Scope: FirDefaultSimpleImportingScope + Element 2 + Scope: FirExplicitStarImportingScope + Element 3 + Scope: FirDefaultSimpleImportingScope + Element 4 + Scope: FirDefaultSimpleImportingScope + Element 5 + Scope: FirPackageMemberScope + Element 6 + Scope: FirExplicitSimpleImportingScope + Element 7 + Scope: FirLocalScope + Properties: + FirValueParameterSymbol a: R|kotlin/Any| + Element 8 + Scope: FirLocalScope +Smart Casts: + FirValueParameterSymbol a: R|kotlin/Any| + Types: + Foo + kotlin/Any + +FILE: [ResolvedTo(IMPORTS)] onReference.kt + public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] a: R|kotlin/Any|): R|kotlin/Unit| { + when () { + (R|/a| !is R|Foo|) -> { + ^test Unit + } + } + + R|/check|(R|/a|.R|/Foo.process|()) + } + public final [ResolvedTo(CONTRACTS)] fun check([ResolvedTo(CONTRACTS)] condition: R|kotlin/Boolean|): R|kotlin/Unit| { + } + public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| { + public abstract [ResolvedTo(CONTRACTS)] fun process(): R|kotlin/Boolean| + + } diff --git a/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.kt b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.kt new file mode 100644 index 00000000000..3771645b806 --- /dev/null +++ b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.kt @@ -0,0 +1,13 @@ +fun test(a: Any) { + if (a !is Foo) { + return + } + + callInt>>() +} + +inline fun call() {} + +interface Foo { + fun process(): Boolean +} \ No newline at end of file diff --git a/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.txt b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.txt new file mode 100644 index 00000000000..462d51d59b7 --- /dev/null +++ b/analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.txt @@ -0,0 +1,43 @@ +Tower Data Context: + Element 0 + Scope: FirDefaultStarImportingScope + Element 1 + Scope: FirDefaultSimpleImportingScope + Element 2 + Scope: FirExplicitStarImportingScope + Element 3 + Scope: FirDefaultSimpleImportingScope + Element 4 + Scope: FirDefaultSimpleImportingScope + Element 5 + Scope: FirPackageMemberScope + Element 6 + Scope: FirExplicitSimpleImportingScope + Element 7 + Scope: FirLocalScope + Properties: + FirValueParameterSymbol a: R|kotlin/Any| + Element 8 + Scope: FirLocalScope +Smart Casts: + FirValueParameterSymbol a: R|kotlin/Any| + Types: + Foo + kotlin/Any + +FILE: [ResolvedTo(IMPORTS)] onType.kt + public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] a: R|kotlin/Any|): R|kotlin/Unit| { + when () { + (R|/a| !is R|Foo|) -> { + ^test Unit + } + } + + R|/call||>() + } + public final inline [ResolvedTo(CONTRACTS)] fun call(): R|kotlin/Unit| { + } + public? final? [ResolvedTo(RAW_FIR)] interface Foo : R|kotlin/Any| { + public? final? [ResolvedTo(RAW_FIR)] fun process(): Boolean + + } diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java index fc56abe5027..970e164e84e 100644 --- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java +++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java @@ -269,6 +269,24 @@ public class ContextCollectorSourceTestGenerated extends AbstractContextCollecto runTest("analysis/low-level-api-fir/testData/contextCollector/smartCasts/insideLoop.kt"); } + @Test + @TestMetadata("onLabel.kt") + public void testOnLabel() throws Exception { + runTest("analysis/low-level-api-fir/testData/contextCollector/smartCasts/onLabel.kt"); + } + + @Test + @TestMetadata("onReference.kt") + public void testOnReference() throws Exception { + runTest("analysis/low-level-api-fir/testData/contextCollector/smartCasts/onReference.kt"); + } + + @Test + @TestMetadata("onType.kt") + public void testOnType() throws Exception { + runTest("analysis/low-level-api-fir/testData/contextCollector/smartCasts/onType.kt"); + } + @Test @TestMetadata("orRight.kt") public void testOrRight() throws Exception {