psi2ir: dynamic array element get/set

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/CallGenerator.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.ir.types.classifierOrFail
 import org.jetbrains.kotlin.ir.util.referenceFunction
 import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.load.java.sam.SamConstructorDescriptor
+import org.jetbrains.kotlin.psi.KtArrayAccessExpression
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -236,13 +237,19 @@ class CallGenerator(statementGenerator: StatementGenerator) : StatementGenerator
         return dispatchReceiver
     }
 
-    private fun ResolvedCall<*>.isImplicitInvokeOnDynamic(): Boolean {
+    private fun ResolvedCall<*>.isImplicitInvoke(): Boolean {
         if (resultingDescriptor.name != OperatorNameConventions.INVOKE) return false
-        val callExression = call.callElement as? KtCallExpression ?: return true
-        val calleeExpression = callExression.calleeExpression as? KtSimpleNameExpression ?: return true
+        val callExpression = call.callElement as? KtCallExpression ?: return true
+        val calleeExpression = callExpression.calleeExpression as? KtSimpleNameExpression ?: return true
         return calleeExpression.getReferencedName() != OperatorNameConventions.INVOKE.asString()
     }
 
+    private fun ResolvedCall<*>.isImplicitGet(): Boolean =
+        resultingDescriptor.name == OperatorNameConventions.GET && call.callElement is KtArrayAccessExpression
+
+    private fun ResolvedCall<*>.isImplicitSet(): Boolean =
+        resultingDescriptor.name == OperatorNameConventions.SET && call.callElement is KtArrayAccessExpression
+
     private fun generateFunctionCall(
         functionDescriptor: FunctionDescriptor,
         startOffset: Int,
@@ -255,27 +262,56 @@ class CallGenerator(statementGenerator: StatementGenerator) : StatementGenerator
             val irType = returnType.toIrType()
 
             if (functionDescriptor.isDynamic()) {
-                IrDynamicOperatorExpressionImpl(
-                    startOffset, endOffset,
-                    irType,
-                    IrDynamicOperator.INVOKE
-                ).apply {
-                    val dispatchReceiver = getDynamicExpressionReceiver(dispatchReceiverValue, extensionReceiverValue, functionDescriptor)
+                fun makeDynamicOperatorExpression(operator: IrDynamicOperator) =
+                    IrDynamicOperatorExpressionImpl(
+                        startOffset, endOffset,
+                        irType,
+                        operator
+                    )
 
-                    receiver =
-                        if (call.original.isImplicitInvokeOnDynamic())
-                            dispatchReceiver
-                        else
+                fun makeDynamicOperatorExpression(operator: IrDynamicOperator, dynamicReceiver: IrExpression) =
+                    IrDynamicOperatorExpressionImpl(
+                        startOffset, endOffset,
+                        irType,
+                        operator
+                    ).apply {
+                        receiver = dynamicReceiver
+                        arguments.addAll(call.getValueArgumentsInParameterOrder().mapIndexed { index: Int, arg: IrExpression? ->
+                            arg ?: throw AssertionError("No argument in dynamic call $functionDescriptor at position $index")
+                        })
+                    }
+
+                val dynamicReceiver = getDynamicExpressionReceiver(dispatchReceiverValue, extensionReceiverValue, functionDescriptor)
+
+                when {
+                    call.original.isImplicitInvoke() ->
+                        makeDynamicOperatorExpression(IrDynamicOperator.INVOKE, dynamicReceiver)
+                    call.original.isImplicitGet() ->
+                        makeDynamicOperatorExpression(IrDynamicOperator.ARRAY_ACCESS, dynamicReceiver)
+                    call.original.isImplicitSet() ->
+                        makeDynamicOperatorExpression(IrDynamicOperator.EQ).apply {
+                            val args = call.getValueArgumentsInParameterOrder()
+                            val arg0 = args[0]
+                                ?: throw AssertionError("No index argument in dynamic array set: ${call.original.call.callElement.text}")
+                            val arg1 = args[1]
+                                ?: throw AssertionError("No value argument in dynamic array set: ${call.original.call.callElement.text}")
+                            receiver =
+                                makeDynamicOperatorExpression(IrDynamicOperator.ARRAY_ACCESS).apply {
+                                    receiver = dynamicReceiver
+                                    arguments.add(arg0)
+                                }
+                            arguments.add(arg1)
+                        }
+                    else ->
+                        makeDynamicOperatorExpression(
+                            IrDynamicOperator.INVOKE,
                             IrDynamicMemberExpressionImpl(
                                 startOffset, endOffset, // TODO obtain more exact start/end offsets for explicit receiver expression
-                                dispatchReceiver.type,
+                                dynamicReceiver.type,
                                 functionDescriptor.name.asString(),
-                                dispatchReceiver
+                                dynamicReceiver
                             )
-
-                    arguments.addAll(call.getValueArgumentsInParameterOrder().mapIndexed { index: Int, arg: IrExpression? ->
-                        arg ?: throw AssertionError("No argument in dynamic call $functionDescriptor at position $index")
-                    })
+                        )
                 }
             } else {
                 val functionSymbol = context.symbolTable.referenceFunction(functionDescriptor.original)

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/OperatorExpressionGenerator.kt

@@ -20,11 +20,13 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.expressions.IrDynamicOperator
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.impl.originalKotlinType
 import org.jetbrains.kotlin.ir.types.makeNotNull
 import org.jetbrains.kotlin.ir.util.referenceClassifier
@@ -45,7 +47,6 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isPrimitiveNumberType
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.types.typeUtil.makeNullable
-import java.lang.AssertionError
 
 
 class OperatorExpressionGenerator(statementGenerator: StatementGenerator) : StatementGeneratorExtension(statementGenerator) {
@@ -433,4 +434,21 @@ class OperatorExpressionGenerator(statementGenerator: StatementGenerator) : Stat
 
         return generateCall(resolvedCall, expression, origin)
     }
+
+    fun generateDynamicOperatorExpression(
+        irOperator: IrDynamicOperator,
+        irType: IrType,
+        ktOperatorExpression: KtExpression,
+        ktReceiverExpression: KtExpression,
+        ktArgumentExpressions: List<KtExpression>
+    ): IrExpression =
+        IrDynamicOperatorExpressionImpl(
+            ktOperatorExpression.startOffsetSkippingComments,
+            ktOperatorExpression.endOffset,
+            irType,
+            irOperator
+        ).apply {
+            receiver = statementGenerator.generateExpression(ktReceiverExpression)
+            ktArgumentExpressions.mapTo(arguments) { statementGenerator.generateExpression(it) }
+        }
 }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/StatementGenerator.kt

@@ -41,6 +41,7 @@ import org.jetbrains.kotlin.resolve.BindingContextUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
+import org.jetbrains.kotlin.resolve.calls.tasks.isDynamic
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.types.KotlinType
@@ -353,10 +354,19 @@ class StatementGenerator(
     override fun visitArrayAccessExpression(expression: KtArrayAccessExpression, data: Nothing?): IrStatement {
         val indexedGetCall = getOrFail(BindingContext.INDEXED_LVALUE_GET, expression)
 
-        return CallGenerator(this).generateCall(
-            expression.startOffsetSkippingComments, expression.endOffset,
-            pregenerateCall(indexedGetCall), IrStatementOrigin.GET_ARRAY_ELEMENT
-        )
+        return if (indexedGetCall.resultingDescriptor.isDynamic())
+            OperatorExpressionGenerator(this).generateDynamicOperatorExpression(
+                IrDynamicOperator.ARRAY_ACCESS,
+                indexedGetCall.resultingDescriptor.returnType!!.toIrType(),
+                expression,
+                expression.arrayExpression ?: throw AssertionError("Array expression not found"),
+                expression.indexExpressions
+            )
+        else
+            CallGenerator(this).generateCall(
+                expression.startOffsetSkippingComments, expression.endOffset,
+                pregenerateCall(indexedGetCall), IrStatementOrigin.GET_ARRAY_ELEMENT
+            )
     }
 
     override fun visitDotQualifiedExpression(expression: KtDotQualifiedExpression, data: Nothing?): IrStatement =