PSI2IR: SAM conversion should be performed once for index variables

Given esoteric code as in 'caoWithAdaptationForSam.kt', we should make sure that we pass same objects to 'get' and 'set'.
2020-01-28 17:12:50 +03:00
parent e750528551
commit bf9673a0a2
14 changed files with 689 additions and 125 deletions
@@ -17,11 +17,18 @@
 package org.jetbrains.kotlin.psi2ir.intermediate

 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.inlineStatement
-import org.jetbrains.kotlin.ir.expressions.isAssignmentOperatorWithResult
+import org.jetbrains.kotlin.ir.types.impl.originalKotlinType
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi2ir.generators.CallGenerator
 import org.jetbrains.kotlin.psi2ir.generators.generateSamConversionForValueArgumentsIfRequired
@@ -29,6 +36,7 @@ import org.jetbrains.kotlin.psi2ir.generators.pregenerateValueArgumentsUsing
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.tasks.isDynamic
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.checker.KotlinTypeChecker

 class ArrayAccessAssignmentReceiver(
    private val irArray: IrExpression,
@@ -47,6 +55,10 @@ class ArrayAccessAssignmentReceiver(
    private val indexedGetDescriptor = indexedGetResolvedCall?.resultingDescriptor
    private val indexedSetDescriptor = indexedSetResolvedCall?.resultingDescriptor

+    private class CompoundAssignmentInfo {
+        val indexVariables = LinkedHashSet<IrVariable>()
+    }
+
    private val descriptor =
        indexedGetDescriptor
            ?: indexedSetDescriptor
@@ -74,17 +86,132 @@ class ArrayAccessAssignmentReceiver(

        val irArrayValue = callGenerator.scope.createTemporaryVariableInBlock(callGenerator.context, irArray, irBlock, "array")

+        val compoundAssignmentInfo = CompoundAssignmentInfo()
+
        irBlock.inlineStatement(
            withLValue(
                createLValue(kotlinType, irArrayValue) { i, irIndex ->
-                    callGenerator.scope.createTemporaryVariableInBlock(callGenerator.context, irIndex, irBlock, "index$i")
+                    val irIndexVar = callGenerator.scope.createTemporaryVariable(irIndex, "index$i")
+                    compoundAssignmentInfo.indexVariables.add(irIndexVar)
+                    irBlock.statements.add(irIndexVar)
+                    VariableLValue(callGenerator.context, irIndexVar)
                }
            )
        )

+        postprocessSamConversionsInCompoundAssignment(irBlock, compoundAssignmentInfo)
        return irBlock
    }

+    private fun postprocessSamConversionsInCompoundAssignment(
+        irBlock: IrBlock,
+        compoundAssignmentInfo: CompoundAssignmentInfo
+    ) {
+        val samConversionsCollector = SamConversionsCollector(compoundAssignmentInfo)
+        irBlock.acceptChildrenVoid(samConversionsCollector)
+
+        if (samConversionsCollector.samConversionsPerVariable.isEmpty()) return
+
+        val samConvertedVars = hashMapOf<IrVariable, IrVariable>()
+        for ((irIndexVar, samConversions) in samConversionsCollector.samConversionsPerVariable) {
+            var mostSpecificSamConversion: IrTypeOperatorCall = samConversions.first()
+            for (samConversion in samConversions) {
+                if (samConversion === mostSpecificSamConversion) continue
+                val lastType = mostSpecificSamConversion.operandKotlinType
+                val nextType = samConversion.operandKotlinType
+                if (KotlinTypeChecker.DEFAULT.isSubtypeOf(nextType, lastType)) {
+                    mostSpecificSamConversion = samConversion
+                } else if (!KotlinTypeChecker.DEFAULT.isSubtypeOf(lastType, nextType)) {
+                    throw AssertionError("Unrelated types in SAM conversion for index variable: $lastType, $nextType")
+                }
+            }
+            val irSamConvertedVarInitializer = createSamConvertedVarInitializer(irIndexVar, mostSpecificSamConversion)
+            val irSamConvertedVar = callGenerator.scope.createTemporaryVariable(irSamConvertedVarInitializer, "sam")
+            val index = irBlock.statements.indexOf(irIndexVar)
+            irBlock.statements[index] = irSamConvertedVar
+            samConvertedVars[irIndexVar] = irSamConvertedVar
+        }
+
+        irBlock.transformChildrenVoid(SamConversionsRewriter(samConvertedVars))
+    }
+
+    private fun createSamConvertedVarInitializer(irIndexVar: IrVariable, mostSpecificSamConversion: IrTypeOperatorCall): IrExpression {
+        val irIndexVarInitializer = irIndexVar.initializer!!
+        val startOffset = irIndexVarInitializer.startOffset
+        val endOffset = irIndexVarInitializer.endOffset
+
+        val implicitCast = mostSpecificSamConversion.argument as IrTypeOperatorCall
+
+        return IrTypeOperatorCallImpl(
+            startOffset, endOffset,
+            mostSpecificSamConversion.type,
+            IrTypeOperator.SAM_CONVERSION,
+            mostSpecificSamConversion.typeOperand,
+            IrTypeOperatorCallImpl(
+                startOffset, endOffset,
+                implicitCast.type,
+                IrTypeOperator.IMPLICIT_CAST,
+                implicitCast.typeOperand,
+                irIndexVarInitializer
+            )
+        )
+    }
+
+    private val IrTypeOperatorCall.operandKotlinType
+        get() = typeOperand.originalKotlinType!!
+
+    private class SamConversionsCollector(
+        private val compoundAssignmentInfo: CompoundAssignmentInfo
+    ) : IrElementVisitorVoid {
+        val samConversionsPerVariable = HashMap<IrVariable, MutableList<IrTypeOperatorCall>>()
+
+        override fun visitElement(element: IrElement) {
+            element.acceptChildrenVoid(this)
+        }
+
+        override fun visitTypeOperator(expression: IrTypeOperatorCall) {
+            expression.acceptChildrenVoid(this)
+            val irGetVar = expression.getSamConvertedGetValue()
+            if (irGetVar != null) {
+                val valueDeclaration = irGetVar.symbol.owner
+                if (valueDeclaration is IrVariable && valueDeclaration in compoundAssignmentInfo.indexVariables) {
+                    samConversionsPerVariable.getOrPut(valueDeclaration) { ArrayList() }.add(expression)
+                }
+            }
+        }
+    }
+
+    private class SamConversionsRewriter(
+        private val replacementVars: Map<IrVariable, IrVariable>
+    ) : IrElementTransformerVoid() {
+        override fun visitElement(element: IrElement): IrElement {
+            return element.apply { transformChildrenVoid() }
+        }
+
+        override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
+            val irGetVar = expression.getSamConvertedGetValue()
+            if (irGetVar != null) {
+                val valueDeclaration = irGetVar.symbol.owner
+                val replacementVar = replacementVars[valueDeclaration]
+                if (replacementVar != null) {
+                    return IrGetValueImpl(expression.startOffset, expression.endOffset, replacementVar.symbol, null)
+                }
+            }
+
+            return expression.apply { transformChildrenVoid() }
+        }
+
+        override fun visitGetValue(expression: IrGetValue): IrExpression {
+            val symbol = expression.symbol
+            if (symbol.owner in replacementVars) {
+                throw AssertionError(
+                    "SAM-converted index variable ${symbol.descriptor} is present in get/set calls in non-converted"
+                )
+            }
+            return expression
+        }
+    }
+
    private fun createLValue(
        kotlinType: KotlinType,
        irArrayValue: IntermediateValue,
@@ -126,4 +253,17 @@ class ArrayAccessAssignmentReceiver(
        value?.let { lastArgument = it }
        callGenerator.statementGenerator.generateSamConversionForValueArgumentsIfRequired(this, resolvedCall)
    }
+
+    companion object {
+        internal fun IrTypeOperatorCall.getSamConvertedGetValue(): IrGetValue? {
+            if (operator != IrTypeOperator.SAM_CONVERSION) return null
+            val arg0 = argument
+            if (arg0 !is IrTypeOperatorCall) return null
+            if (arg0.operator != IrTypeOperator.IMPLICIT_CAST) return null
+            val arg1 = arg0.argument
+            if (arg1 !is IrGetValue) return null
+            if (arg1.symbol.owner !is IrVariable) return null
+            return arg1
+        }
+    }
 }
@@ -23,14 +23,14 @@ import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.psi2ir.generators.CallGenerator

 class LValueWithGetterAndSetterCalls(
-    val callGenerator: CallGenerator,
-    val descriptor: CallableDescriptor,
-    val getterCall: () -> CallBuilder?,
-    val setterCall: (IrExpression) -> CallBuilder?,
+    private val callGenerator: CallGenerator,
+    private val descriptor: CallableDescriptor,
+    private val getterCall: () -> CallBuilder?,
+    private val setterCall: (IrExpression) -> CallBuilder?,
    override val type: IrType,
-    val startOffset: Int,
-    val endOffset: Int,
-    val origin: IrStatementOrigin? = null
+    private val startOffset: Int,
+    private val endOffset: Int,
+    private val origin: IrStatementOrigin? = null
 ) : LValue {

    override fun load(): IrExpression {