JVM KT-49548 progression iterators can be tainted

Author: dnpetrov

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/boxing/BoxedBasicValue.kt

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.resolve.isInlineClass
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
-import java.util.*
 
 abstract class BoxedBasicValue(type: Type) : StrictBasicValue(type) {
     abstract val descriptor: BoxedValueDescriptor

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/boxing/BoxingInterpreter.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
 import org.jetbrains.kotlin.codegen.optimization.common.OptimizationBasicInterpreter
 import org.jetbrains.kotlin.codegen.optimization.common.StrictBasicValue
+import org.jetbrains.kotlin.codegen.range.*
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.codegen.topLevelClassInternalName
@@ -35,13 +36,13 @@ import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
 import org.jetbrains.org.objectweb.asm.tree.InsnList
 import org.jetbrains.org.objectweb.asm.tree.MethodInsnNode
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
-import java.util.*
 
 open class BoxingInterpreter(
     private val insnList: InsnList,
     private val generationState: GenerationState
 ) : OptimizationBasicInterpreter() {
     private val boxingPlaces = HashMap<Int, BoxedBasicValue>()
+    private val progressionIterators = HashMap<AbstractInsnNode, ProgressionIteratorBasicValue>()
 
     protected open fun createNewBoxing(
         insn: AbstractInsnNode,
@@ -89,8 +90,18 @@ open class BoxingInterpreter(
                 onUnboxing(insn, firstArg, value.type)
                 value
             }
-            insn.isIteratorMethodCallOfProgression(values) ->
-                ProgressionIteratorBasicValue.byProgressionClassType(firstArg.type)
+            insn.isIteratorMethodCall() -> {
+                values.markBoxedArgumentValues()
+                val firstArgType = firstArg.type
+                if (isProgressionClass(firstArgType)) {
+                    progressionIterators.getOrPut(insn) {
+                        ProgressionIteratorBasicValue.byProgressionClassType(insn, firstArgType)!!
+                    }
+                } else {
+                    progressionIterators[insn]?.taint()
+                    value
+                }
+            }
             insn.isNextMethodCallOfProgressionIterator(values) -> {
                 val progressionIterator = firstArg as? ProgressionIteratorBasicValue
                     ?: throw AssertionError("firstArg should be progression iterator")
@@ -269,14 +280,27 @@ fun AbstractInsnNode.isNextMethodCallOfProgressionIterator(values: List<BasicVal
                 name == "next"
             }
 
+fun AbstractInsnNode.isIteratorMethodCall() =
+    isMethodInsnWith(Opcodes.INVOKEINTERFACE) {
+        name == "iterator" && desc == "()Ljava/util/Iterator;"
+    }
+
 fun AbstractInsnNode.isIteratorMethodCallOfProgression(values: List<BasicValue>) =
     isMethodInsnWith(Opcodes.INVOKEINTERFACE) {
         val firstArgType = values.firstOrNull()?.type
-        name == "iterator" && firstArgType != null && isProgressionClass(firstArgType)
+        name == "iterator" && desc == "()Ljava/util/Iterator;" &&
+                firstArgType != null && isProgressionClass(firstArgType)
     }
 
+private val PROGRESSION_CLASS_FQNS = setOf(
+    CHAR_RANGE_FQN, CHAR_PROGRESSION_FQN,
+    INT_RANGE_FQN, INT_PROGRESSION_FQN,
+    LONG_RANGE_FQN, LONG_PROGRESSION_FQN
+)
+
 private fun isProgressionClass(type: Type) =
-    ProgressionIteratorBasicValue.byProgressionClassType(type) != null
+    type.className in PROGRESSION_CLASS_FQNS
+
 
 fun AbstractInsnNode.isAreEqualIntrinsicForSameTypedBoxedValues(values: List<BasicValue>) =
     isAreEqualIntrinsic() && areSameTypedPrimitiveBoxedValues(values)

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/boxing/ProgressionIteratorBasicValue.kt

@@ -20,17 +20,23 @@ import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.optimization.common.StrictBasicValue
 import org.jetbrains.kotlin.codegen.range.*
 import org.jetbrains.org.objectweb.asm.Type
+import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
 
 class ProgressionIteratorBasicValue
 private constructor(
+    val iteratorCallInsn: AbstractInsnNode,
     val nextMethodName: String,
     iteratorType: Type,
     private val primitiveElementType: Type,
     val boxedElementType: Type
 ) : StrictBasicValue(iteratorType) {
 
-    private constructor(typeName: String, valuesPrimitiveType: Type, valuesBoxedType: Type = AsmUtil.boxType(valuesPrimitiveType)) :
-            this("next$typeName", Type.getObjectType("kotlin/collections/${typeName}Iterator"), valuesPrimitiveType, valuesBoxedType)
+    var tainted = false
+        private set
+
+    fun taint() {
+        tainted = true
+    }
 
     val nextMethodDesc: String
         get() = "()" + primitiveElementType.descriptor
@@ -47,32 +53,35 @@ private constructor(
         super.hashCode() * 31 + nextMethodName.hashCode()
 
     companion object {
-        private val CHAR_PROGRESSION_ITERATOR_VALUE = ProgressionIteratorBasicValue("Char", Type.CHAR_TYPE)
-        private val INT_PROGRESSION_ITERATOR_VALUE = ProgressionIteratorBasicValue("Int", Type.INT_TYPE)
-        private val LONG_PROGRESSION_ITERATOR_VALUE = ProgressionIteratorBasicValue("Long", Type.LONG_TYPE)
-
         // TODO functions returning inline classes are mangled now, should figure out how to work with UInt/ULong iterators here
-        // private val UINT_PROGRESSION_ITERATOR_VALUE =
         //     ProgressionIteratorBasicValue("UInt", Type.INT_TYPE, Type.getObjectType("kotlin/UInt"))
-        // private val ULONG_PROGRESSION_ITERATOR_VALUE =
         //     ProgressionIteratorBasicValue("ULong", Type.LONG_TYPE, Type.getObjectType("kotlin/ULong"))
 
-        private val PROGRESSION_CLASS_NAME_TO_ITERATOR_VALUE: Map<String, ProgressionIteratorBasicValue> =
-            hashMapOf(
-                CHAR_RANGE_FQN to CHAR_PROGRESSION_ITERATOR_VALUE,
-                CHAR_PROGRESSION_FQN to CHAR_PROGRESSION_ITERATOR_VALUE,
-                INT_RANGE_FQN to INT_PROGRESSION_ITERATOR_VALUE,
-                INT_PROGRESSION_FQN to INT_PROGRESSION_ITERATOR_VALUE,
-                LONG_RANGE_FQN to LONG_PROGRESSION_ITERATOR_VALUE,
-                LONG_PROGRESSION_FQN to LONG_PROGRESSION_ITERATOR_VALUE,
-                // UINT_RANGE_FQN to UINT_PROGRESSION_ITERATOR_VALUE,
-                // UINT_PROGRESSION_FQN to UINT_PROGRESSION_ITERATOR_VALUE,
-                // ULONG_RANGE_FQN to ULONG_PROGRESSION_ITERATOR_VALUE,
-                // ULONG_PROGRESSION_FQN to ULONG_PROGRESSION_ITERATOR_VALUE
+        private fun progressionIteratorValue(
+            iteratorCallInsn: AbstractInsnNode,
+            typeName: String,
+            valuesPrimitiveType: Type,
+            valuesBoxedType: Type = AsmUtil.boxType(valuesPrimitiveType)
+        ) =
+            ProgressionIteratorBasicValue(
+                iteratorCallInsn,
+                "next$typeName",
+                Type.getObjectType("kotlin/collections/${typeName}Iterator"),
+                valuesPrimitiveType,
+                valuesBoxedType
             )
 
-        fun byProgressionClassType(progressionClassType: Type): ProgressionIteratorBasicValue? =
-            PROGRESSION_CLASS_NAME_TO_ITERATOR_VALUE[progressionClassType.className]
+        fun byProgressionClassType(iteratorCallInsn: AbstractInsnNode, progressionClassType: Type): ProgressionIteratorBasicValue? =
+            when (progressionClassType.className) {
+                CHAR_RANGE_FQN, CHAR_PROGRESSION_FQN ->
+                    progressionIteratorValue(iteratorCallInsn, "Char", Type.CHAR_TYPE)
+                INT_RANGE_FQN, INT_PROGRESSION_FQN ->
+                    progressionIteratorValue(iteratorCallInsn, "Int", Type.INT_TYPE)
+                LONG_RANGE_FQN, LONG_PROGRESSION_FQN ->
+                    progressionIteratorValue(iteratorCallInsn, "Long", Type.LONG_TYPE)
+                else ->
+                    null
+            }
     }
 }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/boxing/RedundantBoxingMethodTransformer.kt

@@ -33,7 +33,6 @@ import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.tree.*
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
 import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
-import java.util.*
 
 class RedundantBoxingMethodTransformer(private val generationState: GenerationState) : MethodTransformer() {
 
@@ -50,6 +49,9 @@ class RedundantBoxingMethodTransformer(private val generationState: GenerationSt
         val valuesToOptimize = interpreter.candidatesBoxedValues
 
         if (!valuesToOptimize.isEmpty) {
+            // has side effect on valuesToOptimize
+            removeValuesFromTaintedProgressionIterators(valuesToOptimize)
+
             // has side effect on valuesToOptimize and frames, containing BoxedBasicValues that are unsafe to remove
             removeValuesClashingWithVariables(valuesToOptimize, node, frames)
 
@@ -86,7 +88,7 @@ class RedundantBoxingMethodTransformer(private val generationState: GenerationSt
     private fun removeValuesClashingWithVariables(
         values: RedundantBoxedValuesCollection,
         node: MethodNode,
-        frames: Array<Frame<BasicValue>>
+        frames: Array<Frame<BasicValue>?>
     ) {
         while (removeValuesClashingWithVariablesPass(values, node, frames)) {
             // do nothing
@@ -126,6 +128,15 @@ class RedundantBoxingMethodTransformer(private val generationState: GenerationSt
         return needToRepeat
     }
 
+    private fun removeValuesFromTaintedProgressionIterators(valuesToOptimize: RedundantBoxedValuesCollection) {
+        for (descriptor in valuesToOptimize.toList()) {
+            val progressionIterator = descriptor?.progressionIterator ?: continue
+            if (progressionIterator.tainted) {
+                valuesToOptimize.remove(descriptor)
+            }
+        }
+    }
+
     private fun isUnsafeToRemoveBoxingForConnectedValues(usedValues: List<BasicValue>, unboxedType: Type): Boolean =
         usedValues.any { input ->
             if (input === StrictBasicValue.UNINITIALIZED_VALUE) return@any false
@@ -135,7 +146,7 @@ class RedundantBoxingMethodTransformer(private val generationState: GenerationSt
             !descriptor.isSafeToRemove || descriptor.unboxedType != unboxedType
         }
 
-    private fun adaptLocalVariableTableForBoxedValues(node: MethodNode, frames: Array<Frame<BasicValue>>) {
+    private fun adaptLocalVariableTableForBoxedValues(node: MethodNode, frames: Array<Frame<BasicValue>?>) {
         for (localVariableNode in node.localVariables) {
             if (Type.getType(localVariableNode.desc).sort != Type.OBJECT) {
                 continue

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/nullCheck/NullabilityInterpreter.kt

@@ -84,7 +84,7 @@ class NullabilityInterpreter(private val generationState: GenerationState) : Opt
             insn.isBoxing(generationState) ->
                 NotNullBasicValue(resultType)
             insn.isIteratorMethodCallOfProgression(values) ->
-                ProgressionIteratorBasicValue.byProgressionClassType(values[0].type)
+                ProgressionIteratorBasicValue.byProgressionClassType(insn, values[0].type)
             insn.isNextMethodCallOfProgressionIterator(values) ->
                 NotNullBasicValue(resultType)
             insn.isPseudo(PseudoInsn.AS_NOT_NULL) ->

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -0,0 +1,16 @@
+// IGNORE_BACKEND: WASM
+// WITH_RUNTIME
+
+val p0 = 0..3
+
+fun test(): List<Int> {
+    val progression = if (p0.last == 3) p0 + 1 else p0
+    return progression.map { it }
+}
+
+fun box(): String {
+    val t = test()
+    if (t != listOf(0, 1, 2, 3, 1))
+        return "Failed: t=$t"
+    return "OK"
+}

compiler/testData/codegen/box/boxingOptimization/kt49548.kt

@@ -0,0 +1,14 @@
+// IGNORE_BACKEND: WASM
+// WITH_RUNTIME
+import kotlin.test.assertEquals
+
+fun test() {
+    var i = 0
+    val a = ubyteArrayOf(3u, 2u, 1u)
+    a.forEach { e -> assertEquals(e, a[i++]) }
+}
+
+fun box(): String {
+    test()
+    return "OK"
+}