Specialize unsigned range and progression iterator calls

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

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.codegen.optimization.boxing
 import com.google.common.collect.ImmutableSet
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
-import org.jetbrains.kotlin.codegen.range.isRangeOrProgression
 import org.jetbrains.kotlin.codegen.optimization.common.OptimizationBasicInterpreter
 import org.jetbrains.kotlin.codegen.optimization.common.StrictBasicValue
 import org.jetbrains.kotlin.codegen.state.GenerationState
@@ -91,8 +90,8 @@ open class BoxingInterpreter(
                 ProgressionIteratorBasicValue.byProgressionClassType(firstArg.type)
             insn.isNextMethodCallOfProgressionIterator(values) -> {
                 val progressionIterator = firstArg as? ProgressionIteratorBasicValue
-                        ?: throw AssertionError("firstArg should be progression iterator")
+                    ?: throw AssertionError("firstArg should be progression iterator")
-                createNewBoxing(insn, AsmUtil.boxType(progressionIterator.valuesPrimitiveType), progressionIterator)
+                createNewBoxing(insn, progressionIterator.boxedElementType, progressionIterator)
             }
             insn.isAreEqualIntrinsicForSameTypedBoxedValues(values) && canValuesBeUnboxedForAreEqual(values, generationState) -> {
                 onAreEqual(insn, values[0] as BoxedBasicValue, values[1] as BoxedBasicValue)
@@ -265,13 +264,11 @@ fun AbstractInsnNode.isNextMethodCallOfProgressionIterator(values: List<BasicVal
 fun AbstractInsnNode.isIteratorMethodCallOfProgression(values: List<BasicValue>) =
     isMethodInsnWith(Opcodes.INVOKEINTERFACE) {
         val firstArgType = values.firstOrNull()?.type
-        firstArgType != null &&
+        name == "iterator" && firstArgType != null && isProgressionClass(firstArgType)
-                isProgressionClass(firstArgType) &&
-                name == "iterator"
     }
 
-fun isProgressionClass(type: Type) =
+private fun isProgressionClass(type: Type) =
-    isRangeOrProgression(buildFqNameByInternal(type.internalName))
+    ProgressionIteratorBasicValue.byProgressionClassType(type) != null
 
 fun AbstractInsnNode.isAreEqualIntrinsicForSameTypedBoxedValues(values: List<BasicValue>) =
     isAreEqualIntrinsic() && areSameTypedPrimitiveBoxedValues(values)

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

@@ -16,57 +16,62 @@
 
 package org.jetbrains.kotlin.codegen.optimization.boxing
 
-import org.jetbrains.kotlin.builtins.PrimitiveType
+import org.jetbrains.kotlin.codegen.AsmUtil
-import org.jetbrains.kotlin.codegen.intrinsics.IteratorNext
 import org.jetbrains.kotlin.codegen.optimization.common.StrictBasicValue
-import org.jetbrains.kotlin.codegen.range.getPrimitiveRangeOrProgressionElementType
+import org.jetbrains.kotlin.codegen.range.*
-import org.jetbrains.kotlin.codegen.range.supportedRangeTypes
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.jvm.JvmPrimitiveType
 import org.jetbrains.org.objectweb.asm.Type
 
 class ProgressionIteratorBasicValue
-private constructor(private val valuesPrimitiveTypeName: String) :
+private constructor(
-    StrictBasicValue(IteratorNext.getPrimitiveIteratorType(Name.identifier(valuesPrimitiveTypeName))) {
+    val nextMethodName: String,
+    iteratorType: Type,
+    private val primitiveElementType: Type,
+    val boxedElementType: Type
+) : StrictBasicValue(iteratorType) {
 
-    val valuesPrimitiveType: Type = VALUES_TYPENAME_TO_TYPE[valuesPrimitiveTypeName] ?: error("Unexpected type $valuesPrimitiveTypeName")
+    private constructor(typeName: String, valuesPrimitiveType: Type, valuesBoxedType: Type = AsmUtil.boxType(valuesPrimitiveType)) :
-
+            this("next$typeName", Type.getObjectType("kotlin/collections/${typeName}Iterator"), valuesPrimitiveType, valuesBoxedType)
-    val nextMethodName: String
-        get() = "next$valuesPrimitiveTypeName"
 
     val nextMethodDesc: String
-        get() = "()" + valuesPrimitiveType.descriptor
+        get() = "()" + primitiveElementType.descriptor
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (other == null || javaClass != other.javaClass) return false
         if (!super.equals(other)) return false
-
+        val value = other as ProgressionIteratorBasicValue
-        val value = other as ProgressionIteratorBasicValue?
+        return primitiveElementType == value.primitiveElementType
-
-        return valuesPrimitiveType == value!!.valuesPrimitiveType
     }
 
     override fun hashCode(): Int =
-        super.hashCode() * 31 + valuesPrimitiveType.hashCode()
+        super.hashCode() * 31 + nextMethodName.hashCode()
 
     companion object {
-        private val VALUES_TYPENAME_TO_TYPE: Map<String, Type> =
+        private val CHAR_PROGRESSION_ITERATOR_VALUE = ProgressionIteratorBasicValue("Char", Type.CHAR_TYPE)
-            supportedRangeTypes.associate { primitiveType ->
+        private val INT_PROGRESSION_ITERATOR_VALUE = ProgressionIteratorBasicValue("Int", Type.INT_TYPE)
-                primitiveType.typeName.asString() to Type.getType(JvmPrimitiveType.get(primitiveType).desc)
+        private val LONG_PROGRESSION_ITERATOR_VALUE = ProgressionIteratorBasicValue("Long", Type.LONG_TYPE)
-            }
 
-        private val ITERATOR_VALUE_BY_ELEMENT_PRIMITIVE_TYPE: Map<PrimitiveType, ProgressionIteratorBasicValue> =
+        private val UINT_PROGRESSION_ITERATOR_VALUE =
-            supportedRangeTypes.associate { elementType ->
+            ProgressionIteratorBasicValue("UInt", Type.INT_TYPE, Type.getObjectType("kotlin/UInt"))
-                elementType to ProgressionIteratorBasicValue(elementType.typeName.asString())
+        private val ULONG_PROGRESSION_ITERATOR_VALUE =
-            }
+            ProgressionIteratorBasicValue("ULong", Type.LONG_TYPE, Type.getObjectType("kotlin/ULong"))
 
-        fun byProgressionClassType(progressionClassType: Type): ProgressionIteratorBasicValue? {
+        private val PROGRESSION_CLASS_NAME_TO_ITERATOR_VALUE: Map<String, ProgressionIteratorBasicValue> =
-            val classFqName = FqName(progressionClassType.className)
+            hashMapOf(
-            val elementType = getPrimitiveRangeOrProgressionElementType(classFqName)
+                CHAR_RANGE_FQN to CHAR_PROGRESSION_ITERATOR_VALUE,
-            return ITERATOR_VALUE_BY_ELEMENT_PRIMITIVE_TYPE[elementType]
+                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
+            )
+
+        fun byProgressionClassType(progressionClassType: Type): ProgressionIteratorBasicValue? =
+            PROGRESSION_CLASS_NAME_TO_ITERATOR_VALUE[progressionClassType.className]
     }
 }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/range/RangeCodegenUtil.kt

@@ -6,15 +6,12 @@
 package org.jetbrains.kotlin.codegen.range
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns.RANGES_PACKAGE_FQ_NAME
-import org.jetbrains.kotlin.builtins.PrimitiveType
 import org.jetbrains.kotlin.builtins.UnsignedTypes
 import org.jetbrains.kotlin.codegen.AsmUtil.isPrimitiveNumberClassDescriptor
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.PsiDiagnosticUtils
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
@@ -23,18 +20,6 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.types.KotlinType
 
-val supportedRangeTypes = listOf(PrimitiveType.CHAR, PrimitiveType.INT, PrimitiveType.LONG)
-
-private val RANGE_TO_ELEMENT_TYPE: Map<FqName, PrimitiveType> =
-    supportedRangeTypes.associateBy {
-        RANGES_PACKAGE_FQ_NAME.child(Name.identifier(it.typeName.toString() + "Range"))
-    }
-
-private val PROGRESSION_TO_ELEMENT_TYPE: Map<FqName, PrimitiveType> =
-    supportedRangeTypes.associateBy {
-        RANGES_PACKAGE_FQ_NAME.child(Name.identifier(it.typeName.toString() + "Progression"))
-    }
-
 fun isPrimitiveRange(rangeType: KotlinType) =
     isClassTypeWithFqn(rangeType, PRIMITIVE_RANGE_FQNS)
 
@@ -53,14 +38,14 @@ private fun isClassTypeWithFqn(kotlinType: KotlinType, fqns: Set<String>): Boole
     return fqName.asString() in fqns
 }
 
-private const val CHAR_RANGE_FQN = "kotlin.ranges.CharRange"
+internal const val CHAR_RANGE_FQN = "kotlin.ranges.CharRange"
-private const val INT_RANGE_FQN = "kotlin.ranges.IntRange"
+internal const val INT_RANGE_FQN = "kotlin.ranges.IntRange"
-private const val LONG_RANGE_FQN = "kotlin.ranges.LongRange"
+internal const val LONG_RANGE_FQN = "kotlin.ranges.LongRange"
 private val PRIMITIVE_RANGE_FQNS = setOf(CHAR_RANGE_FQN, INT_RANGE_FQN, LONG_RANGE_FQN)
 
-private const val CHAR_PROGRESSION_FQN = "kotlin.ranges.CharProgression"
+internal const val CHAR_PROGRESSION_FQN = "kotlin.ranges.CharProgression"
-private const val INT_PROGRESSION_FQN = "kotlin.ranges.IntProgression"
+internal const val INT_PROGRESSION_FQN = "kotlin.ranges.IntProgression"
-private const val LONG_PROGRESSION_FQN = "kotlin.ranges.LongProgression"
+internal const val LONG_PROGRESSION_FQN = "kotlin.ranges.LongProgression"
 private val PRIMITIVE_PROGRESSION_FQNS = setOf(CHAR_PROGRESSION_FQN, INT_PROGRESSION_FQN, LONG_PROGRESSION_FQN)
 
 private const val CLOSED_FLOAT_RANGE_FQN = "kotlin.ranges.ClosedFloatRange"
@@ -69,12 +54,12 @@ private const val CLOSED_RANGE_FQN = "kotlin.ranges.ClosedRange"
 private const val CLOSED_FLOATING_POINT_RANGE_FQN = "kotlin.ranges.ClosedFloatingPointRange"
 private const val COMPARABLE_RANGE_FQN = "kotlin.ranges.ComparableRange"
 
-private const val UINT_RANGE_FQN = "kotlin.ranges.UIntRange"
+internal const val UINT_RANGE_FQN = "kotlin.ranges.UIntRange"
-private const val ULONG_RANGE_FQN = "kotlin.ranges.ULongRange"
+internal const val ULONG_RANGE_FQN = "kotlin.ranges.ULongRange"
 private val UNSIGNED_RANGE_FQNS = setOf(UINT_RANGE_FQN, ULONG_RANGE_FQN)
 
-private const val UINT_PROGRESSION_FQN = "kotlin.ranges.UIntProgression"
+internal const val UINT_PROGRESSION_FQN = "kotlin.ranges.UIntProgression"
-private const val ULONG_PROGRESSION_FQN = "kotlin.ranges.ULongProgression"
+internal const val ULONG_PROGRESSION_FQN = "kotlin.ranges.ULongProgression"
 private val UNSIGNED_PROGRESSION_FQNS = setOf(UINT_PROGRESSION_FQN, ULONG_PROGRESSION_FQN)
 
 fun getRangeOrProgressionElementType(rangeType: KotlinType): KotlinType? {
@@ -113,12 +98,6 @@ fun BindingContext.getElementType(forExpression: KtForExpression): KotlinType {
     return nextCall.resultingDescriptor.returnType!!
 }
 
-fun getPrimitiveRangeOrProgressionElementType(rangeOrProgressionName: FqName): PrimitiveType? =
-    RANGE_TO_ELEMENT_TYPE[rangeOrProgressionName] ?: PROGRESSION_TO_ELEMENT_TYPE[rangeOrProgressionName]
-
-fun isRangeOrProgression(className: FqName) =
-    getPrimitiveRangeOrProgressionElementType(className) != null
-
 fun isPrimitiveNumberRangeTo(rangeTo: CallableDescriptor) =
     "rangeTo" == rangeTo.name.asString() && isPrimitiveNumberClassDescriptor(rangeTo.containingDeclaration) ||
             isPrimitiveRangeToExtension(rangeTo)

compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt

@@ -0,0 +1,44 @@
+// WITH_RUNTIME
+// IGNORE_BACKEND: JVM_IR, JS_IR
+
+fun testUIntRangeForEach() {
+    var s = 0
+    (1u .. 5u).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 12345) throw AssertionError("$s")
+}
+
+fun testUIntProgressionForEach() {
+    var s = 0
+    (5u downTo 1u).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 54321) throw AssertionError("$s")
+}
+
+fun testULongRangeForEach() {
+    var s = 0
+    (1UL .. 5UL).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 12345) throw AssertionError("$s")
+}
+
+fun testULongProgressionForEach() {
+    var s = 0
+    (5UL downTo 1UL).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 54321) throw AssertionError("$s")
+}
+
+
+fun box(): String {
+    testUIntRangeForEach()
+    testUIntProgressionForEach()
+    testULongRangeForEach()
+    testULongProgressionForEach()
+
+    return "OK"
+}

compiler/testData/codegen/bytecodeText/boxingOptimization/unsignedRangeIteratorSpecialization.kt

@@ -0,0 +1,38 @@
+// WITH_RUNTIME
+// IGNORE_BACKEND: JVM_IR
+
+fun testUIntRangeForEach() {
+    var s = 0
+    (1u .. 5u).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 12345) throw AssertionError("$s")
+}
+
+fun testUIntProgressionForEach() {
+    var s = 0
+    (5u downTo 1u).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 54321) throw AssertionError("$s")
+}
+
+fun testULongRangeForEach() {
+    var s = 0
+    (1UL .. 5UL).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 12345) throw AssertionError("$s")
+}
+
+fun testULongProgressionForEach() {
+    var s = 0
+    (5UL downTo 1UL).forEach {
+        s = s * 10 + it.toInt()
+    }
+    if (s != 54321) throw AssertionError("$s")
+}
+
+// 0 java/util/Iterator.next \(\)Ljava/lang/Object;
+// 2 nextUInt
+// 2 nextULong

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -24323,6 +24323,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedLongRemainder_jvm18.kt");
         }
 
+        @TestMetadata("unsignedRangeIterator.kt")
+        public void testUnsignedRangeIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt");
+        }
+
         @TestMetadata("unsignedTypePrefixIncrementDecrementBoxing.kt")
         public void testUnsignedTypePrefixIncrementDecrementBoxing() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedTypePrefixIncrementDecrementBoxing.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -542,6 +542,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/boxingOptimization/unsafeRemoving.kt");
         }
 
+        @TestMetadata("unsignedRangeIteratorSpecialization.kt")
+        public void testUnsignedRangeIteratorSpecialization() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/boxingOptimization/unsignedRangeIteratorSpecialization.kt");
+        }
+
         @TestMetadata("variableClash.kt")
         public void testVariableClash() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/boxingOptimization/variableClash.kt");

compiler/tests/org/jetbrains/kotlin/codegen/IrBytecodeTextTestGenerated.java

@@ -542,6 +542,11 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/boxingOptimization/unsafeRemoving.kt");
         }
 
+        @TestMetadata("unsignedRangeIteratorSpecialization.kt")
+        public void testUnsignedRangeIteratorSpecialization() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/boxingOptimization/unsignedRangeIteratorSpecialization.kt");
+        }
+
         @TestMetadata("variableClash.kt")
         public void testVariableClash() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/boxingOptimization/variableClash.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -24323,6 +24323,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedLongRemainder_jvm18.kt");
         }
 
+        @TestMetadata("unsignedRangeIterator.kt")
+        public void testUnsignedRangeIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt");
+        }
+
         @TestMetadata("unsignedTypePrefixIncrementDecrementBoxing.kt")
         public void testUnsignedTypePrefixIncrementDecrementBoxing() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedTypePrefixIncrementDecrementBoxing.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -24328,6 +24328,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedLongRemainder_jvm18.kt");
         }
 
+        @TestMetadata("unsignedRangeIterator.kt")
+        public void testUnsignedRangeIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt");
+        }
+
         @TestMetadata("unsignedTypePrefixIncrementDecrementBoxing.kt")
         public void testUnsignedTypePrefixIncrementDecrementBoxing() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedTypePrefixIncrementDecrementBoxing.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/IrJsCodegenBoxTestGenerated.java

@@ -18763,6 +18763,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedLongRemainder_jvm18.kt");
         }
 
+        @TestMetadata("unsignedRangeIterator.kt")
+        public void testUnsignedRangeIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt");
+        }
+
         @TestMetadata("unsignedTypePrefixIncrementDecrementBoxing.kt")
         public void testUnsignedTypePrefixIncrementDecrementBoxing() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedTypePrefixIncrementDecrementBoxing.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -19813,6 +19813,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedLongRemainder_jvm18.kt");
         }
 
+        @TestMetadata("unsignedRangeIterator.kt")
+        public void testUnsignedRangeIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt");
+        }
+
         @TestMetadata("unsignedTypePrefixIncrementDecrementBoxing.kt")
         public void testUnsignedTypePrefixIncrementDecrementBoxing() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/unsignedTypePrefixIncrementDecrementBoxing.kt");