FIR LL API: resolve PSI elements to more specific FIR element

Compound access and implicit delegated constructor are mapped to FIR elements that are way too broad.
2021-11-22 21:03:21 -08:00
parent e8f1af6140
commit bc95733818
27 changed files with 252 additions and 1 deletions
@@ -120,6 +120,24 @@ public class KtFe10HLExpressionTypeTestGenerated extends AbstractKtFe10HLExpress
        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/nonExpression.kt");
    }

+    @Test
+    @TestMetadata("plusAssign.kt")
+    public void testPlusAssign() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/plusAssign.kt");
+    }
+
+    @Test
+    @TestMetadata("postfixDec.kt")
+    public void testPostfixDec() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/postfixDec.kt");
+    }
+
+    @Test
+    @TestMetadata("prefixInc.kt")
+    public void testPrefixInc() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/prefixInc.kt");
+    }
+
    @Test
    @TestMetadata("property.kt")
    public void testProperty() throws Exception {
@@ -54,6 +54,13 @@ internal class KtFirExpressionTypeProvider(
                    fir.typeRef.coneType.asKtType()
                }
            }
+            is FirVariableAssignment -> {
+                if (expression is KtUnaryExpression && expression.operationToken in KtTokens.INCREMENT_AND_DECREMENT) {
+                    fir.rValue.typeRef.coneType.asKtType()
+                } else {
+                    analysisSession.builtinTypes.UNIT
+                }
+            }
            is FirExpression -> fir.typeRef.coneType.asKtType()
            is FirNamedReference -> fir.getReferencedElementType(firResolveState).asKtType()
            is FirStatement -> with(analysisSession) { builtinTypes.UNIT }
@@ -120,6 +120,24 @@ public class FirHLExpressionTypeTestGenerated extends AbstractFirHLExpressionTyp
        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/nonExpression.kt");
    }

+    @Test
+    @TestMetadata("plusAssign.kt")
+    public void testPlusAssign() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/plusAssign.kt");
+    }
+
+    @Test
+    @TestMetadata("postfixDec.kt")
+    public void testPostfixDec() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/postfixDec.kt");
+    }
+
+    @Test
+    @TestMetadata("prefixInc.kt")
+    public void testPrefixInc() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/prefixInc.kt");
+    }
+
    @Test
    @TestMetadata("property.kt")
    public void testProperty() throws Exception {
@@ -0,0 +1,4 @@
+fun test() {
+    var i = 1
+    <expr>i += 1</expr>
+}
@@ -0,0 +1,2 @@
+expression: i += 1
+type: kotlin.Unit
@@ -0,0 +1,4 @@
+fun test() {
+    var i = 1
+    <expr>i--</expr>
+}
@@ -0,0 +1,2 @@
+expression: i--
+type: kotlin.Int
@@ -0,0 +1,4 @@
+fun test() {
+    var i = 1
+    <expr>++i</expr>
+}
@@ -0,0 +1,2 @@
+expression: ++i
+type: kotlin.Int
@@ -7,9 +7,12 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure

 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.KtRealPsiSourceElement
+import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.DuplicatedFirSourceElementsException
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isErrorElement
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
@@ -19,8 +22,11 @@ import org.jetbrains.kotlin.fir.types.FirUserTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtTypeReference
+import org.jetbrains.kotlin.psi.KtUnaryExpression
+import org.jetbrains.kotlin.util.OperatorNameConventions

 internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement, FirElement>>() {
    private fun cache(psi: KtElement, fir: FirElement, cache: MutableMap<KtElement, FirElement>) {
@@ -95,12 +101,50 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
            ?.takeIf {
                it is KtRealPsiSourceElement ||
                        it.kind == KtFakeSourceElementKind.ReferenceInAtomicQualifiedAccess ||
-                        it.kind == KtFakeSourceElementKind.FromUseSiteTarget
+                        it.kind == KtFakeSourceElementKind.FromUseSiteTarget ||
+                        // For secondary constructors without explicit delegated constructor call, the PSI tree always create an empty
+                        // KtConstructorDelegationCall. In this case, the source in FIR has this fake source kind.
+                        it.kind == KtFakeSourceElementKind.ImplicitConstructor ||
+                        it.isSourceForCompoundAccess(element)
            }.psi as? KtElement
            ?: return
        cache(psi, element, cache)
    }

+    /**
+     * FIR represents compound assignment and inc/dec operations as multiple smaller instructions. Here we choose the write operation as the
+     * resolved FirElement for binary and unary expressions. For example, the `FirVariableAssignment` or the call to `set` or `plusAssign`
+     * function, etc. This is because the write FirElement can be used to retrieve all other information related to this compound operation.
+
+     * On the other hand, if the PSI is the left operand of an assignment or the base expression of a unary expression, we take the read FIR
+     * element so the user of the Analysis API is able to retrieve such read calls reliably.
+     */
+    private fun KtSourceElement.isSourceForCompoundAccess(fir: FirElement): Boolean {
+        val psi = psi
+        val parentPsi = psi?.parent
+        if (kind != KtFakeSourceElementKind.DesugaredCompoundAssignment && kind != KtFakeSourceElementKind.DesugaredIncrementOrDecrement) return false
+        return when {
+            psi is KtBinaryExpression || psi is KtUnaryExpression -> fir.isWriteInCompoundCall()
+            parentPsi is KtBinaryExpression && psi == parentPsi.left -> fir.isReadInCompoundCall()
+            parentPsi is KtUnaryExpression && psi == parentPsi.baseExpression -> fir.isReadInCompoundCall()
+            else -> false
+        }
+    }
+
+    private fun FirElement.isReadInCompoundCall(): Boolean {
+        if (this is FirPropertyAccessExpression) return true
+        if (this !is FirFunctionCall) return false
+        val name = (calleeReference as? FirResolvedNamedReference)?.name
+        return name == OperatorNameConventions.GET
+    }
+
+    private fun FirElement.isWriteInCompoundCall(): Boolean {
+        if (this is FirVariableAssignment) return true
+        if (this !is FirFunctionCall) return false
+        val name = (calleeReference as? FirResolvedNamedReference)?.name
+        return name == OperatorNameConventions.SET || name in OperatorNameConventions.ASSIGNMENT_OPERATIONS
+    }
+
    companion object {
        @OptIn(ExperimentalStdlibApi::class)
        fun recordElementsFrom(firElement: FirElement, recorder: FirElementsRecorder): Map<KtElement, FirElement> =
@@ -0,0 +1,6 @@
+interface A {
+    operator fun plusAssign(i: Int)
+}
+fun test(l: A) {
+    <expr>l += 1</expr>
+}
@@ -0,0 +1,5 @@
+KT element: KtBinaryExpression
+FIR element: FirFunctionCallImpl
+
+FIR element rendered:
+R|<local>/l|.R|/A.plusAssign|(Int(1))
@@ -0,0 +1,6 @@
+interface A {
+    operator fun plusAssign(i: Int)
+}
+fun test(l: A) {
+    <expr>l</expr> += 1
+}
@@ -0,0 +1,5 @@
+KT element: KtNameReferenceExpression
+FIR element: FirPropertyAccessExpressionImpl
+
+FIR element rendered:
+R|<local>/l|
@@ -0,0 +1,4 @@
+fun test() {
+    var i = 1
+    <expr>i += 1</expr>
+}
@@ -0,0 +1,5 @@
+KT element: KtBinaryExpression
+FIR element: FirVariableAssignmentImpl
+
+FIR element rendered:
+R|<local>/i| = R|<local>/i|.R|kotlin/Int.plus|(Int(1))
@@ -0,0 +1,4 @@
+fun test() {
+    var i = 1
+    <expr>i</expr> += 1
+}
@@ -0,0 +1,5 @@
+KT element: KtNameReferenceExpression
+FIR element: FirResolvedNamedReferenceImpl
+
+FIR element rendered:
+R|<local>/i|
@@ -0,0 +1,8 @@
+interface MyMap<K, V> {
+    operator fun get(k: K): V
+    operator fun set(k: K, v: V)
+}
+
+fun test(m: MyMap<String, Int>) {
+    <expr>m["a"] += 1</expr>
+}
@@ -0,0 +1,5 @@
+KT element: KtBinaryExpression
+FIR element: FirFunctionCallImpl
+
+FIR element rendered:
+R|<local>/<<array>>|.R|SubstitutionOverride</MyMap.set: R|kotlin/Unit|>|(R|<local>/<<index_0>>|, R|<local>/<<array>>|.R|SubstitutionOverride</MyMap.get: R|kotlin/Int|>|(R|<local>/<<index_0>>|).R|kotlin/Int.plus|(Int(1)))
@@ -0,0 +1,8 @@
+interface MyMap<K, V> {
+    operator fun get(k: K): V
+    operator fun set(k: K, v: V)
+}
+
+fun test(m: MyMap<String, Int>) {
+    <expr>m["a"]</expr> += 1
+}
@@ -0,0 +1,5 @@
+KT element: KtArrayAccessExpression
+FIR element: FirFunctionCallImpl
+
+FIR element rendered:
+R|<local>/<<array>>|.R|SubstitutionOverride</MyMap.get: R|kotlin/Int|>|(R|<local>/<<index_0>>|)
@@ -0,0 +1,11 @@
+interface A {
+    operator fun plusAssign(i: Int)
+}
+
+interface MyMap<K, V> {
+    operator fun get(k: K): V
+}
+
+fun test(m: MyMap<String, A>) {
+    <expr>m["a"] += 1</expr>
+}
@@ -0,0 +1,5 @@
+KT element: KtBinaryExpression
+FIR element: FirFunctionCallImpl
+
+FIR element rendered:
+R|<local>/m|.R|SubstitutionOverride</MyMap.get: R|A|>|(String(a)).R|/A.plusAssign|(Int(1))
@@ -0,0 +1,11 @@
+interface A {
+    operator fun plusAssign(i: Int)
+}
+
+interface MyMap<K, V> {
+    operator fun get(k: K): V
+}
+
+fun test(m: MyMap<String, A>) {
+    <expr>m["a"]</expr> += 1
+}
@@ -0,0 +1,5 @@
+KT element: KtArrayAccessExpression
+FIR element: FirFunctionCallImpl
+
+FIR element rendered:
+R|<local>/m|.R|SubstitutionOverride</MyMap.get: R|A|>|(String(a))
@@ -97,6 +97,54 @@ public class GetOrBuildFirTestGenerated extends AbstractGetOrBuildFirTest {
            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/calllTypeArguments.kt");
        }

+        @Test
+        @TestMetadata("compoundAssignOnVal.kt")
+        public void testCompoundAssignOnVal() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignOnVal.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignOnVal_lhs.kt")
+        public void testCompoundAssignOnVal_lhs() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignOnVal_lhs.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignOnVar.kt")
+        public void testCompoundAssignOnVar() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignOnVar.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignOnVar_lhs.kt")
+        public void testCompoundAssignOnVar_lhs() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignOnVar_lhs.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignWithArrayAccessConvention.kt")
+        public void testCompoundAssignWithArrayAccessConvention() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayAccessConvention.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignWithArrayAccessConvention_lhs.kt")
+        public void testCompoundAssignWithArrayAccessConvention_lhs() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayAccessConvention_lhs.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignWithArrayGetConvention.kt")
+        public void testCompoundAssignWithArrayGetConvention() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayGetConvention.kt");
+        }
+
+        @Test
+        @TestMetadata("compoundAssignWithArrayGetConvention_lhs.kt")
+        public void testCompoundAssignWithArrayGetConvention_lhs() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayGetConvention_lhs.kt");
+        }
+
        @Test
        @TestMetadata("constructorDelegationSuperCall.kt")
        public void testConstructorDelegationSuperCall() throws Exception {