RangeContainsLowering: Handle unsigned ranges.

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/RangeContainsLowering.kt

@@ -28,18 +28,13 @@ import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addIfNotNull
 
-// TODO:
-// - Handle contains on ClosedFloatingPointRange non-literal expression (similar to DefaultProgressionHandler)
-// - Note for unsigned until, the decremented last is signed (see UntilHandler.kt:81). Rather than converting back to unsigned,
-//     should we remove the conversion call instead?
-// - Unsigned step has similar concerns as well. getProgressionLastElement return value is signed.
-
 val rangeContainsLoweringPhase = makeIrFilePhase(
     ::RangeContainsLowering,
     name = "RangeContainsLowering",
@@ -235,7 +230,23 @@ private class Transformer(
         // **  - Bound with side effect is stored in a temp variable to ensure evaluation even if right side is short-circuited.
         // *** - Bound with side effect is on left side of && to make sure it always gets evaluated.
 
-        var (argVar, argExpression) = createTemporaryVariableIfNecessary(argument, "containsArg")
+        var arg = argument
+        val builtIns = context.irBuiltIns
+        val comparisonClass = if (isNumericRange) {
+            computeComparisonClass(this@Transformer.context.ir.symbols, lower.type, upper.type, arg.type) ?: return null
+        } else {
+            assert(headerInfo is ComparableRangeInfo)
+            this@Transformer.context.ir.symbols.comparable.owner
+        }
+
+        if (isNumericRange) {
+            // Convert argument to the "widest" common numeric type for comparisons.
+            // Note that we do the same for the bounds below. If it is necessary to convert the argument, it's better to do it once and
+            // store in a temp variable, since it is used twice in the transformed expression (bounds are only used once).
+            arg = arg.castIfNecessary(comparisonClass)
+        }
+
+        val (argVar, argExpression) = createTemporaryVariableIfNecessary(arg, "containsArg")
         var lowerExpression: IrExpression
         var upperExpression: IrExpression
         val useLowerClauseOnLeftSide: Boolean
@@ -273,27 +284,16 @@ private class Transformer(
         }
         additionalStatements.addIfNotNull(argVar)
 
-        // TODO: Handle unsigned (comparisonClass is currently null).
-        val builtIns = context.irBuiltIns
-        val comparisonClass = if (isNumericRange) {
-            computeComparisonClass(this@Transformer.context.ir.symbols, lowerExpression.type, upperExpression.type, argExpression.type)
-                ?: return null
-        } else {
-            assert(headerInfo is ComparableRangeInfo)
-            this@Transformer.context.ir.symbols.comparable.owner
-        }
-
         if (isNumericRange) {
             lowerExpression = lowerExpression.castIfNecessary(comparisonClass)
             upperExpression = upperExpression.castIfNecessary(comparisonClass)
-            argExpression = argExpression.castIfNecessary(comparisonClass)
         }
 
         val lessOrEqualFun = builtIns.lessOrEqualFunByOperandType.getValue(if (useCompareTo) builtIns.intClass else comparisonClass.symbol)
         val compareToFun = comparisonClass.functions.singleOrNull {
             it.name == OperatorNameConventions.COMPARE_TO &&
                     it.dispatchReceiverParameter != null && it.extensionReceiverParameter == null &&
-                    it.valueParameters.size == 1 && (!isNumericRange || it.valueParameters[0].type == argument.type.makeNotNull())
+                    it.valueParameters.size == 1 && (!isNumericRange || it.valueParameters[0].type == comparisonClass.defaultType)
         } ?: return null
 
         // contains() function for ComparableRange is implemented as `value >= start && value <= endInclusive` (`value` is the argument).
@@ -355,7 +355,6 @@ private class Transformer(
         }
     }
 
-    // TODO: Handle unsigned.
     private fun computeComparisonClass(
         symbols: Symbols<CommonBackendContext>,
         lowerType: IrType,
@@ -373,25 +372,24 @@ private class Transformer(
         return when {
             pt1.isDouble() || pt2.isDouble() -> symbols.double
             pt1.isFloat() || pt2.isFloat() -> symbols.float
-//            pt1.isULong() || pt2.isULong() -> symbols.uLong!!
+            pt1.isULong() || pt2.isULong() -> symbols.uLong!!
-//            pt1.isUInt() || pt2.isUInt() -> symbols.uInt!!
+            pt1.isUInt() || pt2.isUInt() -> symbols.uInt!!
             pt1.isLong() || pt2.isLong() -> symbols.long
             pt1.isInt() || pt2.isInt() -> symbols.int
             pt1.isChar() || pt2.isChar() -> symbols.char
-            else -> return null
+            else -> error("Unexpected types: t1=${t1.classOrNull?.owner?.name}, t2=${t2.classOrNull?.owner?.name}")
-            // error("Unexpected types: t1=${t1.classOrNull?.owner?.name}, t2=${t2.classOrNull?.owner?.name}")
         }.defaultType
     }
 
     private fun IrType.promoteIntegerTypeToIntIfRequired(symbols: Symbols<CommonBackendContext>): IrType = when {
         isByte() || isShort() -> symbols.int.defaultType
-//        isUByte() || isUShort() -> symbols.uInt!!.defaultType
+        isUByte() || isUShort() -> symbols.uInt!!.defaultType
         else -> this
     }
 }
 
 internal open class RangeHeaderInfoBuilder(context: CommonBackendContext, scopeOwnerSymbol: () -> IrSymbol) :
-    HeaderInfoBuilder(context, scopeOwnerSymbol) {
+    HeaderInfoBuilder(context, scopeOwnerSymbol, allowUnsignedBounds = true) {
 
     override val progressionHandlers = listOf(
         CollectionIndicesHandler(context),
@@ -408,7 +406,7 @@ internal open class RangeHeaderInfoBuilder(context: CommonBackendContext, scopeO
         ReversedHandler(context, this)
     )
 
-    override val expressionHandlers = listOf(DefaultProgressionHandler(context))
+    override val expressionHandlers = listOf(DefaultProgressionHandler(context, allowUnsignedBounds = true))
 }
 
 /** Builds a [HeaderInfo] for closed floating-point ranges built using the `rangeTo` function. */

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderInfo.kt

@@ -255,7 +255,11 @@ internal interface HeaderInfoFromCallHandler<D> : HeaderInfoHandler<IrCall, D> {
 
 internal typealias ProgressionHandler = HeaderInfoFromCallHandler<ProgressionType>
 
-internal abstract class HeaderInfoBuilder(context: CommonBackendContext, private val scopeOwnerSymbol: () -> IrSymbol) :
+internal abstract class HeaderInfoBuilder(
+    context: CommonBackendContext,
+    private val scopeOwnerSymbol: () -> IrSymbol,
+    private val allowUnsignedBounds: Boolean = false
+) :
     IrElementVisitor<HeaderInfo?, IrCall?> {
 
     private val symbols = context.ir.symbols
@@ -284,7 +288,7 @@ internal abstract class HeaderInfoBuilder(context: CommonBackendContext, private
             return callHeaderInfo
 
         // Try to match a call to build a progression (e.g., `.indices`, `downTo`).
-        val progressionType = ProgressionType.fromIrType(expression.type, symbols)
+        val progressionType = ProgressionType.fromIrType(expression.type, symbols, allowUnsignedBounds)
         val progressionHeaderInfo =
             progressionType?.run { progressionHandlers.firstNotNullResult { it.handle(expression, data, this, scopeOwnerSymbol()) } }
 

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ProgressionType.kt

@@ -16,11 +16,12 @@ import org.jetbrains.kotlin.ir.builders.irLong
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
-import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.types.isSubtypeOfClass
 import org.jetbrains.kotlin.ir.util.defaultType
-import org.jetbrains.kotlin.ir.util.functions
-import org.jetbrains.kotlin.name.Name
 
 /** Represents a progression type in the Kotlin stdlib. */
 internal sealed class ProgressionType(
@@ -39,14 +40,21 @@ internal sealed class ProgressionType(
     fun IrExpression.asStepType() = castIfNecessary(stepClass)
 
     companion object {
-        fun fromIrType(irType: IrType, symbols: Symbols<CommonBackendContext>): ProgressionType? = when {
+        fun fromIrType(irType: IrType, symbols: Symbols<CommonBackendContext>, allowUnsignedBounds: Boolean): ProgressionType? =
-            irType.isSubtypeOfClass(symbols.charProgression) -> CharProgressionType(symbols)
+            when {
-            irType.isSubtypeOfClass(symbols.intProgression) -> IntProgressionType(symbols)
+                irType.isSubtypeOfClass(symbols.charProgression) -> CharProgressionType(symbols)
-            irType.isSubtypeOfClass(symbols.longProgression) -> LongProgressionType(symbols)
+                irType.isSubtypeOfClass(symbols.intProgression) -> IntProgressionType(symbols)
-            symbols.uIntProgression != null && irType.isSubtypeOfClass(symbols.uIntProgression) -> UIntProgressionType(symbols)
+                irType.isSubtypeOfClass(symbols.longProgression) -> LongProgressionType(symbols)
-            symbols.uLongProgression != null && irType.isSubtypeOfClass(symbols.uLongProgression) -> ULongProgressionType(symbols)
+                symbols.uIntProgression != null && irType.isSubtypeOfClass(symbols.uIntProgression) -> UIntProgressionType(
-            else -> null
+                    symbols,
-        }
+                    allowUnsignedBounds
+                )
+                symbols.uLongProgression != null && irType.isSubtypeOfClass(symbols.uLongProgression) -> ULongProgressionType(
+                    symbols,
+                    allowUnsignedBounds
+                )
+                else -> null
+            }
     }
 }
 
@@ -171,10 +179,10 @@ internal abstract class UnsignedProgressionType(
     }
 }
 
-internal class UIntProgressionType(symbols: Symbols<CommonBackendContext>) :
+internal class UIntProgressionType(symbols: Symbols<CommonBackendContext>,  allowUnsignedBounds: Boolean) :
     UnsignedProgressionType(
         symbols,
-        elementClass = symbols.int.owner,
+        elementClass = if (allowUnsignedBounds) symbols.uInt!!.owner else symbols.int.owner,
         stepClass = symbols.int.owner,
         minValueAsLong = UInt.MIN_VALUE.toLong(),
         maxValueAsLong = UInt.MAX_VALUE.toLong(),
@@ -184,15 +192,15 @@ internal class UIntProgressionType(symbols: Symbols<CommonBackendContext>) :
         unsignedConversionFunction = symbols.toUIntByExtensionReceiver.getValue(symbols.int)
     ) {
     @OptIn(ExperimentalUnsignedTypes::class)
-    override fun DeclarationIrBuilder.minValueExpression() = irInt(UInt.MIN_VALUE.toInt())
+    override fun DeclarationIrBuilder.minValueExpression() = irInt(UInt.MIN_VALUE.toInt(), elementClass.defaultType)
 
     override fun DeclarationIrBuilder.zeroStepExpression() = irInt(0)
 }
 
-internal class ULongProgressionType(symbols: Symbols<CommonBackendContext>) :
+internal class ULongProgressionType(symbols: Symbols<CommonBackendContext>, private val allowUnsignedBounds: Boolean) :
     UnsignedProgressionType(
         symbols,
-        elementClass = symbols.long.owner,
+        elementClass = if (allowUnsignedBounds) symbols.uLong!!.owner else symbols.long.owner,
         stepClass = symbols.long.owner,
         minValueAsLong = ULong.MIN_VALUE.toLong(),
         maxValueAsLong = ULong.MAX_VALUE.toLong(),
@@ -202,7 +210,7 @@ internal class ULongProgressionType(symbols: Symbols<CommonBackendContext>) :
         unsignedConversionFunction = symbols.toULongByExtensionReceiver.getValue(symbols.long)
     ) {
     @OptIn(ExperimentalUnsignedTypes::class)
-    override fun DeclarationIrBuilder.minValueExpression() = irLong(ULong.MIN_VALUE.toLong())
+    override fun DeclarationIrBuilder.minValueExpression() = irLong(ULong.MIN_VALUE.toLong(), elementClass.defaultType)
 
     override fun DeclarationIrBuilder.zeroStepExpression() = irLong(0)
 }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/handlers/DefaultProgressionHandler.kt

@@ -18,7 +18,7 @@ import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.getPropertyGetter
 
 /** Builds a [HeaderInfo] for progressions not handled by more specialized handlers. */
-internal class DefaultProgressionHandler(private val context: CommonBackendContext) :
+internal class DefaultProgressionHandler(private val context: CommonBackendContext, private val allowUnsignedBounds: Boolean = false) :
     ExpressionHandler {
 
     private val symbols = context.ir.symbols
@@ -26,7 +26,8 @@ internal class DefaultProgressionHandler(private val context: CommonBackendConte
 
     override fun matchIterable(expression: IrExpression) = ProgressionType.fromIrType(
         expression.type,
-        symbols
+        symbols,
+        allowUnsignedBounds
     ) != null
 
     override fun build(expression: IrExpression, scopeOwner: IrSymbol): HeaderInfo? =
@@ -53,7 +54,7 @@ internal class DefaultProgressionHandler(private val context: CommonBackendConte
             val direction = if (isRange) ProgressionDirection.INCREASING else ProgressionDirection.UNKNOWN
 
             ProgressionHeaderInfo(
-                ProgressionType.fromIrType(progressionExpression.type, symbols)!!,
+                ProgressionType.fromIrType(progressionExpression.type, symbols, allowUnsignedBounds)!!,
                 first,
                 last,
                 step,

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/builders/ExpressionHelpers.kt

@@ -329,11 +329,11 @@ fun IrBuilderWithScope.irImplicitCast(argument: IrExpression, type: IrType) =
 fun IrBuilderWithScope.irReinterpretCast(argument: IrExpression, type: IrType) =
     IrTypeOperatorCallImpl(startOffset, endOffset, type, IrTypeOperator.REINTERPRET_CAST, type, argument)
 
-fun IrBuilderWithScope.irInt(value: Int) =
+fun IrBuilderWithScope.irInt(value: Int, type: IrType = context.irBuiltIns.intType) =
-    IrConstImpl.int(startOffset, endOffset, context.irBuiltIns.intType, value)
+    IrConstImpl.int(startOffset, endOffset, type, value)
 
-fun IrBuilderWithScope.irLong(value: Long) =
+fun IrBuilderWithScope.irLong(value: Long, type: IrType = context.irBuiltIns.longType) =
-    IrConstImpl.long(startOffset, endOffset, context.irBuiltIns.longType, value)
+    IrConstImpl.long(startOffset, endOffset, type, value)
 
 fun IrBuilderWithScope.irChar(value: Char) =
     IrConstImpl.char(startOffset, endOffset, context.irBuiltIns.charType, value)

compiler/testData/codegen/box/ranges/contains/inUntilMaxValue.kt

@@ -1,3 +1,4 @@
+// KJS_WITH_FULL_RUNTIME
 // WITH_RUNTIME
 
 fun box(): String {

compiler/testData/codegen/box/ranges/contains/inUntilMinValue.kt

@@ -1,3 +1,4 @@
+// KJS_WITH_FULL_RUNTIME
 // WITH_RUNTIME
 
 fun box(): String {

compiler/testData/codegen/box/ranges/contains/inUntilMinValueNonConst.kt

@@ -1,3 +1,4 @@
+// KJS_WITH_FULL_RUNTIME
 // WITH_RUNTIME
 
 fun box(): String {

compiler/testData/codegen/bytecodeText/ranges/inMixedUnsignedRange.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36829 Optimize 'in' expressions in JVM_IR
-
 fun ub_ub(x: UByte, a: UByte, b: UByte) = x in a..b
 fun ub_us(x: UByte, a: UShort, b: UShort) = x in a..b
 fun ub_ui(x: UByte, a: UInt, b: UInt) = x in a..b
@@ -16,9 +13,7 @@ fun ui_us(x: UInt, a: UShort, b: UShort) = x in a..b
 fun ui_ui(x: UInt, a: UInt, b: UInt) = x in a..b
 fun ui_ul(x: UInt, a: ULong, b: ULong) = x in a..b
 
-fun ul_ub(x: ULong, a: UByte, b: UByte) = x in a..b         // ULong in range of UInt, uses non-intrinsic 'contains'
+// ul_ub, ul_us, ul_ui combinations are tested in inMixedUnsignedRange_2.kt (different behavior in non-IR vs IR backends)
-fun ul_us(x: ULong, a: UShort, b: UShort) = x in a..b       // ULong in range of UInt, uses non-intrinsic 'contains'
-fun ul_ui(x: ULong, a: UInt, b: UInt) = x in a..b           // ULong in range of UInt, uses non-intrinsic 'contains'
 fun ul_ul(x: ULong, a: ULong, b: ULong) = x in a..b
 
 fun n_ub_ub(x: UByte, a: UByte, b: UByte) = x !in a..b
@@ -36,17 +31,16 @@ fun n_ui_us(x: UInt, a: UShort, b: UShort) = x !in a..b
 fun n_ui_ui(x: UInt, a: UInt, b: UInt) = x !in a..b
 fun n_ui_ul(x: UInt, a: ULong, b: ULong) = x !in a..b
 
-fun n_ul_ub(x: ULong, a: UByte, b: UByte) = x !in a..b      // ULong in range of UInt, uses non-intrinsic 'contains'
+// n_ul_ub, n_ul_us, n_ul_ui combinations are tested in inMixedUnsignedRange_2.kt (different behavior in non-IR vs IR backends)
-fun n_ul_us(x: ULong, a: UShort, b: UShort) = x !in a..b    // ULong in range of UInt, uses non-intrinsic 'contains'
-fun n_ul_ui(x: ULong, a: UInt, b: UInt) = x !in a..b        // ULong in range of UInt, uses non-intrinsic 'contains'
 fun n_ul_ul(x: ULong, a: ULong, b: ULong) = x !in a..b
 
-// 6 contains
+// 0 contains
-// 13 IFLE
-// 13 IFLT
-// 13 IFGE
-// 13 IFGT
 // 0 L2I
-// 22 SIPUSH 255
+// 18 SIPUSH 255
-// 24 LDC 65535
+// 2 LDC 255
+// 20 LDC 65535
 // 2 LDC 4294967295
+
+// "SIPUSH/LDC 255" represent conversion from UByte to UInt/ULong
+// "LDC 65535" is UShort to UInt/ULong
+// "LDC 4294967295" is UInt to ULong

compiler/testData/codegen/bytecodeText/ranges/inMixedUnsignedRange_2.kt

@@ -0,0 +1,22 @@
+// ULong in range of UInt, uses non-intrinsic 'contains' for non-IR backend
+
+fun ul_ub(x: ULong, a: UByte, b: UByte) = x in a..b
+fun ul_us(x: ULong, a: UShort, b: UShort) = x in a..b
+fun ul_ui(x: ULong, a: UInt, b: UInt) = x in a..b
+
+fun n_ul_ub(x: ULong, a: UByte, b: UByte) = x !in a..b
+fun n_ul_us(x: ULong, a: UShort, b: UShort) = x !in a..b
+fun n_ul_ui(x: ULong, a: UInt, b: UInt) = x !in a..b
+
+// JVM_TEMPLATES
+// 6 contains
+
+// JVM_IR_TEMPLATES
+// 0 contains
+// 4 LDC 255
+// 4 LDC 65535
+// 4 LDC 4294967295
+
+// "LDC 255" represent conversion from UByte to ULong
+// "LDC 65535" is UShort to ULong
+// "LDC 4294967295" is UInt to ULong

compiler/testData/codegen/bytecodeText/ranges/inNonMatchingRangeIntrinsified.kt

@@ -14,7 +14,6 @@ fun inDouble(x: Float): Boolean {
     return x in 1.0..2.0
 }
 
-// 2 I2L
 // 0 INVOKESPECIAL
 // 0 NEW
 // 0 rangeTo
@@ -25,7 +24,9 @@ fun inDouble(x: Float): Boolean {
 // 0 contains
 
 // JVM_TEMPLATES
+// 2 I2L
 // 3 F2D
 
 // JVM_IR_TEMPLATES
-// 2 F2D
+// 1 I2L
+// 1 F2D

compiler/testData/codegen/bytecodeText/ranges/inOptimizableUnsignedRange.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36829 Optimize 'in' expressions in JVM_IR
-
 fun testUIntRangeLiteral(a: UInt, b: UInt) = 42u in a .. b
 
 fun testULongRangeLiteral(a: ULong, b: ULong) = 42UL in a .. b

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

@@ -4283,6 +4283,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/ranges/inMixedUnsignedRange.kt");
         }
 
+        @TestMetadata("inMixedUnsignedRange_2.kt")
+        public void testInMixedUnsignedRange_2() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/ranges/inMixedUnsignedRange_2.kt");
+        }
+
         @TestMetadata("inNonMatchingRangeIntrinsified.kt")
         public void testInNonMatchingRangeIntrinsified() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/ranges/inNonMatchingRangeIntrinsified.kt");

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

@@ -4251,6 +4251,11 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/ranges/inMixedUnsignedRange.kt");
         }
 
+        @TestMetadata("inMixedUnsignedRange_2.kt")
+        public void testInMixedUnsignedRange_2() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/ranges/inMixedUnsignedRange_2.kt");
+        }
+
         @TestMetadata("inNonMatchingRangeIntrinsified.kt")
         public void testInNonMatchingRangeIntrinsified() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/ranges/inNonMatchingRangeIntrinsified.kt");