Support const-bound counter loop generation for 'downTo'

compiler/backend/src/org/jetbrains/kotlin/codegen/codegenUtil.kt

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils.isSubclass
 import org.jetbrains.kotlin.resolve.annotations.hasJvmStaticAnnotation
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoBefore
+import org.jetbrains.kotlin.resolve.calls.callUtil.getFirstArgumentExpression
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
@@ -420,11 +421,11 @@ fun extractReificationArgument(type: KotlinType): Pair<TypeParameterDescriptor,
 fun unwrapInitialSignatureDescriptor(function: FunctionDescriptor): FunctionDescriptor =
         function.initialSignatureDescriptor ?: function
 
-fun ExpressionCodegen.generateCallReceiver(rangeCall: ResolvedCall<out CallableDescriptor>): StackValue =
+fun ExpressionCodegen.generateCallReceiver(call: ResolvedCall<out CallableDescriptor>): StackValue =
-        generateReceiverValue(rangeCall.extensionReceiver ?: rangeCall.dispatchReceiver!!, false)
+        generateReceiverValue(call.extensionReceiver ?: call.dispatchReceiver!!, false)
 
-fun ExpressionCodegen.generateCallSingleArgument(rangeCall: ResolvedCall<out CallableDescriptor>): StackValue =
+fun ExpressionCodegen.generateCallSingleArgument(call: ResolvedCall<out CallableDescriptor>): StackValue =
-        gen(ExpressionCodegen.getSingleArgumentExpression(rangeCall)!!)
+        gen(call.getFirstArgumentExpression()!!)
 
 fun ClassDescriptor.isPossiblyUninitializedSingleton() =
         DescriptorUtils.isEnumEntry(this) ||

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

@@ -20,8 +20,11 @@ import org.jetbrains.kotlin.codegen.ExpressionCodegen
 import org.jetbrains.kotlin.codegen.generateCallReceiver
 import org.jetbrains.kotlin.codegen.generateCallSingleArgument
 import org.jetbrains.kotlin.codegen.range.forLoop.ForInSimpleProgressionLoopGenerator
+import org.jetbrains.kotlin.codegen.range.forLoop.ForLoopGenerator
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.resolve.calls.callUtil.getFirstArgumentExpression
+import org.jetbrains.kotlin.resolve.calls.callUtil.getReceiverExpression
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 
 class DownToProgressionRangeValue(rangeCall: ResolvedCall<out CallableDescriptor>) :
@@ -35,8 +38,36 @@ class DownToProgressionRangeValue(rangeCall: ResolvedCall<out CallableDescriptor
             )
 
     override fun createForLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression) =
+            createConstBoundedForInDownToGenerator(codegen, forExpression) ?:
             ForInSimpleProgressionLoopGenerator.fromBoundedValueWithStepMinus1(codegen, forExpression, getBoundedValue(codegen))
 
     override fun createForInReversedLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression) =
+            createConstBoundedForInReversedDownToGenerator(codegen, forExpression) ?:
             ForInSimpleProgressionLoopGenerator.fromBoundedValueWithStep1(codegen, forExpression, getBoundedValue(codegen))
+
+    private fun createConstBoundedForInDownToGenerator(
+            codegen: ExpressionCodegen,
+            forExpression: KtForExpression
+    ): ForLoopGenerator? {
+        val endExpression = rangeCall.getFirstArgumentExpression() ?: return null
+        return createConstBoundedForLoopGeneratorOrNull(
+                codegen, forExpression,
+                codegen.generateCallReceiver(rangeCall),
+                endExpression,
+                -1
+        )
+    }
+
+    private fun createConstBoundedForInReversedDownToGenerator(
+            codegen: ExpressionCodegen,
+            forExpression: KtForExpression
+    ): ForLoopGenerator? {
+        val endExpression = rangeCall.getReceiverExpression() ?: return null
+        return createConstBoundedForLoopGeneratorOrNull(
+                codegen, forExpression,
+                codegen.generateCallSingleArgument(rangeCall),
+                endExpression,
+                1
+        )
+    }
 }

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

@@ -23,9 +23,9 @@ import org.jetbrains.kotlin.codegen.range.forLoop.ForInSimpleProgressionLoopGene
 import org.jetbrains.kotlin.codegen.range.forLoop.ForLoopGenerator
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.resolve.calls.callUtil.getReceiverExpression
+import org.jetbrains.kotlin.resolve.calls.callUtil.getFirstArgumentExpression
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
-import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class PrimitiveNumberRangeLiteralRangeValue(
         rangeCall: ResolvedCall<out CallableDescriptor>
@@ -47,7 +47,7 @@ class PrimitiveNumberRangeLiteralRangeValue(
             codegen: ExpressionCodegen,
             forExpression: KtForExpression
     ): ForLoopGenerator? {
-        val endExpression = rangeCall.valueArgumentsByIndex?.run { get(0).arguments[0].getArgumentExpression() } ?: return null
+        val endExpression = rangeCall.getFirstArgumentExpression() ?: return null
         return createConstBoundedForLoopGeneratorOrNull(
                 codegen, forExpression,
                 codegen.generateCallReceiver(rangeCall),
@@ -60,7 +60,7 @@ class PrimitiveNumberRangeLiteralRangeValue(
             codegen: ExpressionCodegen,
             forExpression: KtForExpression
     ): ForLoopGenerator? {
-        val endExpression = rangeCall.extensionReceiver.safeAs<ExpressionReceiver>()?.expression ?: return null
+        val endExpression = rangeCall.getReceiverExpression() ?: return null
         return createConstBoundedForLoopGeneratorOrNull(
                 codegen, forExpression,
                 codegen.generateCallSingleArgument(rangeCall),

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/util/callUtil.kt

@@ -28,8 +28,10 @@ import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.CallTransformer
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.*
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.isError
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.sure
 
 // resolved call
@@ -256,3 +258,10 @@ fun Call.createLookupLocation(): KotlinLookupLocation {
             else callElement
     return KotlinLookupLocation(element)
 }
+
+fun ResolvedCall<*>.getFirstArgumentExpression(): KtExpression? =
+        valueArgumentsByIndex?.run { get(0).arguments[0].getArgumentExpression() }
+
+fun ResolvedCall<*>.getReceiverExpression(): KtExpression? =
+        extensionReceiver.safeAs<ExpressionReceiver>()?.expression ?:
+        dispatchReceiver.safeAs<ExpressionReceiver>()?.expression

compiler/testData/codegen/box/ranges/forInDownTo/forIntInDownToWithNonConstBounds.kt

@@ -0,0 +1,16 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun low() = 4
+fun high() = 1
+
+fun box(): String {
+    var sum = 0
+    for (i in low() downTo high()) {
+        sum = sum * 10 + i
+    }
+    assertEquals(4321, sum)
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToCharMinValue.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+
+const val M = 0.toChar()
+
+fun box(): String {
+    var count = 0
+    for (i in M downTo M) {
+        ++count
+        if (count > 1) {
+            throw AssertionError("Loop should be executed once")
+        }
+    }
+    if (count != 1) throw AssertionError("Should be executed once")
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToIntMinValue.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+
+const val M = Int.MIN_VALUE
+
+fun box(): String {
+    var count = 0
+    for (i in M downTo M) {
+        ++count
+        if (count > 1) {
+            throw AssertionError("Loop should be executed once")
+        }
+    }
+    if (count != 1) throw AssertionError("Should be executed once")
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToLongMinValue.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+
+const val M = Long.MIN_VALUE
+
+fun box(): String {
+    var count = 0
+    for (i in M downTo M) {
+        ++count
+        if (count > 1) {
+            throw AssertionError("Loop should be executed once")
+        }
+    }
+    if (count != 1) throw AssertionError("Should be executed once")
+    return "OK"
+}

compiler/testData/codegen/bytecodeText/forLoop/forInReversed/forInReversedRangeLiteral.kt

@@ -29,4 +29,7 @@ fun box(): String {
 // 0 getFirst
 // 0 getLast
 // 0 getStep
+// 0 IFEQ
 // 0 IF_ICMPEQ
+// 2 IF_ICMPLT
+// 1 LCMP

compiler/testData/codegen/bytecodeText/forLoop/forIntInDownTo.kt

@@ -13,3 +13,5 @@ fun test(): Int {
 // 0 getEnd
 // 0 getFirst
 // 0 getLast
+// 0 IF_ICMPEQ
+// 1 IF_ICMPLT

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

@@ -15103,6 +15103,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 doTest(fileName);
             }
 
+            @TestMetadata("forIntInDownToWithNonConstBounds.kt")
+            public void testForIntInDownToWithNonConstBounds() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInDownToWithNonConstBounds.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forIntInNonOptimizedDownTo.kt")
             public void testForIntInNonOptimizedDownTo() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInNonOptimizedDownTo.kt");
@@ -15379,6 +15385,24 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/ranges/forWithPossibleOverflow"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
             }
 
+            @TestMetadata("forInDownToCharMinValue.kt")
+            public void testForInDownToCharMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToCharMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToIntMinValue.kt")
+            public void testForInDownToIntMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToIntMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToLongMinValue.kt")
+            public void testForInDownToLongMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToLongMinValue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forInRangeToCharMaxValue.kt")
             public void testForInRangeToCharMaxValue() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInRangeToCharMaxValue.kt");

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

@@ -15103,6 +15103,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("forIntInDownToWithNonConstBounds.kt")
+            public void testForIntInDownToWithNonConstBounds() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInDownToWithNonConstBounds.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forIntInNonOptimizedDownTo.kt")
             public void testForIntInNonOptimizedDownTo() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInNonOptimizedDownTo.kt");
@@ -15379,6 +15385,24 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/ranges/forWithPossibleOverflow"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
             }
 
+            @TestMetadata("forInDownToCharMinValue.kt")
+            public void testForInDownToCharMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToCharMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToIntMinValue.kt")
+            public void testForInDownToIntMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToIntMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToLongMinValue.kt")
+            public void testForInDownToLongMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToLongMinValue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forInRangeToCharMaxValue.kt")
             public void testForInRangeToCharMaxValue() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInRangeToCharMaxValue.kt");

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

@@ -15103,6 +15103,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 doTest(fileName);
             }
 
+            @TestMetadata("forIntInDownToWithNonConstBounds.kt")
+            public void testForIntInDownToWithNonConstBounds() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInDownToWithNonConstBounds.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forIntInNonOptimizedDownTo.kt")
             public void testForIntInNonOptimizedDownTo() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInNonOptimizedDownTo.kt");
@@ -15379,6 +15385,24 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/ranges/forWithPossibleOverflow"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
             }
 
+            @TestMetadata("forInDownToCharMinValue.kt")
+            public void testForInDownToCharMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToCharMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToIntMinValue.kt")
+            public void testForInDownToIntMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToIntMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToLongMinValue.kt")
+            public void testForInDownToLongMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToLongMinValue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forInRangeToCharMaxValue.kt")
             public void testForInRangeToCharMaxValue() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInRangeToCharMaxValue.kt");

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

@@ -16519,6 +16519,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("forIntInDownToWithNonConstBounds.kt")
+            public void testForIntInDownToWithNonConstBounds() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInDownToWithNonConstBounds.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forIntInNonOptimizedDownTo.kt")
             public void testForIntInNonOptimizedDownTo() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInDownTo/forIntInNonOptimizedDownTo.kt");
@@ -16795,6 +16801,24 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/ranges/forWithPossibleOverflow"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JS, true);
             }
 
+            @TestMetadata("forInDownToCharMinValue.kt")
+            public void testForInDownToCharMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToCharMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToIntMinValue.kt")
+            public void testForInDownToIntMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToIntMinValue.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInDownToLongMinValue.kt")
+            public void testForInDownToLongMinValue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInDownToLongMinValue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forInRangeToCharMaxValue.kt")
             public void testForInRangeToCharMaxValue() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forWithPossibleOverflow/forInRangeToCharMaxValue.kt");