FIR: Ensure that the array set argument (on RHS of =) is always mapped

to the last parameter of the set operator function, even if there are
missing or too many index arguments.

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -4437,6 +4437,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSet.kt");
             }
 
+            @Test
+            @TestMetadata("arrayAccessSetNotEnoughArgs.kt")
+            public void testArrayAccessSetNotEnoughArgs() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetNotEnoughArgs.kt");
+            }
+
             @Test
             @TestMetadata("arrayAccessSetTooManyArgs.kt")
             public void testArrayAccessSetTooManyArgs() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -4437,6 +4437,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSet.kt");
             }
 
+            @Test
+            @TestMetadata("arrayAccessSetNotEnoughArgs.kt")
+            public void testArrayAccessSetNotEnoughArgs() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetNotEnoughArgs.kt");
+            }
+
             @Test
             @TestMetadata("arrayAccessSetTooManyArgs.kt")
             public void testArrayAccessSetTooManyArgs() throws Exception {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt

@@ -437,6 +437,16 @@ object LightTreePositioningStrategies {
             endOffset: Int,
             tree: FlyweightCapableTreeStructure<LighterASTNode>
         ): List<TextRange> {
+            if (node.tokenType == KtNodeTypes.BINARY_EXPRESSION &&
+                tree.findDescendantByTypes(node, KtTokens.ALL_ASSIGNMENTS) != null
+            ) {
+                val lhs = tree.firstChildExpression(node)
+                lhs?.let {
+                    tree.unwrapParenthesesLabelsAndAnnotations(it)?.let { unwrapped ->
+                        return markElement(unwrapped, startOffset, endOffset, tree, node)
+                    }
+                }
+            }
             val nodeToStart = when (node.tokenType) {
                 in KtTokens.QUALIFIED_ACCESS -> tree.findLastChildByType(node, KtNodeTypes.CALL_EXPRESSION) ?: node
                 else -> node

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -2182,9 +2182,9 @@ open class RawFirBuilder(
 
         override fun visitArrayAccessExpression(expression: KtArrayAccessExpression, data: Unit): FirElement {
             val arrayExpression = expression.arrayExpression
-            val getArgument = context.arraySetArgument.remove(expression)
+            val setArgument = context.arraySetArgument.remove(expression)
             return buildFunctionCall {
-                val isGet = getArgument == null
+                val isGet = setArgument == null
                 source = (if (isGet) expression else expression.parent).toFirSourceElement()
                 calleeReference = buildSimpleNamedReference {
                     source = expression.toFirSourceElement().fakeElement(FirFakeSourceElementKind.ArrayAccessNameReference)
@@ -2195,8 +2195,8 @@ open class RawFirBuilder(
                     for (indexExpression in expression.indexExpressions) {
                         arguments += indexExpression.toFirExpression("Incorrect index expression")
                     }
-                    if (getArgument != null) {
+                    if (setArgument != null) {
-                        arguments += getArgument
+                        arguments += setArgument
                     }
                 }
                 origin = FirFunctionCallOrigin.Operator

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirArgumentsToParametersMapper.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.resolve.getAsForbiddenNamedArgumentsTarget
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import kotlin.collections.component1
 import kotlin.collections.component2
 import kotlin.collections.set
@@ -41,7 +42,8 @@ data class ArgumentMapping(
                 is ResolvedCallArgument.VarargArgument -> resolvedArgument.arguments.forEach {
                     argumentToParameterMapping[it] = valueParameter
                 }
-                ResolvedCallArgument.DefaultArgument -> {}
+                ResolvedCallArgument.DefaultArgument -> {
+                }
             }
         }
         return argumentToParameterMapping
@@ -63,7 +65,7 @@ fun BodyResolveComponents.mapArguments(
         return EmptyArgumentMapping
     }
 
-    val argumentsInParenthesis: MutableList<FirExpression> = mutableListOf()
+    val nonLambdaArguments: MutableList<FirExpression> = mutableListOf()
     val excessLambdaArguments: MutableList<FirExpression> = mutableListOf()
     var externalArgument: FirExpression? = null
     for (argument in arguments) {
@@ -74,17 +76,17 @@ fun BodyResolveComponents.mapArguments(
                 excessLambdaArguments.add(argument)
             }
         } else {
-            argumentsInParenthesis.add(argument)
+            nonLambdaArguments.add(argument)
         }
     }
 
-    // If this is an overloading indexed access operator, it could have default values or a vararg parameter in the middle.
+    // If this is an indexed access set operator, it could have default values or a vararg parameter in the middle.
     // For proper argument mapping, wrap the last one, which is supposed to be the updated value, as a named argument.
-    if ((function as? FirSimpleFunction)?.isOperator == true &&
+    val isIndexedSetOperator = function is FirSimpleFunction && function.isOperator && function.name == OperatorNameConventions.SET
-        function.name == Name.identifier("set") &&
+    if (isIndexedSetOperator &&
         function.valueParameters.any { it.defaultValue != null || it.isVararg }
     ) {
-        val v = argumentsInParenthesis.last()
+        val v = nonLambdaArguments.last()
         if (v !is FirNamedArgumentExpression) {
             val namedV = buildNamedArgumentExpression {
                 source = v.source
@@ -92,13 +94,13 @@ fun BodyResolveComponents.mapArguments(
                 isSpread = false
                 name = function.valueParameters.last().name
             }
-            argumentsInParenthesis.removeAt(argumentsInParenthesis.size - 1)
+            nonLambdaArguments.removeAt(nonLambdaArguments.size - 1)
-            argumentsInParenthesis.add(namedV)
+            nonLambdaArguments.add(namedV)
         }
     }
 
-    val processor = FirCallArgumentsProcessor(session, function, this, originScope)
+    val processor = FirCallArgumentsProcessor(session, function, this, originScope, isIndexedSetOperator)
-    processor.processArgumentsInParenthesis(argumentsInParenthesis)
+    processor.processNonLambdaArguments(nonLambdaArguments)
     if (externalArgument != null) {
         processor.processExternalArgument(externalArgument)
     }
@@ -113,6 +115,7 @@ private class FirCallArgumentsProcessor(
     private val function: FirFunction,
     private val bodyResolveComponents: BodyResolveComponents,
     private val originScope: FirScope?,
+    private val isIndexedSetOperator: Boolean
 ) {
     private var state = State.POSITION_ARGUMENTS
     private var currentPositionedParameterIndex = 0
@@ -132,11 +135,11 @@ private class FirCallArgumentsProcessor(
         NAMED_ONLY_ARGUMENTS
     }
 
-    fun processArgumentsInParenthesis(arguments: List<FirExpression>) {
+    fun processNonLambdaArguments(arguments: List<FirExpression>) {
         for (argument in arguments) {
             // process position argument
             if (argument !is FirNamedArgumentExpression) {
-                if (processPositionArgument(argument)) {
+                if (processPositionArgument(argument, isLastArgument = argument === arguments.last())) {
                     state = State.VARARG_POSITION
                 }
             }
@@ -155,13 +158,31 @@ private class FirCallArgumentsProcessor(
     }
 
     // return true, if it was mapped to vararg parameter
-    private fun processPositionArgument(argument: FirExpression): Boolean {
+    private fun processPositionArgument(argument: FirExpression, isLastArgument: Boolean): Boolean {
         if (state == State.NAMED_ONLY_ARGUMENTS) {
             addDiagnostic(MixingNamedAndPositionArguments(argument))
             return false
         }
 
-        val parameter = parameters.getOrNull(currentPositionedParameterIndex)
+        // The last parameter of an indexed set operator should be reserved for the last argument (the assigned value).
+        // We don't want the assigned value mapped to an index parameter if some of the index arguments are absent.
+        val assignedParameterIndex = if (isIndexedSetOperator) {
+            val lastParameterIndex = parameters.lastIndex
+            when {
+                isLastArgument -> lastParameterIndex
+                currentPositionedParameterIndex >= lastParameterIndex -> {
+                    // This is an extra index argument that should NOT be mapped to the parameter for the assigned value.
+                    -1
+                }
+                else -> {
+                    // This is an index argument that can be properly mapped.
+                    currentPositionedParameterIndex
+                }
+            }
+        } else {
+            currentPositionedParameterIndex
+        }
+        val parameter = parameters.getOrNull(assignedParameterIndex)
         if (parameter == null) {
             addDiagnostic(TooManyArguments(argument, function))
             return false

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/PositioningStrategies.kt

@@ -642,6 +642,9 @@ object PositioningStrategies {
     @JvmField
     val VALUE_ARGUMENTS: PositioningStrategy<KtElement> = object : PositioningStrategy<KtElement>() {
         override fun mark(element: KtElement): List<TextRange> {
+            if (element is KtBinaryExpression && element.operationToken in KtTokens.ALL_ASSIGNMENTS) {
+                element.left.let { left -> left.unwrapParenthesesLabelsAndAnnotations()?.let { return markElement(it) } }
+            }
             val qualifiedAccess = when (element) {
                 is KtQualifiedExpression -> element.selectorExpression ?: element
                 else -> element

compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetNotEnoughArgs.kt

@@ -0,0 +1,10 @@
+// FIR_IDENTICAL
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+class A {
+    operator fun set(x: String, y: Boolean, value: Int) {}
+
+    fun d(x: Int) {
+        <!NO_VALUE_FOR_PARAMETER("y")!>this[""]<!> = 1
+    }
+}

compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetNotEnoughArgs.txt

@@ -0,0 +1,10 @@
+package
+
+public final class A {
+    public constructor A()
+    public final fun d(/*0*/ x: kotlin.Int): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final operator fun set(/*0*/ x: kotlin.String, /*1*/ y: kotlin.Boolean, /*2*/ value: kotlin.Int): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetTooManyArgs.fir.kt

@@ -0,0 +1,10 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+class A {
+    operator fun get(x: Int) {}
+    operator fun set(x: String, value: Int) {}
+
+    fun d(x: Int) {
+        this["", <!TOO_MANY_ARGUMENTS!>1<!>] = 1
+    }
+}

compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetTooManyArgs.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 
 class A {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -4443,6 +4443,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSet.kt");
             }
 
+            @Test
+            @TestMetadata("arrayAccessSetNotEnoughArgs.kt")
+            public void testArrayAccessSetNotEnoughArgs() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetNotEnoughArgs.kt");
+            }
+
             @Test
             @TestMetadata("arrayAccessSetTooManyArgs.kt")
             public void testArrayAccessSetTooManyArgs() throws Exception {

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -4437,6 +4437,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSet.kt");
             }
 
+            @Test
+            @TestMetadata("arrayAccessSetNotEnoughArgs.kt")
+            public void testArrayAccessSetNotEnoughArgs() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/checkArguments/arrayAccessSetNotEnoughArgs.kt");
+            }
+
             @Test
             @TestMetadata("arrayAccessSetTooManyArgs.kt")
             public void testArrayAccessSetTooManyArgs() throws Exception {

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGet.kt

@@ -0,0 +1,7 @@
+class C {
+    operator fun get(a: Int, b: String): Boolean = true
+}
+
+fun call(c: C) {
+    val res = <expr>c[1, "foo"]</expr>
+}

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGet.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (a: kotlin.Int), "foo" -> (b: kotlin.String) }
+targetFunction = /C.get(a: kotlin.Int, b: kotlin.String): kotlin.Boolean

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGetWithNotEnoughArgs.kt

@@ -0,0 +1,7 @@
+class C {
+    operator fun get(a: Int, b: String): Boolean = true
+}
+
+fun call(c: C) {
+    val res = <expr>c[1]</expr>
+}

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGetWithNotEnoughArgs.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (a: kotlin.Int) }
+targetFunction = ERR<No value passed for parameter 'b', [/C.get(a: kotlin.Int, b: kotlin.String): kotlin.Boolean]>

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGetWithTooManyArgs.kt

@@ -0,0 +1,7 @@
+class C {
+    operator fun get(a: Int, b: String): Boolean = true
+}
+
+fun call(c: C) {
+    val res = <expr>c[1, "foo", false]</expr>
+}

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGetWithTooManyArgs.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (a: kotlin.Int), "foo" -> (b: kotlin.String) }
+targetFunction = ERR<Too many arguments for public final operator fun /C.get(a: R|kotlin/Int|, b: R|kotlin/String|): R|kotlin/Boolean|, [/C.get(a: kotlin.Int, b: kotlin.String): kotlin.Boolean]>

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSet.kt

@@ -0,0 +1,7 @@
+class C {
+    operator fun set(a: Int, b: String, value: Boolean) {}
+}
+
+fun call(c: C) {
+    <expr>c[1, "foo"]</expr> = false
+}

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSet.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (a: kotlin.Int), "foo" -> (b: kotlin.String), false -> (value: kotlin.Boolean) }
+targetFunction = /C.set(a: kotlin.Int, b: kotlin.String, value: kotlin.Boolean): kotlin.Unit

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSetWithNotEnoughArgs.kt

@@ -0,0 +1,7 @@
+class C {
+    operator fun set(a: Int, b: String, value: Boolean) {}
+}
+
+fun call(c: C) {
+    <expr>c[1]</expr> = false
+}

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSetWithNotEnoughArgs.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (a: kotlin.Int), false -> (value: kotlin.Boolean) }
+targetFunction = ERR<No value passed for parameter 'b', [/C.set(a: kotlin.Int, b: kotlin.String, value: kotlin.Boolean): kotlin.Unit]>

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSetWithTooManyArgs.kt

@@ -0,0 +1,7 @@
+class C {
+    operator fun set(a: Int, b: String, value: Boolean) {}
+}
+
+fun call(c: C) {
+    <expr>c[1, "foo", 3.14]</expr> = false
+}

idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSetWithTooManyArgs.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (a: kotlin.Int), "foo" -> (b: kotlin.String), false -> (value: kotlin.Boolean) }
+targetFunction = ERR<Too many arguments for public final operator fun /C.set(a: R|kotlin/Int|, b: R|kotlin/String|, value: R|kotlin/Boolean|): R|kotlin/Unit|, [/C.set(a: kotlin.Int, b: kotlin.String, value: kotlin.Boolean): kotlin.Unit]>

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/fir/AbstractResolveCallTest.kt

@@ -45,8 +45,9 @@ abstract class AbstractResolveCallTest : AbstractHLApiSingleModuleTest() {
         is KtCallElement -> element.resolveCall()
         is KtBinaryExpression -> element.resolveCall()
         is KtUnaryExpression -> element.resolveCall()
+        is KtArrayAccessExpression -> element.resolveCall()
         is KtValueArgument -> resolveCall(element.getArgumentExpression()!!)
-        else -> error("Selected should be either KtCallElement, KtBinaryExpression, or KtUnaryExpression, but was $element")
+        else -> error("Selected element type (${element::class.simpleName}) is not supported for resolveCall()")
     }
 
 }

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/fir/ResolveCallTestGenerated.java

@@ -114,6 +114,42 @@ public class ResolveCallTestGenerated extends AbstractResolveCallTest {
         runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/implicitJavaConstuctorCall.kt");
     }
 
+    @Test
+    @TestMetadata("indexedGet.kt")
+    public void testIndexedGet() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGet.kt");
+    }
+
+    @Test
+    @TestMetadata("indexedGetWithNotEnoughArgs.kt")
+    public void testIndexedGetWithNotEnoughArgs() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGetWithNotEnoughArgs.kt");
+    }
+
+    @Test
+    @TestMetadata("indexedGetWithTooManyArgs.kt")
+    public void testIndexedGetWithTooManyArgs() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedGetWithTooManyArgs.kt");
+    }
+
+    @Test
+    @TestMetadata("indexedSet.kt")
+    public void testIndexedSet() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSet.kt");
+    }
+
+    @Test
+    @TestMetadata("indexedSetWithNotEnoughArgs.kt")
+    public void testIndexedSetWithNotEnoughArgs() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSetWithNotEnoughArgs.kt");
+    }
+
+    @Test
+    @TestMetadata("indexedSetWithTooManyArgs.kt")
+    public void testIndexedSetWithTooManyArgs() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/analysisSession/resolveCall/indexedSetWithTooManyArgs.kt");
+    }
+
     @Test
     @TestMetadata("javaFunctionCall.kt")
     public void testJavaFunctionCall() throws Exception {