RangeContainsLowering: Handle Comparable ranges.

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/Ir.kt

@@ -223,6 +223,7 @@ open class BuiltinSymbolsBase(protected val irBuiltIns: IrBuiltIns, protected va
     val mutableIterable = symbolTable.referenceClass(builtIns.mutableIterable)
     val mutableIterator = symbolTable.referenceClass(builtIns.mutableIterator)
     val mutableListIterator = symbolTable.referenceClass(builtIns.mutableListIterator)
+    val comparable = symbolTable.referenceClass(builtIns.comparable)
 
     private val binaryOperatorCache = mutableMapOf<Triple<Name, KotlinType, KotlinType>, IrSimpleFunctionSymbol>()
 

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

@@ -35,12 +35,10 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addIfNotNull
 
 // TODO:
-// 1. Handle (kotlin.ranges.)ClosedRange.contains on (kotlin.ranges.)Comparable.rangeTo(Comparable).
+// - Handle contains on ClosedFloatingPointRange non-literal expression (similar to DefaultProgressionHandler)
-//    Make sure right extension function was called!
+// - Note for unsigned until, the decremented last is signed (see UntilHandler.kt:81). Rather than converting back to unsigned,
-// 2. Handle contains on ClosedFloatingPointRange non-literal expression (similar to DefaultProgressionHandler)
-// 3. 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?
-// 4. Unsigned step has similar concerns as well. getProgressionLastElement return value is signed.
+// - Unsigned step has similar concerns as well. getProgressionLastElement return value is signed.
 
 val rangeContainsLoweringPhase = makeIrFilePhase(
     ::RangeContainsLowering,
@@ -133,6 +131,7 @@ private class Transformer(
         val upper: IrExpression
         val shouldUpperComeFirst: Boolean
         val useCompareTo: Boolean
+        val isNumericRange: Boolean
         val additionalNotEmptyCondition: IrExpression?
         val additionalStatements = mutableListOf<IrStatement>()
 
@@ -184,6 +183,7 @@ private class Transformer(
 
                 // `compareTo` must be used for UInt/ULong; they don't have intrinsic comparison operators.
                 useCompareTo = headerInfo.progressionType is UnsignedProgressionType
+                isNumericRange = true
                 additionalNotEmptyCondition = headerInfo.additionalNotEmptyCondition
             }
             is FloatingPointRangeHeaderInfo -> {
@@ -191,6 +191,15 @@ private class Transformer(
                 upper = headerInfo.endInclusive
                 shouldUpperComeFirst = false
                 useCompareTo = false
+                isNumericRange = true
+                additionalNotEmptyCondition = null
+            }
+            is ComparableRangeInfo -> {
+                lower = headerInfo.start
+                upper = headerInfo.endInclusive
+                shouldUpperComeFirst = false
+                useCompareTo = true
+                isNumericRange = false
                 additionalNotEmptyCondition = null
             }
             else -> return null
@@ -226,9 +235,9 @@ 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.
 
-        val (argVar, argExpression) = createTemporaryVariableIfNecessary(argument, "containsArg")
+        var (argVar, argExpression) = createTemporaryVariableIfNecessary(argument, "containsArg")
-        val lowerExpression: IrExpression
+        var lowerExpression: IrExpression
-        val upperExpression: IrExpression
+        var upperExpression: IrExpression
         val useLowerClauseOnLeftSide: Boolean
         if (argVar != null) {
             val (lowerVar, tmpLowerExpression) = createTemporaryVariableIfNecessary(lower, "containsLower")
@@ -266,44 +275,57 @@ private class Transformer(
 
         // TODO: Handle unsigned (comparisonClass is currently null).
         val builtIns = context.irBuiltIns
-        val comparisonClass =
+        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 && it.valueParameters[0].type == argument.type.makeNotNull()
+                    it.valueParameters.size == 1 && (!isNumericRange || it.valueParameters[0].type == argument.type.makeNotNull())
         } ?: return null
-        // TODO: Handle null. This can happen with an extension fun, e.g., IntRange.compareTo(String) (See rangeContainsString.kt).
 
+        // contains() function for ComparableRange is implemented as `value >= start && value <= endInclusive` (`value` is the argument).
+        // Therefore the dispatch receiver for the compareTo() calls should be the argument. This is important since the implementation
+        // for compareTo() may have side effects dependent on which expressions are the receiver and argument
+        // (see evaluationOrderForComparableRange.kt test).
         val lowerClause = if (useCompareTo) {
+            irCall(lessOrEqualFun).apply {
+                putValueArgument(0, irInt(0))
+                putValueArgument(1, irCall(compareToFun).apply {
+                    dispatchReceiver = argExpression
+                    putValueArgument(0, lowerExpression)
+                })
+            }
+        } else {
+            irCall(lessOrEqualFun).apply {
+                putValueArgument(0, lowerExpression)
+                putValueArgument(1, argExpression)
+            }
+        }
+        val upperClause = if (useCompareTo) {
             irCall(lessOrEqualFun).apply {
                 putValueArgument(0, irCall(compareToFun).apply {
-                    dispatchReceiver = lowerExpression.castIfNecessary(comparisonClass)
+                    dispatchReceiver = argExpression.deepCopyWithSymbols()
-                    putValueArgument(0, argExpression.castIfNecessary(comparisonClass))
+                    putValueArgument(0, upperExpression)
                 })
                 putValueArgument(1, irInt(0))
             }
         } else {
             irCall(lessOrEqualFun).apply {
-                putValueArgument(0, lowerExpression.castIfNecessary(comparisonClass))
+                putValueArgument(0, argExpression.deepCopyWithSymbols())
-                putValueArgument(1, argExpression.castIfNecessary(comparisonClass))
+                putValueArgument(1, upperExpression)
-            }
-        }
-        val upperClause = if (useCompareTo) {
-            irCall(lessOrEqualFun).apply {
-                putValueArgument(0, irInt(0))
-                putValueArgument(1, irCall(compareToFun).apply {
-                    dispatchReceiver = upperExpression.castIfNecessary(comparisonClass)
-                    putValueArgument(0, argExpression.deepCopyWithSymbols().castIfNecessary(comparisonClass))
-                })
-            }
-        } else {
-            irCall(lessOrEqualFun).apply {
-                putValueArgument(0, argExpression.deepCopyWithSymbols().castIfNecessary(comparisonClass))
-                putValueArgument(1, upperExpression.castIfNecessary(comparisonClass))
             }
         }
 
@@ -378,7 +400,11 @@ internal open class RangeHeaderInfoBuilder(context: CommonBackendContext, scopeO
         RangeToHandler(context)
     )
 
-    override val callHandlers = listOf(ReversedHandler(context, this), FloatingPointRangeToHandler)
+    override val callHandlers = listOf(
+        FloatingPointRangeToHandler,
+        ComparableRangeToHandler(context),
+        ReversedHandler(context, this)
+    )
 
     override val expressionHandlers = listOf(DefaultProgressionHandler(context))
 }
@@ -387,7 +413,7 @@ internal open class RangeHeaderInfoBuilder(context: CommonBackendContext, scopeO
 internal object FloatingPointRangeToHandler : HeaderInfoFromCallHandler<Nothing?> {
 
     override val matcher = SimpleCalleeMatcher {
-        fqName { it == FqName("kotlin.ranges.rangeTo") }
+        fqName { it == FqName("kotlin.ranges.${OperatorNameConventions.RANGE_TO}") }
         extensionReceiver { it != null && it.type.run { isFloat() || isDouble() } }
         parameterCount { it == 1 }
         parameter(0) { it.type.run { isFloat() || isDouble() } }
@@ -399,3 +425,19 @@ internal object FloatingPointRangeToHandler : HeaderInfoFromCallHandler<Nothing?
             endInclusive = expression.getValueArgument(0)!!
         )
 }
+
+/** Builds a [HeaderInfo] for ranges of Comparables built using the `rangeTo` extension function. */
+internal class ComparableRangeToHandler(context: CommonBackendContext) : HeaderInfoFromCallHandler<Nothing?> {
+
+    override val matcher = SimpleCalleeMatcher {
+        fqName { it == FqName("kotlin.ranges.${OperatorNameConventions.RANGE_TO}") }
+        extensionReceiver { it != null && it.type.isSubtypeOfClass(context.ir.symbols.comparable) }
+        parameterCount { it == 1 }
+    }
+
+    override fun build(expression: IrCall, data: Nothing?, scopeOwner: IrSymbol) =
+        ComparableRangeInfo(
+            start = expression.extensionReceiver!!,
+            endInclusive = expression.getValueArgument(0)!!
+        )
+}

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

@@ -55,6 +55,13 @@ internal class FloatingPointRangeHeaderInfo(
     override fun asReversed(): HeaderInfo? = null
 }
 
+internal class ComparableRangeInfo(
+    val start: IrExpression,
+    val endInclusive: IrExpression
+) : HeaderInfo() {
+    override fun asReversed(): HeaderInfo? = null
+}
+
 internal sealed class NumericHeaderInfo(
     val progressionType: ProgressionType,
     val first: IrExpression,

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

@@ -440,7 +440,7 @@ internal class WithIndexLoopHeader(
                 is ProgressionHeaderInfo -> ProgressionLoopHeader(nestedInfo, this@with, context)
                 is IterableHeaderInfo -> IterableLoopHeader(nestedInfo)
                 is WithIndexHeaderInfo -> throw IllegalStateException("Nested WithIndexHeaderInfo not allowed for WithIndexLoopHeader")
-                is FloatingPointRangeHeaderInfo -> error("Unexpected FloatingPointRangeHeaderInfo for loops")
+                is FloatingPointRangeHeaderInfo, is ComparableRangeInfo -> error("Unexpected ${nestedInfo::class.simpleName} for loops")
             }
 
             // Do not build own indexVariable if the nested loop header has an inductionVariable == 0 and step == 1.
@@ -669,7 +669,7 @@ internal class HeaderProcessor(
             is ProgressionHeaderInfo -> ProgressionLoopHeader(headerInfo, builder, context)
             is WithIndexHeaderInfo -> WithIndexLoopHeader(headerInfo, builder, context)
             is IterableHeaderInfo -> IterableLoopHeader(headerInfo)
-            is FloatingPointRangeHeaderInfo -> error("Unexpected FloatingPointRangeHeaderInfo for loops")
+            is FloatingPointRangeHeaderInfo, is ComparableRangeInfo -> error("Unexpected ${headerInfo::class.simpleName} for loops")
         }
     }
 }

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

@@ -1,8 +1,8 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36829 Optimize 'in' expressions in JVM_IR
 fun test1(a: String) = a in "alpha" .. "omega"
 fun test2(a: String) = a !in "alpha" .. "omega"
+fun <T : Comparable<T>> test3(x: T, left: T, right: T) = x in left .. right
+fun <T : Enum<T>> test4(x: T, left: T, right: T) = x in left .. right
 
 // 0 rangeTo
 // 0 contains
-// 4 compareTo
+// 8 compareTo