KT-32044 Fix unsigned integer coercion in loop intrinsics

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

@@ -33,8 +33,8 @@ class DownToProgressionRangeValue(rangeCall: ResolvedCall<out CallableDescriptor
 
     override fun getBoundedValue(codegen: ExpressionCodegen) =
         BoundedValue(
-            lowBound = codegen.generateCallSingleArgument(rangeCall),
+            lowBound = codegen.generateCallSingleArgument(rangeCall).coerceToRangeElementTypeIfRequired(),
-            highBound = codegen.generateCallReceiver(rangeCall)
+            highBound = codegen.generateCallReceiver(rangeCall).coerceToRangeElementTypeIfRequired()
         )
 
     override fun createForLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression) =

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

@@ -16,6 +16,9 @@
 
 package org.jetbrains.kotlin.codegen.range
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.UnsignedType
+import org.jetbrains.kotlin.builtins.UnsignedTypes
 import org.jetbrains.kotlin.codegen.ExpressionCodegen
 import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.range.comparison.getComparisonGeneratorForKotlinType
@@ -32,6 +35,7 @@ import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.constants.*
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.Type
 
@@ -77,6 +81,84 @@ abstract class PrimitiveNumberRangeIntrinsicRangeValue(
 
     protected abstract fun getBoundedValue(codegen: ExpressionCodegen): BoundedValue
 
+    protected fun StackValue.coerceToRangeElementTypeIfRequired(): StackValue {
+        val rangeKotlinType = rangeCall.resultingDescriptor.returnType!!
+        val rangeElementKotlinType = getRangeOrProgressionElementType(rangeKotlinType)!!
+        return when {
+            KotlinBuiltIns.isUInt(rangeElementKotlinType) ->
+                coerceUnsignedToUInt(this, rangeElementKotlinType)
+            KotlinBuiltIns.isULong(rangeElementKotlinType) ->
+                coerceUnsignedToULong(this, rangeElementKotlinType)
+            else ->
+                this
+        }
+    }
+
+    private val StackValue.unsignedType: UnsignedType?
+        get() = kotlinType?.let { UnsignedTypes.toUnsignedType(it) }
+
+    private fun coerceUnsignedToUInt(stackValue: StackValue, uIntKotlinType: KotlinType): StackValue {
+        val valueKotlinType = stackValue.kotlinType
+        val valueUnsignedType = stackValue.unsignedType
+            ?: throw AssertionError("Unsigned type expected: $valueKotlinType")
+
+        if (valueUnsignedType == UnsignedType.UINT) return stackValue
+
+        return StackValue.operation(Type.INT_TYPE, uIntKotlinType) { v ->
+            stackValue.put(stackValue.type, valueKotlinType, v)
+            when (valueUnsignedType) {
+                UnsignedType.UBYTE -> {
+                    v.iconst(0xFF)
+                    v.and(Type.INT_TYPE)
+                }
+
+                UnsignedType.USHORT -> {
+                    v.iconst(0xFFFF)
+                    v.and(Type.INT_TYPE)
+                }
+
+                UnsignedType.ULONG -> {
+                    v.cast(Type.LONG_TYPE, Type.INT_TYPE)
+                }
+
+                else -> throw AssertionError("Unexpected value type: $valueKotlinType")
+            }
+        }
+    }
+
+    private fun coerceUnsignedToULong(stackValue: StackValue, uLongKotlinType: KotlinType): StackValue {
+        val valueKotlinType = stackValue.kotlinType
+        val valueUnsignedType = stackValue.unsignedType
+            ?: throw AssertionError("Unsigned type expected: $valueKotlinType")
+
+        if (valueUnsignedType == UnsignedType.ULONG) return stackValue
+
+        return StackValue.operation(Type.LONG_TYPE, uLongKotlinType) { v ->
+            stackValue.put(stackValue.type, valueKotlinType, v)
+            when (valueUnsignedType) {
+                UnsignedType.UBYTE -> {
+                    v.cast(Type.INT_TYPE, Type.LONG_TYPE)
+                    v.lconst(0xFF)
+                    v.and(Type.LONG_TYPE)
+                }
+
+                UnsignedType.USHORT -> {
+                    v.cast(Type.INT_TYPE, Type.LONG_TYPE)
+                    v.lconst(0xFFFF)
+                    v.and(Type.LONG_TYPE)
+                }
+
+                UnsignedType.UINT -> {
+                    v.cast(Type.INT_TYPE, Type.LONG_TYPE)
+                    v.lconst(0xFFFF_FFFFL)
+                    v.and(Type.LONG_TYPE)
+                }
+
+                else -> throw AssertionError("Unexpected value type: $valueKotlinType")
+            }
+        }
+    }
+
     protected fun createConstBoundedForLoopGeneratorOrNull(
         codegen: ExpressionCodegen,
         forExpression: KtForExpression,

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.codegen.range
 
 import org.jetbrains.kotlin.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.generateCallReceiver
 import org.jetbrains.kotlin.codegen.generateCallSingleArgument
 import org.jetbrains.kotlin.codegen.range.comparison.getComparisonGeneratorForKotlinType
@@ -37,15 +38,20 @@ class PrimitiveNumberRangeLiteralRangeValue(
     ReversableRangeValue {
 
     override fun getBoundedValue(codegen: ExpressionCodegen): BoundedValue {
-        val instanceType = codegen.asmType(rangeCall.resultingDescriptor.returnType!!)
         val lowBound = codegen.generateCallReceiver(rangeCall)
+
         if (codegen.canBeSpecializedByExcludingHighBound(rangeCall)) {
-            val highBound = (rangeCall.getFirstArgumentExpression() as KtBinaryExpression).left
+            return BoundedValue(
-            return BoundedValue(lowBound, true, codegen.gen(highBound), false)
+                lowBound = lowBound,
+                isLowInclusive = true,
+                highBound = codegen.gen((rangeCall.getFirstArgumentExpression() as KtBinaryExpression).left),
+                isHighInclusive = false
+            )
         }
+
         return BoundedValue(
-            lowBound = lowBound,
+            lowBound = lowBound.coerceToRangeElementTypeIfRequired(),
-            highBound = codegen.generateCallSingleArgument(rangeCall)
+            highBound = codegen.generateCallSingleArgument(rangeCall).coerceToRangeElementTypeIfRequired()
         )
     }
 

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

@@ -32,10 +32,10 @@ class PrimitiveNumberUntilRangeValue(rangeCall: ResolvedCall<out CallableDescrip
 
     override fun getBoundedValue(codegen: ExpressionCodegen) =
         BoundedValue(
-            codegen.generateCallReceiver(rangeCall),
+            lowBound = codegen.generateCallReceiver(rangeCall).coerceToRangeElementTypeIfRequired(),
-            true,
+            isLowInclusive = true,
-            codegen.generateCallSingleArgument(rangeCall),
+            highBound = codegen.generateCallSingleArgument(rangeCall).coerceToRangeElementTypeIfRequired(),
-            false
+            isHighInclusive = false
         )
 
     override fun createForLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression) =

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

@@ -32,11 +32,15 @@ fun isPrimitiveProgression(rangeType: KotlinType) =
 fun isUnsignedProgression(rangeType: KotlinType) =
     isClassTypeWithFqn(rangeType, UNSIGNED_PROGRESSION_FQNS)
 
-private fun isClassTypeWithFqn(kotlinType: KotlinType, fqns: Set<String>): Boolean {
+private val KotlinType.classFqnString: String?
-    val declarationDescriptor = kotlinType.constructor.declarationDescriptor as? ClassDescriptor ?: return false
+    get() {
-    val fqName = DescriptorUtils.getFqName(declarationDescriptor).takeIf { it.isSafe } ?: return false
+        val declarationDescriptor = constructor.declarationDescriptor as? ClassDescriptor ?: return null
-    return fqName.asString() in fqns
+        val fqn = DescriptorUtils.getFqName(declarationDescriptor)
-}
+        return if (fqn.isSafe) fqn.asString() else null
+    }
+
+private fun isClassTypeWithFqn(kotlinType: KotlinType, fqns: Set<String>): Boolean =
+    kotlinType.classFqnString in fqns
 
 internal const val CHAR_RANGE_FQN = "kotlin.ranges.CharRange"
 internal const val INT_RANGE_FQN = "kotlin.ranges.IntRange"

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

@@ -0,0 +1,58 @@
+// KJS_WITH_FULL_RUNTIME
+// IGNORE_BACKEND: JVM_IR
+// WITH_RUNTIME
+
+val UB_MAX = UByte.MAX_VALUE
+val UB_START = (UB_MAX - 10u).toUByte()
+
+val US_MAX = UShort.MAX_VALUE
+val US_START = (US_MAX - 10u).toUShort()
+
+fun testUByteLoopWithCoercion1() {
+    for (x in UB_START..UB_MAX) {
+        if (x > UB_MAX.toUInt()) throw AssertionError()
+    }
+}
+
+fun testUByteLoopWithCoercion2() {
+    for (x in UB_START until UB_MAX) {
+        if (x > UB_MAX.toUInt()) throw AssertionError()
+    }
+}
+
+fun testUByteLoopWithCoercion3() {
+    for (x in UB_MAX downTo UB_START) {
+        if (x > UB_MAX.toUInt()) throw AssertionError()
+    }
+}
+
+
+fun testUShortLoopWithCoercion1() {
+    for (x in US_START..US_MAX) {
+        if (x > US_MAX.toUInt()) throw AssertionError()
+    }
+}
+
+fun testUShortLoopWithCoercion2() {
+    for (x in US_START until US_MAX) {
+        if (x > US_MAX.toUInt()) throw AssertionError()
+    }
+}
+
+fun testUShortLoopWithCoercion3() {
+    for (x in US_MAX downTo US_START) {
+        if (x > US_MAX.toUInt()) throw AssertionError()
+    }
+}
+
+
+fun box(): String {
+    testUByteLoopWithCoercion1()
+    testUByteLoopWithCoercion2()
+    testUByteLoopWithCoercion3()
+    testUShortLoopWithCoercion1()
+    testUShortLoopWithCoercion2()
+    testUShortLoopWithCoercion3()
+
+    return "OK"
+}

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

@@ -25100,6 +25100,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeLiteral.kt");
         }
 
+        @TestMetadata("forInUnsignedRangeWithCoercion.kt")
+        public void testForInUnsignedRangeWithCoercion() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeWithCoercion.kt");
+        }
+
         @TestMetadata("forInUnsignedUntil.kt")
         public void testForInUnsignedUntil() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedUntil.kt");

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

@@ -25100,6 +25100,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeLiteral.kt");
         }
 
+        @TestMetadata("forInUnsignedRangeWithCoercion.kt")
+        public void testForInUnsignedRangeWithCoercion() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeWithCoercion.kt");
+        }
+
         @TestMetadata("forInUnsignedUntil.kt")
         public void testForInUnsignedUntil() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedUntil.kt");

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

@@ -25120,6 +25120,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeLiteral.kt");
         }
 
+        @TestMetadata("forInUnsignedRangeWithCoercion.kt")
+        public void testForInUnsignedRangeWithCoercion() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeWithCoercion.kt");
+        }
+
         @TestMetadata("forInUnsignedUntil.kt")
         public void testForInUnsignedUntil() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedUntil.kt");

core/descriptors/src/org/jetbrains/kotlin/builtins/UnsignedType.kt

@@ -50,6 +50,15 @@ object UnsignedTypes {
         return isUnsignedClass(descriptor)
     }
 
+    fun toUnsignedType(type: KotlinType): UnsignedType? =
+        when {
+            KotlinBuiltIns.isUByte(type) -> UnsignedType.UBYTE
+            KotlinBuiltIns.isUShort(type) -> UnsignedType.USHORT
+            KotlinBuiltIns.isUInt(type) -> UnsignedType.UINT
+            KotlinBuiltIns.isULong(type) -> UnsignedType.ULONG
+            else -> null
+        }
+
     fun isUnsignedClass(descriptor: DeclarationDescriptor): Boolean {
         val container = descriptor.containingDeclaration
         return container is PackageFragmentDescriptor &&

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

@@ -19280,6 +19280,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeLiteral.kt");
         }
 
+        @TestMetadata("forInUnsignedRangeWithCoercion.kt")
+        public void testForInUnsignedRangeWithCoercion() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeWithCoercion.kt");
+        }
+
         @TestMetadata("forInUnsignedUntil.kt")
         public void testForInUnsignedUntil() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedUntil.kt");

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

@@ -20435,6 +20435,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeLiteral.kt");
         }
 
+        @TestMetadata("forInUnsignedRangeWithCoercion.kt")
+        public void testForInUnsignedRangeWithCoercion() throws Exception {
+            runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedRangeWithCoercion.kt");
+        }
+
         @TestMetadata("forInUnsignedUntil.kt")
         public void testForInUnsignedUntil() throws Exception {
             runTest("compiler/testData/codegen/box/unsignedTypes/forInUnsignedUntil.kt");