Analysis API: fix getKtExpressionType for array assignment target

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirExpressionTypeProvider.kt

@@ -39,6 +39,10 @@ internal class KtFirExpressionTypeProvider(
 
     override fun getKtExpressionType(expression: KtExpression): KtType? = withValidityAssertion {
         when (val fir = expression.unwrap().getOrBuildFir(firResolveState)) {
+            is FirFunctionCall -> {
+                getReturnTypeForArrayStyleAssignmentTarget(expression, fir)
+                    ?: fir.typeRef.coneType.asKtType()
+            }
             is FirPropertyAccessExpression -> {
                 // For unresolved `super`, we manually create an intersection type so that IDE features like completion can work correctly.
                 val containingClass =
@@ -78,6 +82,18 @@ internal class KtFirExpressionTypeProvider(
         }
     }
 
+    private fun getReturnTypeForArrayStyleAssignmentTarget(
+        expression: KtExpression,
+        fir: FirFunctionCall
+    ): KtType? {
+        if (fir.calleeReference !is FirResolvedNamedReference) return null
+        if (expression !is KtArrayAccessExpression) return null
+        val assignment = expression.parent as? KtBinaryExpression ?: return null
+        if (assignment.operationToken !in KtTokens.ALL_ASSIGNMENTS) return null
+        val setTargetArgumentParameter = fir.argumentMapping?.entries?.last()?.value ?: return null
+        return setTargetArgumentParameter.returnTypeRef.coneType.asKtType()
+    }
+
     override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType = withValidityAssertion {
         val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration>(firResolveState)
         firDeclaration.returnTypeRef.coneType.asKtType()

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/components/expressionTypeProvider/FirHLExpressionTypeTestGenerated.java

@@ -209,4 +209,32 @@ public class FirHLExpressionTypeTestGenerated extends AbstractFirHLExpressionTyp
     public void testWhileExpression() throws Exception {
         runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/whileExpression.kt");
     }
+
+    @Nested
+    @TestMetadata("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment")
+    @TestDataPath("$PROJECT_ROOT")
+    public class Assignment {
+        @Test
+        public void testAllFilesPresentInAssignment() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment"), Pattern.compile("^(.+)\\.kt$"), null, true);
+        }
+
+        @Test
+        @TestMetadata("arrayAssignementTarget.kt")
+        public void testArrayAssignementTarget() throws Exception {
+            runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/arrayAssignementTarget.kt");
+        }
+
+        @Test
+        @TestMetadata("arrayAssignmentTargetUnresovledSet.kt")
+        public void testArrayAssignmentTargetUnresovledSet() throws Exception {
+            runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/arrayAssignmentTargetUnresovledSet.kt");
+        }
+
+        @Test
+        @TestMetadata("arrayCompoundAssignementTarget.kt")
+        public void testArrayCompoundAssignementTarget() throws Exception {
+            runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/arrayCompoundAssignementTarget.kt");
+        }
+    }
 }