[LL API] Fix smart-cast collection for special 'FirElement's

'DFANode' is not created for certain elements, such as types or
references. Before the change, only the target element itself was
queried for DFA, making the 'ContextCollector' return an empty map of
smart casts.

^KTIJ-26973 Fixed
This commit is contained in:
Yan Zhulanow
2023-09-09 02:04:48 +09:00
committed by Space Team
parent 6c1711c52a
commit 4690a430f4
9 changed files with 251 additions and 15 deletions
@@ -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<FirElement>()
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<Pair<Int, ConeKotlinType>>()
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()
@@ -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?| {
@@ -0,0 +1,17 @@
fun test(a: Any) {
if (a !is Foo) {
return
}
var result: Int = 0
<expr>loop@</expr> while (true) {
if (!a.process()) {
break@loop
}
}
}
interface Foo {
fun process(): Boolean
}
@@ -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|<local>/a| !is R|Foo|) -> {
^test Unit
}
}
[ResolvedTo(BODY_RESOLVE)] lvar result: R|kotlin/Int| = Int(0)
loop@while(Boolean(true)) {
when () {
R|<local>/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|
}
@@ -0,0 +1,13 @@
fun test(a: Any) {
if (a !is Foo) {
return
}
<expr>check</expr>(a.process())
}
fun check(condition: Boolean) {}
interface Foo {
fun process(): Boolean
}
@@ -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|<local>/a| !is R|Foo|) -> {
^test Unit
}
}
R|/check|(R|<local>/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|
}
@@ -0,0 +1,13 @@
fun test(a: Any) {
if (a !is Foo) {
return
}
call<List<<expr>Int</expr>>>()
}
inline fun <reified T : Any> call() {}
interface Foo {
fun process(): Boolean
}
@@ -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|<local>/a| !is R|Foo|) -> {
^test Unit
}
}
R|/call|<R|kotlin/collections/List<kotlin/Int>|>()
}
public final inline [ResolvedTo(CONTRACTS)] fun <reified [ResolvedTo(CONTRACTS)] T : R|kotlin/Any|> call(): R|kotlin/Unit| {
}
public? final? [ResolvedTo(RAW_FIR)] interface Foo : R|kotlin/Any| {
public? final? [ResolvedTo(RAW_FIR)] fun process(): Boolean
}
@@ -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 {