[AA] Fix getKtExpressionType for augmented array assignment

This previously worked on accident because the get call in an
augmented array assignment wouldn't have a resolved argument list, and
so `argumentsToSubstitutedValueParameters` would return null.
Now, we additionally verify that we're in a `set` call.

#KT-66124

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

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getOutermostParenthesizerOrThis
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 import org.jetbrains.kotlin.utils.exceptions.rethrowExceptionWithDetails
 import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
@@ -150,7 +151,8 @@ internal class KtFirExpressionTypeProvider(
         expression: KtExpression,
         fir: FirFunctionCall,
     ): KtType? {
-        if (fir.calleeReference !is FirResolvedNamedReference) return null
+        // When we're in a call like `a[x] = y`, we want to get the `set` call's last argument's type.
+        if (fir.calleeReference !is FirResolvedNamedReference || fir.calleeReference.name != OperatorNameConventions.SET) 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

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/expressionTypeProvider/FirIdeDependentAnalysisSourceModuleHLExpressionTypeTestGenerated.java

@@ -313,6 +313,12 @@ public class FirIdeDependentAnalysisSourceModuleHLExpressionTypeTestGenerated ex
       runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/arrayCompoundAssignementTarget.kt");
     }
 
+    @Test
+    @TestMetadata("augmentedArrayAssigment.kt")
+    public void testAugmentedArrayAssigment() {
+      runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/augmentedArrayAssigment.kt");
+    }
+
     @Test
     @TestMetadata("readArrayElement.kt")
     public void testReadArrayElement() {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/expressionTypeProvider/FirIdeNormalAnalysisSourceModuleHLExpressionTypeTestGenerated.java

@@ -313,6 +313,12 @@ public class FirIdeNormalAnalysisSourceModuleHLExpressionTypeTestGenerated exten
       runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/arrayCompoundAssignementTarget.kt");
     }
 
+    @Test
+    @TestMetadata("augmentedArrayAssigment.kt")
+    public void testAugmentedArrayAssigment() {
+      runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expressionType/assignment/augmentedArrayAssigment.kt");
+    }
+
     @Test
     @TestMetadata("readArrayElement.kt")
     public void testReadArrayElement() {