JVM KT-49548 progression iterators can be tainted

2021-11-09 13:50:06 +03:00
parent 646d156bde
commit 9a4cff0dc7
15 changed files with 193 additions and 32 deletions
@@ -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
@@ -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)
@@ -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
+            }
    }
 }

@@ -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
@@ -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) ->